diff --git a/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/accuracy/accuracy.json b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/accuracy/accuracy.json new file mode 100644 index 00000000..c35ff064 --- /dev/null +++ b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/accuracy/accuracy.json @@ -0,0 +1,8 @@ +{ + "subset_score": 0.41, + "baseline_delta": 0.03, + "valid": true, + "framework": "SGLang", + "precision": "BF16", + "notes": "Integrated accuracy check \u2014 used same SGLang instance as benchmark." +} \ No newline at end of file diff --git a/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/env_info.json b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/env_info.json new file mode 100644 index 00000000..049d327a --- /dev/null +++ b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/env_info.json @@ -0,0 +1,64 @@ +{ + "collected_at": "2026-07-31T02:37:50.313608+00:00", + "accelerators": [ + { + "index": 5, + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "memory_gb": 80.0, + "driver_version": "570.133.20", + "firmware_version": null, + "compute_capability": "8.0", + "supports_bf16": true + }, + { + "index": 7, + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "memory_gb": 80.0, + "driver_version": "570.133.20", + "firmware_version": null, + "compute_capability": "8.0", + "supports_bf16": true + } + ], + "accelerator_platform": "nvidia", + "accelerator_topology": null, + "intra_node_interconnect": null, + "cpu": { + "model": "Intel(R) Xeon(R) Processor @ 2.90GHz", + "physical_cores": 96, + "logical_cores": 96, + "numa_nodes": 2 + }, + "system_memory_gb": 1000.0, + "pcie_generation": "PCIe Gen 4", + "cpu_accelerator_bandwidth_gbs": null, + "network_interfaces": [ + { + "name": "mlx5_0", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_1", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_2", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_3", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + } + ], + "os": "Ubuntu 22.04.5 LTS", + "python_version": "3.11.13", + "kernel_version": "5.10.134-013.9.kangaroo.al8.x86_64", + "runtime_version": "CUDA 12.8", + "pytorch_version": "2.9.1+cu128" +} \ No newline at end of file diff --git a/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/interactive/result.json b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/interactive/result.json new file mode 100644 index 00000000..bbafd4ba --- /dev/null +++ b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/interactive/result.json @@ -0,0 +1,169 @@ +{ + "schema_version": "1.0", + "suite_id": "suite_F", + "implementation_id": "nvidia_sglang_inference_opt_4fd1e59e", + "chip": { + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "count": 1, + "memory_gb": 80.0, + "interconnect_intra_node": null, + "interconnect_inter_node": null + }, + "environment": { + "collected_at": "2026-07-31T02:37:50.313608+00:00", + "accelerators": [ + { + "index": 5, + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "memory_gb": 80.0, + "driver_version": "570.133.20", + "firmware_version": null, + "compute_capability": "8.0", + "supports_bf16": true + }, + { + "index": 7, + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "memory_gb": 80.0, + "driver_version": "570.133.20", + "firmware_version": null, + "compute_capability": "8.0", + "supports_bf16": true + } + ], + "accelerator_platform": "nvidia", + "accelerator_topology": null, + "intra_node_interconnect": null, + "cpu": { + "model": "Intel(R) Xeon(R) Processor @ 2.90GHz", + "physical_cores": 96, + "logical_cores": 96, + "numa_nodes": 2 + }, + "system_memory_gb": 1000.0, + "pcie_generation": "PCIe Gen 4", + "cpu_accelerator_bandwidth_gbs": null, + "network_interfaces": [ + { + "name": "mlx5_0", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_1", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_2", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_3", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + } + ], + "os": "Ubuntu 22.04.5 LTS", + "python_version": "3.11.13", + "kernel_version": "5.10.134-013.9.kangaroo.al8.x86_64", + "runtime_version": "CUDA 12.8", + "pytorch_version": "2.9.1+cu128" + }, + "software": { + "framework": "SGLang", + "framework_version": "0.5.6", + "driver_version": "570.133.20", + "runtime_version": "CUDA 12.8", + "os": "Ubuntu 22.04.5 LTS", + "python_version": "3.11.13" + }, + "model": { + "model_id": "Qwen/Qwen2.5-0.5B-Instruct", + "model_revision": "7ae557604adf67be50417f59c2c2f167def9a775", + "model_name": null, + "model_note": "Local ModelScope snapshot for Suite F", + "model_source": "local", + "architecture": "dense", + "parameter_count_b": 0.5, + "precision": "BF16", + "effective_dtype": "bfloat16", + "quantization_method": null, + "model_format": "HuggingFace original" + }, + "task": { + "scenario": "interactive", + "num_runs": 3, + "warmup_runs": 1, + "parallelism": { + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "expert_parallel_size": 1, + "data_parallel_size": 1 + }, + "extra_config": { + "enable_tokenizer_batch_encode": true, + "continuous_decode_steps": 4, + "engine_kwargs": { + "attention_backend": "triton", + "disable_cuda_graph_padding": true, + "enable_piecewise_cuda_graph": true, + "piecewise_cuda_graph_max_tokens": 1024, + "prefill_attention_backend": "fa3" + } + }, + "runtime_metrics": null + }, + "metrics": { + "interactive": { + "ttft_ms_p50": 7.19, + "ttft_ms_p90": 7.54, + "ttft_ms_p99": 15.06, + "tpot_ms_p50": 1.62, + "tpot_ms_p90": 1.64, + "tpot_ms_p99": 1.71, + "peak_memory_gb": null, + "elapsed_seconds_median": 46.2, + "ttft_ms_p99_reliability": { + "n": 3, + "mean": 23.42, + "std": 22.39, + "cv_pct": 95.62, + "stability": "high-variance", + "runs": [ + 49.2, + 12.23, + 8.83 + ] + } + } + }, + "accuracy": { + "subset_score": null, + "baseline_delta": null, + "valid": false, + "notes": "Run --scenario accuracy to check model accuracy." + }, + "meta": { + "submitted_by": "ChiangYuhsin", + "submission_type": "individual", + "date": "2026-07-31", + "time": "02:46:54", + "run_id": "1189951b", + "run_name": "nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b", + "flagged": null, + "reproduce_script": "runners/nvidia_sglang_inference_opt_4fd1e59e/runner.py", + "env_info_file": "../env_info.json", + "log_file": "run.log", + "samples_file": "samples.jsonl", + "notes": null, + "benchmark_start_time": "2026-07-31T02:44:33.515731+00:00", + "benchmark_end_time": "2026-07-31T02:46:54.781708+00:00", + "benchmark_elapsed_minutes": 2.4, + "model_load_seconds": 47.7 + } +} \ No newline at end of file diff --git a/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/offline/result.json b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/offline/result.json new file mode 100644 index 00000000..277d1dfd --- /dev/null +++ b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/offline/result.json @@ -0,0 +1,228 @@ +{ + "schema_version": "1.0", + "suite_id": "suite_F", + "implementation_id": "nvidia_sglang_inference_opt_4fd1e59e", + "chip": { + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "count": 1, + "memory_gb": 80.0, + "interconnect_intra_node": null, + "interconnect_inter_node": null + }, + "environment": { + "collected_at": "2026-07-31T02:37:50.313608+00:00", + "accelerators": [ + { + "index": 5, + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "memory_gb": 80.0, + "driver_version": "570.133.20", + "firmware_version": null, + "compute_capability": "8.0", + "supports_bf16": true + }, + { + "index": 7, + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "memory_gb": 80.0, + "driver_version": "570.133.20", + "firmware_version": null, + "compute_capability": "8.0", + "supports_bf16": true + } + ], + "accelerator_platform": "nvidia", + "accelerator_topology": null, + "intra_node_interconnect": null, + "cpu": { + "model": "Intel(R) Xeon(R) Processor @ 2.90GHz", + "physical_cores": 96, + "logical_cores": 96, + "numa_nodes": 2 + }, + "system_memory_gb": 1000.0, + "pcie_generation": "PCIe Gen 4", + "cpu_accelerator_bandwidth_gbs": null, + "network_interfaces": [ + { + "name": "mlx5_0", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_1", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_2", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_3", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + } + ], + "os": "Ubuntu 22.04.5 LTS", + "python_version": "3.11.13", + "kernel_version": "5.10.134-013.9.kangaroo.al8.x86_64", + "runtime_version": "CUDA 12.8", + "pytorch_version": "2.9.1+cu128" + }, + "software": { + "framework": "SGLang", + "framework_version": "0.5.6", + "driver_version": "570.133.20", + "runtime_version": "CUDA 12.8", + "os": "Ubuntu 22.04.5 LTS", + "python_version": "3.11.13" + }, + "model": { + "model_id": "Qwen/Qwen2.5-0.5B-Instruct", + "model_revision": "7ae557604adf67be50417f59c2c2f167def9a775", + "model_name": null, + "model_note": "Local ModelScope snapshot for Suite F", + "model_source": "local", + "architecture": "dense", + "parameter_count_b": 0.5, + "precision": "BF16", + "effective_dtype": "bfloat16", + "quantization_method": null, + "model_format": "HuggingFace original" + }, + "task": { + "scenario": "offline", + "num_runs": 3, + "warmup_runs": 1, + "parallelism": { + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "expert_parallel_size": 1, + "data_parallel_size": 1 + }, + "extra_config": { + "enable_tokenizer_batch_encode": true, + "continuous_decode_steps": 4, + "engine_kwargs": { + "attention_backend": "triton", + "disable_cuda_graph_padding": true, + "enable_piecewise_cuda_graph": true, + "piecewise_cuda_graph_max_tokens": 1024, + "prefill_attention_backend": "fa3", + "enable_tokenizer_batch_encode": true, + "num_continuous_decode_steps": 4 + } + }, + "runtime_metrics": null + }, + "metrics": { + "offline": { + "results_by_concurrency": [ + { + "client_concurrency": 4, + "throughput_tokens_per_sec": 46128.3, + "throughput_tokens_per_sec_per_chip": 46128.3, + "throughput_tokens_per_sec_total": 62323.45, + "elapsed_seconds_median": 0.9, + "peak_memory_gb": null, + "power_watts_avg": null, + "power_watts_peak": null, + "oom": false, + "throughput_tokens_per_sec_reliability": { + "n": 3, + "mean": 46035.52, + "std": 177.27, + "cv_pct": 0.39, + "stability": "stable", + "runs": [ + 45831.12, + 46147.16, + 46128.3 + ] + }, + "_throughput_note": "output_only", + "_concurrency_note": "client_concurrency is the number of requests sent simultaneously. The inference engine batches internally; this does not directly set engine parameters like max_num_seqs." + }, + { + "client_concurrency": 16, + "throughput_tokens_per_sec": 46406.6, + "throughput_tokens_per_sec_per_chip": 46406.6, + "throughput_tokens_per_sec_total": 62699.46, + "elapsed_seconds_median": 0.9, + "peak_memory_gb": null, + "power_watts_avg": null, + "power_watts_peak": null, + "oom": false, + "throughput_tokens_per_sec_reliability": { + "n": 3, + "mean": 46306.02, + "std": 204.98, + "cv_pct": 0.44, + "stability": "stable", + "runs": [ + 46070.18, + 46406.6, + 46441.29 + ] + }, + "_throughput_note": "output_only", + "_concurrency_note": "client_concurrency is the number of requests sent simultaneously. The inference engine batches internally; this does not directly set engine parameters like max_num_seqs." + }, + { + "client_concurrency": 64, + "throughput_tokens_per_sec": 46500.82, + "throughput_tokens_per_sec_per_chip": 46500.82, + "throughput_tokens_per_sec_total": 62826.76, + "elapsed_seconds_median": 0.9, + "peak_memory_gb": null, + "power_watts_avg": null, + "power_watts_peak": null, + "oom": false, + "throughput_tokens_per_sec_reliability": { + "n": 3, + "mean": 46028.86, + "std": 890.06, + "cv_pct": 1.93, + "stability": "stable", + "runs": [ + 45002.22, + 46500.82, + 46583.55 + ] + }, + "_throughput_note": "output_only", + "_concurrency_note": "client_concurrency is the number of requests sent simultaneously. The inference engine batches internally; this does not directly set engine parameters like max_num_seqs." + } + ] + } + }, + "accuracy": { + "subset_score": null, + "baseline_delta": null, + "valid": false, + "notes": "Run --scenario accuracy to check model accuracy." + }, + "meta": { + "submitted_by": "ChiangYuhsin", + "submission_type": "individual", + "date": "2026-07-31", + "time": "02:40:05", + "run_id": "1189951b", + "run_name": "nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b", + "flagged": null, + "reproduce_script": "runners/nvidia_sglang_inference_opt_4fd1e59e/runner.py", + "env_info_file": "../env_info.json", + "log_file": "run.log", + "samples_file": "samples.jsonl", + "notes": null, + "benchmark_start_time": "2026-07-31T02:39:52.638757+00:00", + "benchmark_end_time": "2026-07-31T02:40:05.992733+00:00", + "benchmark_elapsed_minutes": 0.2, + "model_load_seconds": 47.1 + } +} \ No newline at end of file diff --git a/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/online/result.json b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/online/result.json new file mode 100644 index 00000000..9b503059 --- /dev/null +++ b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/online/result.json @@ -0,0 +1,201 @@ +{ + "schema_version": "1.0", + "suite_id": "suite_F", + "implementation_id": "nvidia_sglang_inference_opt_4fd1e59e", + "chip": { + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "count": 1, + "memory_gb": 80.0, + "interconnect_intra_node": null, + "interconnect_inter_node": null + }, + "environment": { + "collected_at": "2026-07-31T02:37:50.313608+00:00", + "accelerators": [ + { + "index": 5, + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "memory_gb": 80.0, + "driver_version": "570.133.20", + "firmware_version": null, + "compute_capability": "8.0", + "supports_bf16": true + }, + { + "index": 7, + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "memory_gb": 80.0, + "driver_version": "570.133.20", + "firmware_version": null, + "compute_capability": "8.0", + "supports_bf16": true + } + ], + "accelerator_platform": "nvidia", + "accelerator_topology": null, + "intra_node_interconnect": null, + "cpu": { + "model": "Intel(R) Xeon(R) Processor @ 2.90GHz", + "physical_cores": 96, + "logical_cores": 96, + "numa_nodes": 2 + }, + "system_memory_gb": 1000.0, + "pcie_generation": "PCIe Gen 4", + "cpu_accelerator_bandwidth_gbs": null, + "network_interfaces": [ + { + "name": "mlx5_0", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_1", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_2", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_3", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + } + ], + "os": "Ubuntu 22.04.5 LTS", + "python_version": "3.11.13", + "kernel_version": "5.10.134-013.9.kangaroo.al8.x86_64", + "runtime_version": "CUDA 12.8", + "pytorch_version": "2.9.1+cu128" + }, + "software": { + "framework": "SGLang", + "framework_version": "0.5.6", + "driver_version": "570.133.20", + "runtime_version": "CUDA 12.8", + "os": "Ubuntu 22.04.5 LTS", + "python_version": "3.11.13" + }, + "model": { + "model_id": "Qwen/Qwen2.5-0.5B-Instruct", + "model_revision": "7ae557604adf67be50417f59c2c2f167def9a775", + "model_name": null, + "model_note": "Local ModelScope snapshot for Suite F", + "model_source": "local", + "architecture": "dense", + "parameter_count_b": 0.5, + "precision": "BF16", + "effective_dtype": "bfloat16", + "quantization_method": null, + "model_format": "HuggingFace original" + }, + "task": { + "scenario": "online", + "num_runs": 3, + "warmup_runs": 1, + "parallelism": { + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "expert_parallel_size": 1, + "data_parallel_size": 1 + }, + "extra_config": { + "enable_tokenizer_batch_encode": true, + "continuous_decode_steps": 4, + "engine_kwargs": { + "attention_backend": "triton", + "disable_cuda_graph_padding": true, + "enable_piecewise_cuda_graph": true, + "piecewise_cuda_graph_max_tokens": 1024, + "prefill_attention_backend": "fa3" + } + }, + "runtime_metrics": null + }, + "metrics": { + "online": { + "sla_ttft_ms": 500, + "max_valid_qps": 40, + "results_by_qps": [ + { + "target_qps": 10, + "achieved_qps": 10.0, + "ttft_ms_p50": 8.39, + "ttft_ms_p90": 9.37, + "ttft_ms_p99": 14.12, + "tpot_ms_p50": 1.94, + "tpot_ms_p90": 2.05, + "tpot_ms_p99": 2.28, + "elapsed_seconds_median": 31.8, + "sla_met": true, + "ttft_ms_p99_reliability": { + "n": 3, + "mean": 17.9, + "std": 7.1, + "cv_pct": 39.67, + "stability": "high-variance", + "runs": [ + 25.95, + 15.2, + 12.54 + ] + } + }, + { + "target_qps": 40, + "achieved_qps": 40.0, + "ttft_ms_p50": 8.38, + "ttft_ms_p90": 11.9, + "ttft_ms_p99": 14.53, + "tpot_ms_p50": 2.4, + "tpot_ms_p90": 2.64, + "tpot_ms_p99": 3.18, + "elapsed_seconds_median": 7.8, + "sla_met": true, + "ttft_ms_p99_reliability": { + "n": 3, + "mean": 14.39, + "std": 0.16, + "cv_pct": 1.12, + "stability": "stable", + "runs": [ + 14.39, + 14.23, + 14.56 + ] + } + } + ] + } + }, + "accuracy": { + "subset_score": null, + "baseline_delta": null, + "valid": false, + "notes": "Run --scenario accuracy to check model accuracy." + }, + "meta": { + "submitted_by": "ChiangYuhsin", + "submission_type": "individual", + "date": "2026-07-31", + "time": "02:43:24", + "run_id": "1189951b", + "run_name": "nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b", + "flagged": null, + "reproduce_script": "runners/nvidia_sglang_inference_opt_4fd1e59e/runner.py", + "env_info_file": "../env_info.json", + "log_file": "run.log", + "samples_file": "samples.jsonl", + "notes": null, + "benchmark_start_time": "2026-07-31T02:41:19.824971+00:00", + "benchmark_end_time": "2026-07-31T02:43:24.313106+00:00", + "benchmark_elapsed_minutes": 2.1, + "model_load_seconds": 52.5 + } +} \ No newline at end of file diff --git a/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/result.json b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/result.json new file mode 100644 index 00000000..b4d1df95 --- /dev/null +++ b/results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/result.json @@ -0,0 +1,315 @@ +{ + "schema_version": "1.0", + "suite_id": "suite_F", + "implementation_id": "nvidia_sglang_inference_opt_4fd1e59e", + "chip": { + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "count": 1, + "memory_gb": 80.0, + "interconnect_intra_node": null, + "interconnect_inter_node": null + }, + "environment": { + "collected_at": "2026-07-31T02:37:50.313608+00:00", + "accelerators": [ + { + "index": 5, + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "memory_gb": 80.0, + "driver_version": "570.133.20", + "firmware_version": null, + "compute_capability": "8.0", + "supports_bf16": true + }, + { + "index": 7, + "name": "NVIDIA A100-SXM4-80GB", + "vendor": "NVIDIA", + "memory_gb": 80.0, + "driver_version": "570.133.20", + "firmware_version": null, + "compute_capability": "8.0", + "supports_bf16": true + } + ], + "accelerator_platform": "nvidia", + "accelerator_topology": null, + "intra_node_interconnect": null, + "cpu": { + "model": "Intel(R) Xeon(R) Processor @ 2.90GHz", + "physical_cores": 96, + "logical_cores": 96, + "numa_nodes": 2 + }, + "system_memory_gb": 1000.0, + "pcie_generation": "PCIe Gen 4", + "cpu_accelerator_bandwidth_gbs": null, + "network_interfaces": [ + { + "name": "mlx5_0", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_1", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_2", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + }, + { + "name": "mlx5_3", + "type": "InfiniBand/RoCE", + "bandwidth_gbps": null + } + ], + "os": "Ubuntu 22.04.5 LTS", + "python_version": "3.11.13", + "kernel_version": "5.10.134-013.9.kangaroo.al8.x86_64", + "runtime_version": "CUDA 12.8", + "pytorch_version": "2.9.1+cu128" + }, + "software": { + "framework": "SGLang", + "framework_version": "0.5.6", + "driver_version": "570.133.20", + "runtime_version": "CUDA 12.8", + "os": "Ubuntu 22.04.5 LTS", + "python_version": "3.11.13" + }, + "model": { + "model_id": "Qwen/Qwen2.5-0.5B-Instruct", + "model_revision": "7ae557604adf67be50417f59c2c2f167def9a775", + "model_name": null, + "model_note": "Local ModelScope snapshot for Suite F", + "model_source": "local", + "architecture": "dense", + "parameter_count_b": 0.5, + "precision": "BF16", + "effective_dtype": "bfloat16", + "quantization_method": null, + "model_format": "HuggingFace original" + }, + "task": { + "scenarios_run": [ + "offline", + "online", + "interactive" + ], + "parallelism": { + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "expert_parallel_size": 1, + "data_parallel_size": 1 + }, + "num_runs": 3, + "extra_config": { + "enable_tokenizer_batch_encode": true, + "continuous_decode_steps": 4, + "engine_kwargs": { + "attention_backend": "triton", + "disable_cuda_graph_padding": true, + "enable_piecewise_cuda_graph": true, + "piecewise_cuda_graph_max_tokens": 1024, + "prefill_attention_backend": "fa3", + "enable_tokenizer_batch_encode": true, + "num_continuous_decode_steps": 4 + } + } + }, + "metrics": { + "derived": {}, + "offline": { + "results_by_concurrency": [ + { + "client_concurrency": 4, + "throughput_tokens_per_sec": 46128.3, + "throughput_tokens_per_sec_per_chip": 46128.3, + "throughput_tokens_per_sec_total": 62323.45, + "elapsed_seconds_median": 0.9, + "peak_memory_gb": null, + "power_watts_avg": null, + "power_watts_peak": null, + "oom": false, + "throughput_tokens_per_sec_reliability": { + "n": 3, + "mean": 46035.52, + "std": 177.27, + "cv_pct": 0.39, + "stability": "stable", + "runs": [ + 45831.12, + 46147.16, + 46128.3 + ] + }, + "_throughput_note": "output_only", + "_concurrency_note": "client_concurrency is the number of requests sent simultaneously. The inference engine batches internally; this does not directly set engine parameters like max_num_seqs." + }, + { + "client_concurrency": 16, + "throughput_tokens_per_sec": 46406.6, + "throughput_tokens_per_sec_per_chip": 46406.6, + "throughput_tokens_per_sec_total": 62699.46, + "elapsed_seconds_median": 0.9, + "peak_memory_gb": null, + "power_watts_avg": null, + "power_watts_peak": null, + "oom": false, + "throughput_tokens_per_sec_reliability": { + "n": 3, + "mean": 46306.02, + "std": 204.98, + "cv_pct": 0.44, + "stability": "stable", + "runs": [ + 46070.18, + 46406.6, + 46441.29 + ] + }, + "_throughput_note": "output_only", + "_concurrency_note": "client_concurrency is the number of requests sent simultaneously. The inference engine batches internally; this does not directly set engine parameters like max_num_seqs." + }, + { + "client_concurrency": 64, + "throughput_tokens_per_sec": 46500.82, + "throughput_tokens_per_sec_per_chip": 46500.82, + "throughput_tokens_per_sec_total": 62826.76, + "elapsed_seconds_median": 0.9, + "peak_memory_gb": null, + "power_watts_avg": null, + "power_watts_peak": null, + "oom": false, + "throughput_tokens_per_sec_reliability": { + "n": 3, + "mean": 46028.86, + "std": 890.06, + "cv_pct": 1.93, + "stability": "stable", + "runs": [ + 45002.22, + 46500.82, + 46583.55 + ] + }, + "_throughput_note": "output_only", + "_concurrency_note": "client_concurrency is the number of requests sent simultaneously. The inference engine batches internally; this does not directly set engine parameters like max_num_seqs." + } + ] + }, + "online": { + "sla_ttft_ms": 500, + "max_valid_qps": 40, + "results_by_qps": [ + { + "target_qps": 10, + "achieved_qps": 10.0, + "ttft_ms_p50": 8.39, + "ttft_ms_p90": 9.37, + "ttft_ms_p99": 14.12, + "tpot_ms_p50": 1.94, + "tpot_ms_p90": 2.05, + "tpot_ms_p99": 2.28, + "elapsed_seconds_median": 31.8, + "sla_met": true, + "ttft_ms_p99_reliability": { + "n": 3, + "mean": 17.9, + "std": 7.1, + "cv_pct": 39.67, + "stability": "high-variance", + "runs": [ + 25.95, + 15.2, + 12.54 + ] + } + }, + { + "target_qps": 40, + "achieved_qps": 40.0, + "ttft_ms_p50": 8.38, + "ttft_ms_p90": 11.9, + "ttft_ms_p99": 14.53, + "tpot_ms_p50": 2.4, + "tpot_ms_p90": 2.64, + "tpot_ms_p99": 3.18, + "elapsed_seconds_median": 7.8, + "sla_met": true, + "ttft_ms_p99_reliability": { + "n": 3, + "mean": 14.39, + "std": 0.16, + "cv_pct": 1.12, + "stability": "stable", + "runs": [ + 14.39, + 14.23, + 14.56 + ] + } + } + ] + }, + "interactive": { + "ttft_ms_p50": 7.19, + "ttft_ms_p90": 7.54, + "ttft_ms_p99": 15.06, + "tpot_ms_p50": 1.62, + "tpot_ms_p90": 1.64, + "tpot_ms_p99": 1.71, + "peak_memory_gb": null, + "elapsed_seconds_median": 46.2, + "ttft_ms_p99_reliability": { + "n": 3, + "mean": 23.42, + "std": 22.39, + "cv_pct": 95.62, + "stability": "high-variance", + "runs": [ + 49.2, + 12.23, + 8.83 + ] + } + } + }, + "accuracy": { + "subset_score": 0.41, + "baseline_delta": 0.03, + "valid": true, + "framework": "SGLang", + "precision": "BF16", + "notes": "Integrated accuracy check \u2014 used same SGLang instance as benchmark." + }, + "meta": { + "submitted_by": "ChiangYuhsin", + "submission_type": "individual", + "date": "2026-07-31", + "time": "02:40:05", + "run_id": "1189951b", + "run_name": "nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b", + "flagged": null, + "reproduce_script": "runners/nvidia_sglang_inference_opt_4fd1e59e/runner.py", + "env_info_file": "../env_info.json", + "log_file": "run.log", + "samples_file": "samples.jsonl", + "notes": null, + "benchmark_start_time": "2026-07-31T02:39:52.638757+00:00", + "benchmark_end_time": "2026-07-31T02:40:05.992733+00:00", + "benchmark_elapsed_minutes": 4.7, + "model_load_seconds": 47.1, + "benchmark_elapsed_minutes_note": "Total across ['offline', 'online', 'interactive'] scenarios.", + "scenario_dirs": { + "offline": "results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/offline", + "online": "results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/online", + "interactive": "results/community/nvidia_a100_sxm4_80gbx1_suite_F_nvidia_sglang_inference_opt_4fd1e59e_1189951b/interactive" + } + } +} \ No newline at end of file diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.clang-format-ignore b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.clang-format-ignore new file mode 100644 index 00000000..15c76cc4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.clang-format-ignore @@ -0,0 +1 @@ +sgl-kernel/3rdparty/tensorrt_llm/* diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.devcontainer/Dockerfile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.devcontainer/Dockerfile new file mode 100644 index 00000000..3c7b67ca --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.devcontainer/Dockerfile @@ -0,0 +1,35 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.devcontainer/devcontainer.json b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.devcontainer/devcontainer.json new file mode 100644 index 00000000..338b10fe --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.devcontainer/devcontainer.json @@ -0,0 +1,30 @@ +{ + "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" +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.editorconfig b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.editorconfig new file mode 100644 index 00000000..030a7293 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.editorconfig @@ -0,0 +1,25 @@ +# https://editorconfig.org/ + +root = true + +[*] +charset = utf-8 +end_of_line = lf +indent_style = space +indent_size = 4 +trim_trailing_whitespace = true +insert_final_newline = true + +[*.{json,yaml,yml}] +indent_size = 2 + +[*.md] +indent_size = 2 +x-soft-wrap-text = true + +[*.rst] +indent_size = 4 +x-soft-wrap-text = true + +[Makefile] +indent_style = tab diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/CI_PERMISSIONS.json b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/CI_PERMISSIONS.json new file mode 100644 index 00000000..79d99a23 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/CI_PERMISSIONS.json @@ -0,0 +1,989 @@ +{ + "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 + }, + "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": 60, + "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 + }, + "Oasis-Git": { + "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 + }, + "ShangmingCai": { + "can_tag_run_ci_label": true, + "can_rerun_failed_ci": true, + "cooldown_interval_minutes": 0, + "reason": "top contributor", + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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": 60, + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "sglang-bot": { + "can_tag_run_ci_label": true, + "can_rerun_failed_ci": true, + "cooldown_interval_minutes": 0, + "reason": "top contributor", + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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 + } +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/CODEOWNERS b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/CODEOWNERS new file mode 100644 index 00000000..1f009435 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/CODEOWNERS @@ -0,0 +1,47 @@ +.github @merrymercy @Fridge003 @ispobock @Kangyan-Zhou +/docker @Fridge003 @ispobock @HaiShaw @ishandhanani +/docker/npu.Dockerfile @ping1jing2 @iforgetmyname +/python/pyproject.toml @merrymercy @Fridge003 @ispobock +/python/sglang/multimodal_gen @mickqian @yhyang201 +/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/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/layers @merrymercy @Ying1123 @Fridge003 @ispobock @HaiShaw @ch-wan @BBuf @Edwardf0t1 +/python/sglang/srt/layers/quantization @ch-wan @BBuf @Edwardf0t1 @FlamingoPg @AniZpZ +/python/sglang/srt/layers/attention/ascend_backend.py @ping1jing2 @iforgetmyname +/python/sglang/srt/lora @Ying1123 @Fridge003 @lifuhuang +/python/sglang/srt/managers @merrymercy @Ying1123 @hnyls2002 @xiezhq-hermann @zhyncs +/python/sglang/srt/mem_cache @merrymercy @Ying1123 @hnyls2002 @xiezhq-hermann +/python/sglang/srt/mem_cache/allocator_ascend.py @ping1jing2 @iforgetmyname +/python/sglang/srt/model_executor @merrymercy @Ying1123 @hnyls2002 @Fridge003 @ispobock +/python/sglang/srt/model_executor/npu_graph_runner.py @ping1jing2 @iforgetmyname +/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @hebiao064 +/python/sglang/srt/multimodal @mickqian @JustinTong0323 @yhyang201 +/python/sglang/srt/speculative @Ying1123 @merrymercy @hnyls2002 +/sgl-kernel @zhyncs @ispobock @BBuf @yizhang2077 @merrymercy @FlamingoPg @HaiShaw +/sgl-router @slin1237 @CatherineSue +/sgl-router/benches @slin1237 +/sgl-router/bindings/python @CatherineSue @key4ng @slin1237 +/sgl-router/py_test @CatherineSue @key4ng +/sgl-router/src/config @slin1237 +/sgl-router/src/core @slin1237 +/sgl-router/src/data_connector @key4ng +/sgl-router/src/grpc_client @CatherineSue @slin1237 +/sgl-router/src/mcp @key4ng @slin1237 +/sgl-router/src/policies @slin1237 @ByronHsu +/sgl-router/src/proto @CatherineSue @slin1237 +/sgl-router/src/protocols @CatherineSue @key4ng +/sgl-router/src/reasoning_parser @CatherineSue +/sgl-router/src/routers @CatherineSue @key4ng @slin1237 +/sgl-router/src/tokenizer @slin1237 @CatherineSue +/sgl-router/src/tool_parser @slin1237 @CatherineSue +/test/srt/ascend @ping1jing2 @iforgetmyname +/test/srt/test_modelopt* @Edwardf0t1 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/FOLDER_README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/FOLDER_README.md new file mode 100644 index 00000000..ccbf94ec --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/FOLDER_README.md @@ -0,0 +1,12 @@ +# 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/ISSUE_TEMPLATE/1-bug-report.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/ISSUE_TEMPLATE/1-bug-report.yml new file mode 100644 index 00000000..6e3d9a83 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/ISSUE_TEMPLATE/1-bug-report.yml @@ -0,0 +1,35 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/ISSUE_TEMPLATE/2-feature-request.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/ISSUE_TEMPLATE/2-feature-request.yml new file mode 100644 index 00000000..99f1f4d5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/ISSUE_TEMPLATE/2-feature-request.yml @@ -0,0 +1,23 @@ +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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/MAINTAINER.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/MAINTAINER.md new file mode 100644 index 00000000..634f11cd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/MAINTAINER.md @@ -0,0 +1,67 @@ +# 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.ai/): #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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/labeler.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/labeler.yml new file mode 100644 index 00000000..5151e5e2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/labeler.yml @@ -0,0 +1,110 @@ +# Configuration for the GitHub Labeler action +# Automatically adds labels to PRs based on the files changed + +# Router specific (Rust code in sgl-router) +model-gateway: + - changed-files: + - any-glob-to-any-file: 'sgl-router/**/*' + +# Kernel specific +sgl-kernel: + - changed-files: + - any-glob-to-any-file: 'sgl-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/**/*' diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/pull_request_template.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/pull_request_template.md new file mode 100644 index 00000000..7696bc51 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/pull_request_template.md @@ -0,0 +1,26 @@ + + +## 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). +- [ ] Work with maintainers to merge your PR. See the [PR Merge Process](https://github.com/sgl-project/sglang/blob/main/.github/MAINTAINER.md#pull-request-merge-process) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/update_ci_permission.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/update_ci_permission.py new file mode 100644 index 00000000..2ed84667 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/update_ci_permission.py @@ -0,0 +1,196 @@ +""" +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 +""" + +import json +import os +from collections import Counter +from datetime import datetime, timedelta, timezone + +import requests + +# Configuration +REPO_OWNER = "sgl-project" +REPO_NAME = "sglang" +FILE_NAME = "CI_PERMISSIONS.json" +GH_TOKEN = os.getenv("GH_TOKEN") + +if not GH_TOKEN: + raise ValueError("Error: GH_TOKEN environment variable is not set.") + +HEADERS = { + "Authorization": f"Bearer {GH_TOKEN}", + "Accept": "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28", +} + + +def github_api_get(endpoint, params=None): + """Helper to make paginated GitHub API requests.""" + 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 main(): + # 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/auto-format.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/auto-format.yml new file mode 100644 index 00000000..7466572a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/auto-format.yml @@ -0,0 +1,71 @@ +name: Auto Format Code + +on: + pull_request: + types: [labeled] + +permissions: + contents: write + pull-requests: write + +jobs: + auto-format: + if: github.event.label.name == 'format' + runs-on: ubuntu-latest + steps: + - name: Checkout PR branch + uses: actions/checkout@v4 + with: + ref: ${{ github.event.pull_request.head.ref }} + repository: ${{ github.event.pull_request.head.repo.full_name }} + token: ${{ secrets.GITHUB_TOKEN }} + fetch-depth: 0 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + + - name: Install pre-commit hook + run: | + python -m pip install pre-commit + pre-commit install + + - name: Run pre-commit to format code + run: SKIP=no-commit-to-branch pre-commit run --all-files + continue-on-error: true + + - name: Check for changes + id: check_changes + run: | + if [[ -n $(git status -s) ]]; then + echo "has_changes=true" >> $GITHUB_OUTPUT + else + echo "has_changes=false" >> $GITHUB_OUTPUT + fi + + - name: Commit and push changes + if: steps.check_changes.outputs.has_changes == 'true' + run: | + git config --local user.email "github-actions[bot]@users.noreply.github.com" + git config --local user.name "github-actions[bot]" + git add . + git commit -m "🤖 Auto-format code with isort, black, ruff, and clang-format" + git push + + - name: Remove format label + if: always() + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + try { + await github.rest.issues.removeLabel({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: context.issue.number, + name: 'format' + }); + } catch (error) { + console.log('Label may have already been removed'); + } diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/auto-tune.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/auto-tune.yml new file mode 100644 index 00000000..0afc79bb --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/auto-tune.yml @@ -0,0 +1,10 @@ +name: Auto tune + +on: + workflow_dispatch: + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/bot-bump-kernel-version-to-sglang.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/bot-bump-kernel-version-to-sglang.yml new file mode 100644 index 00000000..1621c924 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/bot-bump-kernel-version-to-sglang.yml @@ -0,0 +1,68 @@ +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: + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/bot-bump-kernel-version.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/bot-bump-kernel-version.yml new file mode 100644 index 00000000..91a808c6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/bot-bump-kernel-version.yml @@ -0,0 +1,50 @@ +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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/bot-bump-sglang-version.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/bot-bump-sglang-version.yml new file mode 100644 index 00000000..cbfd8e76 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/bot-bump-sglang-version.yml @@ -0,0 +1,61 @@ +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: + needs: bump-sglang-version + uses: ./.github/workflows/nightly-test-nvidia.yml + with: + ref: ${{ needs.bump-sglang-version.outputs.branch_name }} + secrets: inherit diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/cancel-all-pending-pr-test-runs.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/cancel-all-pending-pr-test-runs.yml new file mode 100644 index 00000000..6217542e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/cancel-all-pending-pr-test-runs.yml @@ -0,0 +1,45 @@ +name: Cancel All Pending PR Test Runs + +on: + workflow_dispatch: + inputs: + workflows: + description: 'Space-separated list of workflow filenames to cancel' + required: true + type: string + default: 'pr-test.yml pr-test-xeon.yml' + +permissions: + actions: write # Needed to cancel runs + contents: read # Needed to read repo info + +jobs: + cancel-pending: + runs-on: ubuntu-latest + steps: + - name: Install GitHub CLI + run: sudo apt-get install -y gh jq + + - name: Cancel all pending/waiting runs for specified workflows + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + REPO: ${{ github.repository }} + run: | + # Read the space-separated string from the input into a bash array + WORKFLOW_FILES=(${{ github.event.inputs.workflows }}) + + echo "Targeting ${#WORKFLOW_FILES[@]} workflow(s): ${{ github.event.inputs.workflows }}" + + for workflow_file in "${WORKFLOW_FILES[@]}"; do + echo "--- Checking workflow: $workflow_file ---" + gh run list \ + --repo "$REPO" \ + --workflow "$workflow_file" \ + --json databaseId,status \ + --limit 1000 \ + | jq -r '.[] | select(.status=="queued" or .status=="in_progress") | .databaseId' \ + | while read run_id; do + echo "Cancelling run ID: $run_id for workflow: $workflow_file" + gh run cancel "$run_id" --repo "$REPO" + done + done diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/cancel-pr-workflow-on-merge.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/cancel-pr-workflow-on-merge.yml new file mode 100644 index 00000000..535884ba --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/cancel-pr-workflow-on-merge.yml @@ -0,0 +1,22 @@ +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 }} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/ci-failure-monitor.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/ci-failure-monitor.yml new file mode 100644 index 00000000..89e770ee --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/ci-failure-monitor.yml @@ -0,0 +1,64 @@ +name: CI Failure Monitor + +on: + schedule: + - cron: '*/30 * * * *' # Every 30 minutes + workflow_dispatch: + inputs: + limit: + description: 'Number of workflow runs to analyze (across all workflows)' + required: false + default: '1000' + type: string + threshold: + description: 'Alert threshold for consecutive failures' + required: false + default: '4' + type: string + +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 + + - name: Run Failure Analysis + env: + GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} + PYTHONUNBUFFERED: 1 + PYTHONIOENCODING: utf-8 + run: | + cd scripts/ci_monitor + python ci_failures_analysis.py \ + --token $GITHUB_TOKEN \ + --limit ${{ inputs.limit || '1000' }} \ + --threshold ${{ inputs.threshold || '4' }} \ + --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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/ci-monitor.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/ci-monitor.yml new file mode 100644 index 00000000..28a198a3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/ci-monitor.yml @@ -0,0 +1,111 @@ +name: CI Monitor + +on: + schedule: + - cron: '0 */12 * * *' # Every 12 hours for main analysis + workflow_dispatch: + inputs: + limit: + description: 'Number of CI runs to analyze' + required: false + default: '1000' + type: string + +concurrency: + group: ci-monitor-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: write + actions: read + +jobs: + ci-monitor: + 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.9' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install requests matplotlib pandas + + - name: Run CI Analysis + env: + GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} + PYTHONUNBUFFERED: 1 + PYTHONIOENCODING: utf-8 + run: | + cd scripts/ci_monitor + python ci_analyzer.py --token $GITHUB_TOKEN --limit ${{ inputs.limit || '1000' }} --output ci_analysis_$(date +%Y%m%d_%H%M%S).json + + - name: Run Nightly Test Analysis + env: + GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} + PYTHONUNBUFFERED: 1 + PYTHONIOENCODING: utf-8 + run: | + cd scripts/ci_monitor + python ci_analyzer.py --token $GITHUB_TOKEN --mode nightly --days 2 --output nightly_analysis_$(date +%Y%m%d_%H%M%S).json + + - name: Run Performance Analysis + env: + GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} + PYTHONUNBUFFERED: 1 + PYTHONIOENCODING: utf-8 + run: | + cd scripts/ci_monitor + python ci_analyzer_perf.py --token $GITHUB_TOKEN --limit ${{ inputs.limit || '1000' }} --output-dir performance_tables_$(date +%Y%m%d_%H%M%S) --upload-to-github + + - name: Upload Analysis Results + uses: actions/upload-artifact@v4 + with: + name: ci-analysis-results-${{ github.run_number }} + path: | + scripts/ci_monitor/ci_analysis_*.json + scripts/ci_monitor/nightly_analysis_*.json + scripts/ci_monitor/performance_tables_* + retention-days: 30 + + ci-monitor-balance: + needs: ci-monitor + 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.9' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install requests + + - name: Run Test Balance Analysis + env: + GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} + PYTHONUNBUFFERED: 1 + PYTHONIOENCODING: utf-8 + run: | + cd scripts/ci_monitor + python ci_analyzer_balance.py --token $GITHUB_TOKEN --limit ${{ inputs.limit || '1000' }} --output test_balance_report_$(date +%Y%m%d_%H%M%S).json + + - name: Upload Balance Analysis Results + uses: actions/upload-artifact@v4 + with: + name: test-balance-results-${{ github.run_number }} + path: | + scripts/ci_monitor/test_balance_report_*.json + scripts/ci_monitor/test_balance_report_*.csv + retention-days: 30 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/close-inactive-issues.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/close-inactive-issues.yml new file mode 100644 index 00000000..048e6c44 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/close-inactive-issues.yml @@ -0,0 +1,96 @@ +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(); diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/execute-notebook.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/execute-notebook.yml new file mode 100644 index 00000000..52942c77 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/execute-notebook.yml @@ -0,0 +1,60 @@ +name: Execute Notebooks + +on: + pull_request: + branches: [ main ] + paths: + - "python/sglang/**" + - "docs/**" + workflow_dispatch: + + +concurrency: + group: execute-notebook-${{ github.ref }} + cancel-in-progress: true + + +jobs: + run-all-notebooks: + runs-on: 1-gpu-runner + if: github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'run-ci') + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/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: [ + run-all-notebooks + ] + runs-on: ubuntu-latest + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/labeler.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/labeler.yml new file mode 100644 index 00000000..5509bd41 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/labeler.yml @@ -0,0 +1,20 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/lint.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/lint.yml new file mode 100644 index 00000000..56598470 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/lint.yml @@ -0,0 +1,50 @@ +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.10" + + - 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.18.1 + with: + source: sgl-kernel + extensions: h,c,cpp,hpp,cu,cuh,cc + clangFormatVersion: 18 + style: file + + - name: Check proto files are in sync + run: | + if ! diff -q python/sglang/srt/grpc/sglang_scheduler.proto sgl-router/src/proto/sglang_scheduler.proto; then + echo "❌ ERROR: Proto files are out of sync!" + echo "" + echo "The following files must be kept identical:" + echo " - python/sglang/srt/grpc/sglang_scheduler.proto" + echo " - sgl-router/src/proto/sglang_scheduler.proto" + echo "" + echo "Please ensure both files have the same content." + echo "" + echo "Differences:" + diff python/sglang/srt/grpc/sglang_scheduler.proto sgl-router/src/proto/sglang_scheduler.proto || true + exit 1 + fi diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-release-gateway.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-release-gateway.yml new file mode 100644 index 00000000..7b5226ba --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-release-gateway.yml @@ -0,0 +1,196 @@ +# 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-router folder to root and delete sglang-repo + run: | + mv sglang-repo/sgl-router/* . + 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-router folder to root and delete sglang-repo + run: | + mv sglang-repo/sgl-router/* . + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-test-amd.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-test-amd.yml new file mode 100644 index 00000000..932aafe8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-test-amd.yml @@ -0,0 +1,59 @@ +name: Nightly Test (AMD) + +on: + schedule: + - cron: '0 0 * * *' + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +concurrency: + group: nightly-test-amd-${{ github.ref }} + cancel-in-progress: true + +jobs: + nightly-test: + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + strategy: + matrix: + runner: [linux-mi300-gpu-2, linux-mi325-gpu-2-nightly] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Setup docker + run: | + touch github_summary.md + bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Nightly Test + run: | + bash scripts/ci/amd_ci_exec.sh -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" python3 run_suite.py --suite nightly-amd --timeout-per-file 7200 + echo "$(> $GITHUB_STEP_SUMMARY + + check-all-jobs: + if: always() && (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch') + needs: + - nightly-test + 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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-test-intel.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-test-intel.yml new file mode 100644 index 00000000..b32735dd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-test-intel.yml @@ -0,0 +1,26 @@ +name: Nightly Test (Intel) + +on: + schedule: + - cron: '0 0 * * *' + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +concurrency: + group: nightly-test-intel-${{ github.ref }} + cancel-in-progress: true + +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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-test-nvidia.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-test-nvidia.yml new file mode 100644 index 00000000..37d36439 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/nightly-test-nvidia.yml @@ -0,0 +1,570 @@ +name: Nightly Test (Nvidia) + +on: + schedule: + - cron: '0 0 * * *' + push: + branches: + - main + paths: + - "python/sglang/version.py" + - "test/run_suite.py" + 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-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-${{ github.ref }} + cancel-in-progress: true + +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/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 + + # 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/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 + + # 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 + 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/ci_install_dependency.sh + + - 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: Run Qwen3-235B nightly performance test + 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: "8-gpu-h200" + run: | + rm -rf test/performance_profiles_qwen3_235b/ + cd test + python3 nightly/test_qwen3_235b_perf.py + + - name: Publish Qwen3-235B 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/publish_traces.py --traces-dir test/performance_profiles_qwen3_235b + + - name: Run Kimi-K2-Thinking nightly performance test + 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: "8-gpu-h200" + run: | + rm -rf test/performance_profiles_kimi_k2_thinking/ + cd test + python3 nightly/test_kimi_k2_thinking_perf.py + + - name: Publish Kimi-K2-Thinking 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/publish_traces.py --traces-dir test/performance_profiles_kimi_k2_thinking + + - name: Run GLM-4.6 nightly performance test + 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: "8-gpu-h200" + run: | + rm -rf test/performance_profiles_glm_4_6/ + cd test + IS_BLACKWELL=1 python3 nightly/test_glm_4_6_perf.py + + - name: Publish GLM-4.6 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/publish_traces.py --traces-dir test/performance_profiles_glm_4_6 + + - name: Run MiniMax-M2 nightly performance test + 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: "8-gpu-h200" + run: | + rm -rf test/performance_profiles_minimax_m2/ + cd test + python3 nightly/test_minimax_m2_perf.py + + - name: Publish MiniMax-M2 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/publish_traces.py --traces-dir test/performance_profiles_minimax_m2 + + # 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/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 + + # 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/ci_install_dependency.sh + + - name: Run eval test for text models + timeout-minutes: 120 + run: | + cd test + python3 nightly/test_text_models_gsm8k_eval.py + + # 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/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 nightly/test_text_models_perf.py + + - 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/publish_traces.py --traces-dir test/performance_profiles_text_models + + # 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/ci_install_dependency.sh + + - name: Run eval test for VLM models (fixed MMMU-100) + timeout-minutes: 240 + run: | + cd test + python3 nightly/test_vlms_mmmu_eval.py + + # 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/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 nightly/test_vlms_perf.py + + - 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/publish_traces.py --traces-dir test/performance_profiles_vlms + + # 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/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 }} + + timeout-minutes: 60 + run: | + cd python + python3 sglang/multimodal_gen/test/run_suite.py \ + --suite 1-gpu \ + --partition-id ${{ matrix.part }} \ + --total-partitions 2 + + + 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/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 }} + + timeout-minutes: 60 + run: | + cd python + python3 sglang/multimodal_gen/test/run_suite.py \ + --suite 2-gpu \ + --partition-id ${{ matrix.part }} \ + --total-partitions 2 + + # 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/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 60 + run: | + cd test + python3 run_suite.py --hw cuda --suite nightly-4-gpu-b200 --nightly --continue-on-error + + # B200 Performance tests - 8 GPU + nightly-test-perf-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/ci_install_dependency.sh + + - name: Run DeepSeek v3.1 nightly performance test + 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: "8-gpu-b200" + run: | + rm -rf test/performance_profiles_deepseek_v31/ + cd test + IS_BLACKWELL=1 python3 nightly/test_deepseek_v31_perf.py + + - name: Publish DeepSeek v3.1 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/publish_traces.py --traces-dir test/performance_profiles_deepseek_v31 + + - name: Run DeepSeek v3.2 nightly performance test + if: always() + 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: "8-gpu-b200" + run: | + rm -rf test/performance_profiles_deepseek_v32/ + cd test + IS_BLACKWELL=1 python3 nightly/test_deepseek_v32_perf.py + + - name: Publish DeepSeek v3.2 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/publish_traces.py --traces-dir test/performance_profiles_deepseek_v32 + + - name: Run Kimi-K2-Thinking nightly performance test + if: always() + 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: "8-gpu-b200" + run: | + rm -rf test/performance_profiles_kimi_k2_thinking/ + cd test + IS_BLACKWELL=1 python3 nightly/test_kimi_k2_thinking_perf.py + + - name: Publish Kimi-K2-Thinking 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/publish_traces.py --traces-dir test/performance_profiles_kimi_k2_thinking + + - name: Run Qwen3-235B nightly performance test + if: always() + 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: "8-gpu-b200" + run: | + rm -rf test/performance_profiles_qwen3_235b/ + cd test + IS_BLACKWELL=1 python3 nightly/test_qwen3_235b_perf.py + + - name: Publish Qwen3-235B 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/publish_traces.py --traces-dir test/performance_profiles_qwen3_235b + + - name: Run GLM-4.6 nightly performance test + if: always() + 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: "8-gpu-b200" + run: | + rm -rf test/performance_profiles_glm_4_6/ + cd test + IS_BLACKWELL=1 python3 nightly/test_glm_4_6_perf.py + + - name: Publish GLM-4.6 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/publish_traces.py --traces-dir test/performance_profiles_glm_4_6 + + - name: Run MiniMax-M2 nightly performance test + if: always() + 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: "8-gpu-b200" + run: | + rm -rf test/performance_profiles_minimax_m2/ + cd test + IS_BLACKWELL=1 python3 nightly/test_minimax_m2_perf.py + + - name: Publish MiniMax-M2 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/publish_traces.py --traces-dir test/performance_profiles_minimax_m2 + + # 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-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 + 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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/open-pr-copy-from-oss.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/open-pr-copy-from-oss.yml new file mode 100644 index 00000000..05af6ea4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/open-pr-copy-from-oss.yml @@ -0,0 +1,28 @@ +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.PAT_FOR_CODE_SYNC_FROM_LIANMIN }} + run: | + python3 scripts/code_sync/copy_from_oss.py diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/open-pr-copy-to-oss.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/open-pr-copy-to-oss.yml new file mode 100644 index 00000000..b3bb6aae --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/open-pr-copy-to-oss.yml @@ -0,0 +1,31 @@ +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.PAT_FOR_CODE_SYNC_FROM_LIANMIN }} + run: | + python3 scripts/code_sync/copy_to_oss.py --commit ${{ github.event.inputs.commit_sha }} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-benchmark-rust.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-benchmark-rust.yml new file mode 100644 index 00000000..0b98b774 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-benchmark-rust.yml @@ -0,0 +1,321 @@ +name: PR Benchmark (SMG Components) + +on: + push: + branches: [ main ] + paths: + - "sgl-router/**" + pull_request: + branches: [ main ] + paths: + - "sgl-router/**" + 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: + # Quick check job that always runs on PRs + 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/ci_install_rust.sh + + - name: Configure sccache + uses: mozilla-actions/sccache-action@v0.0.9 + with: + version: "v0.10.0" + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + # Share cache across all benchmark jobs + shared-key: "rust-cache" + # Save cache even on failure + save-if: true + cache-all-crates: true + cache-on-failure: true + + - name: Check benchmarks compile + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + cargo check --benches + + - name: Show sccache stats + if: always() + run: sccache --show-stats + + # Full benchmark jobs that only run with label or on main branch + benchmark-request-processing: + name: Request Processing Benchmark + 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'))) + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + # Fetch enough history for baseline comparison + fetch-depth: 100 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Configure sccache + uses: mozilla-actions/sccache-action@v0.0.9 + with: + version: "v0.10.0" + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + # Share cache across all benchmark jobs + shared-key: "rust-cache" + cache-all-crates: true + cache-on-failure: true + # Save cache even on failure + save-if: true + + - name: Run request processing benchmark + timeout-minutes: 30 + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + # Try to use sccache, but disable if it fails + if command -v sccache &> /dev/null; then + echo "Testing sccache availability..." + # Try to start sccache and check if it works + 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 + # Run only the summary benchmark for quick validation in PRs + cargo bench --bench request_processing -- benchmark_summary --exact + + - name: Upload benchmark results + if: always() + uses: actions/upload-artifact@v4 + with: + name: request-processing-results-${{ github.sha }} + path: | + sgl-router/target/criterion/benchmark_summary/ + retention-days: 30 + + - name: Show sccache stats + if: always() + run: sccache --show-stats + + benchmark-tokenizer: + name: Tokenizer Benchmark + 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'))) + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 100 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Configure sccache + uses: mozilla-actions/sccache-action@v0.0.9 + with: + version: "v0.10.0" + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + # Share cache across all benchmark jobs + shared-key: "rust-cache" + cache-all-crates: true + cache-on-failure: true + # Save cache even on failure + save-if: true + + - name: Run tokenizer benchmark + timeout-minutes: 30 + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + # Try to use sccache, but disable if it fails + if command -v sccache &> /dev/null; then + echo "Testing sccache availability..." + # Try to start sccache and check if it works + 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 tokenizer_benchmark + + - name: Upload benchmark results + if: always() + uses: actions/upload-artifact@v4 + with: + name: tokenizer-results-${{ github.sha }} + path: | + sgl-router/target/criterion/tokenizer*/ + retention-days: 30 + + benchmark-tool-parser: + name: Tool Parser Benchmark + 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'))) + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 100 + + - name: Install dependencies + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Configure sccache + uses: mozilla-actions/sccache-action@v0.0.9 + with: + version: "v0.10.0" + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + # Share cache across all benchmark jobs + shared-key: "rust-cache" + cache-all-crates: true + cache-on-failure: true + # Save cache even on failure + save-if: true + + - name: Run tool parser benchmark + timeout-minutes: 30 + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + # Try to use sccache, but disable if it fails + if command -v sccache &> /dev/null; then + echo "Testing sccache availability..." + # Try to start sccache and check if it works + 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 tool_parser_benchmark + + - name: Upload benchmark results + if: always() + uses: actions/upload-artifact@v4 + with: + name: tool-parser-results-${{ github.sha }} + path: | + sgl-router/target/criterion/tool_parser*/ + retention-days: 30 + + - name: Show sccache stats + if: always() + run: sccache --show-stats + + benchmark-summary: + name: Benchmark Summary + needs: [benchmark-request-processing, benchmark-tokenizer, benchmark-tool-parser] + 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: | + echo "## Benchmark Results Summary" > summary.md + echo "" >> summary.md + echo "### Request Processing" >> summary.md + if [ -d "benchmark-results/request-processing-results-${{ github.sha }}" ]; then + echo "✅ Completed" >> summary.md + else + echo "❌ Failed or skipped" >> summary.md + fi + echo "" >> summary.md + echo "### Tokenizer" >> summary.md + if [ -d "benchmark-results/tokenizer-results-${{ github.sha }}" ]; then + echo "✅ Completed" >> summary.md + else + echo "❌ Failed or skipped" >> summary.md + fi + echo "" >> summary.md + echo "### Tool Parser" >> summary.md + if [ -d "benchmark-results/tool-parser-results-${{ github.sha }}" ]; then + echo "✅ Completed" >> summary.md + else + echo "❌ Failed or skipped" >> summary.md + fi + cat summary.md + + - name: Upload summary + uses: actions/upload-artifact@v4 + with: + name: benchmark-summary-${{ github.sha }} + path: summary.md + retention-days: 30 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-gate.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-gate.yml new file mode 100644 index 00000000..cffc8f5d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-gate.yml @@ -0,0 +1,173 @@ +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'." + 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 || []; + const recentFound = runs.find((run) => { + if (String(run.id) === String(context.runId)) return false; + if (new Date(run.created_at) < cutoff) return false; + return (run.actor?.login === triggeringActor) || (run.triggering_actor?.login === triggeringActor); + }); + + 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.` + ); + } diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-amd.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-amd.yml new file mode 100644 index 00000000..8a48f7cb --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-amd.yml @@ -0,0 +1,451 @@ +name: PR Test (AMD) + +on: + push: + branches: [ main ] + paths: + - "python/**" + - "!python/sglang/multimodal_gen/**" + - "scripts/ci/**" + - "test/**" + - "sgl-kernel/**" + - ".github/workflows/pr-test-amd.yml" + pull_request: + branches: [ main ] + paths: + - "python/**" + - "!python/sglang/multimodal_gen/**" + - "scripts/ci/**" + - "test/**" + - "sgl-kernel/**" + - ".github/workflows/pr-test-amd.yml" + workflow_dispatch: + +concurrency: + group: pr-test-amd-${{ github.ref }} + cancel-in-progress: true + +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 }} + sgl_kernel: ${{ steps.filter.outputs.sgl_kernel }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + - name: Detect file changes + id: filter + uses: dorny/paths-filter@v3 + with: + filters: | + main_package: + - "python/**" + - "scripts/ci/**" + - "test/**" + - ".github/workflows/pr-test-amd.yml" + sgl_kernel: + - "sgl-kernel/**" + + # =============================================== sgl-kernel ==================================================== + sgl-kernel-unit-test-amd: + needs: [check-changes] + if: needs.check-changes.outputs.sgl_kernel == 'true' + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: | + bash scripts/ci/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_kvcacheio.py + + # =============================================== primary ==================================================== + + stage-a-test-1-amd: + needs: [check-changes] + if: always() && !failure() && !cancelled() && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: | + bash scripts/ci/amd_ci_install_dependency.sh + + - name: Run test + timeout-minutes: 10 + run: | + docker exec -w /sglang-checkout/test ci_sglang python3 run_suite.py --hw amd --suite stage-a-test-1 + + unit-test-backend-1-gpu-amd: + needs: [check-changes, stage-a-test-1-amd] + if: always() && !failure() && !cancelled() && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1] + part: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Run test + timeout-minutes: 30 + run: | + bash scripts/ci/amd_ci_exec.sh python3 run_suite.py --suite per-commit-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 12 + + unit-test-backend-2-gpu-amd: + needs: [check-changes, stage-a-test-1-amd] + if: always() && !failure() && !cancelled() && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-2] + part: [0, 1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Run test + timeout-minutes: 30 + run: | + bash scripts/ci/amd_ci_exec.sh python3 run_suite.py --suite per-commit-2-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 + + unit-test-backend-8-gpu-amd: + needs: [check-changes, unit-test-backend-2-gpu-amd] + if: always() && !failure() && !cancelled() && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + env: + RUNNER_LABELS: linux-mi300-gpu-8 + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-8] + part: [0, 1, 2] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Run test + timeout-minutes: 60 + run: | + bash scripts/ci/amd_ci_exec.sh python3 run_suite.py --suite per-commit-8-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 --timeout-per-file 3600 + + performance-test-1-gpu-part-1-amd: + needs: [check-changes, stage-a-test-1-amd] + if: always() && !failure() && !cancelled() && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Benchmark single latency + timeout-minutes: 20 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_bs1_small + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_bs1_default + + - name: Benchmark online latency + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_online_latency_default + + - name: Benchmark offline throughput + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_default + + - name: Benchmark offline throughput (Non-streaming, small batch size) + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_non_stream_small_batch_size + + performance-test-1-gpu-part-2-amd: + needs: [check-changes, stage-a-test-1-amd] + if: always() && !failure() && !cancelled() && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Benchmark offline throughput (w/o RadixAttention) + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_without_radix_cache + + - name: Benchmark offline throughput (w/ Triton) + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_with_triton_attention_backend + + - name: Benchmark offline throughput (w/ FP8) + timeout-minutes: 15 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_default_fp8 + + performance-test-2-gpu-amd: + needs: [check-changes, unit-test-backend-2-gpu-amd] + if: always() && !failure() && !cancelled() && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-2] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Benchmark dummy grok (TP=2) + timeout-minutes: 30 + run: | + bash scripts/ci/amd_ci_exec.sh python3 models/test_dummy_grok_models.py + + - name: Benchmark single latency (TP=2) + timeout-minutes: 25 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_moe_tp2_bs1 + + - name: Benchmark single latency + torch.compile (TP=2) + timeout-minutes: 25 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_torch_compile_tp2_bs1 + + - name: Benchmark offline throughput (TP=2) + timeout-minutes: 25 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_moe_offline_throughput_default + + - name: Benchmark offline throughput (w/o RadixAttention) (TP=2) + timeout-minutes: 25 + run: | + bash scripts/ci/amd_ci_exec.sh python3 -m unittest test_bench_serving.TestBenchServing.test_moe_offline_throughput_without_radix_cache + + accuracy-test-1-gpu-amd: + needs: [check-changes, stage-a-test-1-amd] + if: always() && !failure() && !cancelled() && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Evaluate Accuracy + timeout-minutes: 30 + run: | + bash scripts/ci/amd_ci_exec.sh -e SGLANG_USE_AITER=0 python3 test_eval_accuracy_large.py + bash scripts/ci/amd_ci_exec.sh python3 test_eval_fp8_accuracy.py + bash scripts/ci/amd_ci_exec.sh python3 models/test_qwen_models.py + + accuracy-test-2-gpu-amd: + needs: [check-changes, accuracy-test-1-gpu-amd] + if: always() && !failure() && !cancelled() && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + strategy: + fail-fast: false + matrix: + runner: [linux-mi300-gpu-2] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: bash scripts/ci/amd_ci_install_dependency.sh + + - name: Evaluate accuracy (TP=2) + timeout-minutes: 30 + run: | + bash scripts/ci/amd_ci_exec.sh -e SGLANG_USE_AITER=0 python3 test_moe_eval_accuracy_large.py + + pr-test-amd-finish: + needs: + [ + call-gate, + check-changes, + + sgl-kernel-unit-test-amd, + + stage-a-test-1-amd, + unit-test-backend-1-gpu-amd, + unit-test-backend-2-gpu-amd, + unit-test-backend-8-gpu-amd, + performance-test-1-gpu-part-1-amd, + performance-test-1-gpu-part-2-amd, + performance-test-2-gpu-amd, + accuracy-test-1-gpu-amd, + accuracy-test-2-gpu-amd, + ] + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-npu.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-npu.yml new file mode 100644 index 00000000..af19e754 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-npu.yml @@ -0,0 +1,205 @@ +name: PR Test (NPU) + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + workflow_dispatch: + +concurrency: + group: pr-test-npu-${{ github.ref }} + cancel-in-progress: true + +jobs: + # ==================== Check Changes ==================== # + check-changes: + runs-on: ubuntu-latest + outputs: + main_package: ${{ steps.filter.outputs.main_package }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Detect file changes + id: filter + uses: dorny/paths-filter@v3 + with: + filters: | + main_package: + - "python/sglang/!(multimodal_gen)/**" + - "python/*.toml" + - "scripts/ci/npu_ci_install_dependency.sh" + - "test/srt/ascend/**" + - ".github/workflows/pr-test-npu.yml" + + # ==================== PR Gate ==================== # + pr-gate: + needs: check-changes + if: needs.check-changes.outputs.main_package == '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-arm64-npu-1 + container: + image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.3.rc1-910b-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_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://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 test/srt + python3 run_suite.py --suite per-commit-1-npu-a2 + + per-commit-2-npu-a2: + needs: [check-changes, pr-gate] + if: needs.check-changes.outputs.main_package == 'true' + runs-on: linux-arm64-npu-2 + strategy: + fail-fast: true + matrix: + part: [0, 1, 2] + container: + image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.3.rc1-910b-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_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://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 test/srt + python3 run_suite.py --suite per-commit-2-npu-a2 --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 + + per-commit-4-npu-a2: + needs: [check-changes, pr-gate] + if: needs.check-changes.outputs.main_package == 'true' + runs-on: linux-arm64-npu-4 + container: + image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.3.rc1-910b-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_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://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 test/srt + python3 run_suite.py --suite per-commit-4-npu-a2 --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 + strategy: + fail-fast: true + matrix: + part: [0, 1] + container: + image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.3.rc1-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_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: | + export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}" + cd test/srt + python3 run_suite.py --suite per-commit-16-npu-a3 --timeout-per-file 3600 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-pd-router.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-pd-router.yml new file mode 100644 index 00000000..d95e7153 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-pd-router.yml @@ -0,0 +1,607 @@ +name: PR Benchmark (SMG PD Router) + +on: + push: + branches: [ main ] + paths: + - 'python/sglang/srt/disaggregation/**' + - 'scripts/ci/ci_start_disaggregation_servers.sh' + - 'sgl-router/**' + pull_request: + branches: [ main ] + types: [opened, synchronize, reopened, labeled] + paths: + - 'python/sglang/srt/disaggregation/**' + - 'scripts/ci/ci_start_disaggregation_servers.sh' + - 'sgl-router/**' + workflow_dispatch: + +concurrency: + group: test-disaggregation-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + pull-requests: write + issues: write + +jobs: + test-disaggregation: + if: | + github.event_name != 'pull_request' || + (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci') && contains(github.event.pull_request.labels.*.name, 'router-benchmark')) || + (github.event.action == 'labeled' && github.event.label.name == 'run-ci' && contains(github.event.pull_request.labels.*.name, 'router-benchmark')) || + (github.event.action == 'labeled' && github.event.label.name == 'router-benchmark' && contains(github.event.pull_request.labels.*.name, 'run-ci')) + runs-on: [8-gpu-h200-oracle] + timeout-minutes: 45 + + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 10 + + - name: Setup Python + uses: actions/setup-python@v4 + with: + python-version: '3.12' + + - name: Setup Rust + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Cache Rust dependencies + uses: actions/cache@v4 + with: + path: | + ~/.cargo/bin/ + ~/.cargo/registry/index/ + ~/.cargo/registry/cache/ + ~/.cargo/git/db/ + sgl-router/target/ + key: ${{ runner.os }}-cargo-${{ hashFiles('sgl-router/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-cargo- + + - name: Cache pip dependencies + uses: actions/cache@v4 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ hashFiles('python/pyproject.toml') }} + restore-keys: | + ${{ runner.os }}-pip- + + - name: Validate environment + run: | + echo "=== System Validation ===" + nvidia-smi + echo "GPU count: $(nvidia-smi -L | wc -l)" + if [ $(nvidia-smi -L | wc -l) -lt 8 ]; then + echo "Error: This test requires at least 8 GPUs" + exit 1 + fi + + echo "=== GPU Process Check ===" + # Fail fast if any GPU compute processes are active + if command -v nvidia-smi >/dev/null 2>&1; then + # Try to query compute apps first (preferred and concise) + gpu_procs=$(nvidia-smi --query-compute-apps=pid,process_name,gpu_uuid --format=csv,noheader 2>/dev/null | sed '/^$/d' || true) + + # Fallback to detailed PIDS report if the query returns nothing but there might still be processes + if [ -z "$gpu_procs" ]; then + gpu_procs=$(nvidia-smi -q -d PIDS 2>/dev/null | awk '/Processes/{flag=1;next}/^$/{flag=0}flag' | sed '/^\s*Processes:/d' | sed '/^\s*$/d' || true) + fi + + if [ -n "$gpu_procs" ]; then + echo "Error: Found active GPU processes using the device(s):" + echo "$gpu_procs" + exit 1 + else + echo "No active GPU compute processes detected." + fi + else + echo "Error: nvidia-smi not found; skipping GPU process check." + exit 1 + fi + + echo "=== RDMA Validation ===" + if ! command -v ibv_devices >/dev/null 2>&1; then + echo "Error: InfiniBand tools not found" + exit 1 + fi + + # Check for active IB devices + found_active_device=false + for device in mlx5_{0..11}; do + if ibv_devinfo $device >/dev/null 2>&1; then + state=$(ibv_devinfo $device | grep "state:" | head -1 | awk '{print $2}') + if [[ "$state" == "PORT_ACTIVE" ]]; then + echo "✓ Found active device: $device" + found_active_device=true + break + fi + fi + done + + if [ "$found_active_device" = false ]; then + echo "Error: No active IB devices found" + echo "Available devices:" + ibv_devices || true + exit 1 + fi + + echo "=== Model Validation ===" + if [ ! -d "/raid/models/meta-llama/Llama-3.1-8B-Instruct" ]; then + echo "Error: Model not found" + ls -la /raid/models/ || echo "No models directory" + exit 1 + fi + echo "✓ Model found" + + - name: Install SGLang dependencies + run: | + echo "Installing SGLang with all extras..." + python3 -m pip --no-cache-dir install --upgrade pip + python3 -m pip --no-cache-dir install torch==2.9.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu128 + python3 -m pip --no-cache-dir install -e "python[all]" --break-system-packages + python3 -m pip --no-cache-dir install mooncake-transfer-engine==0.3.7.post2 + python3 -m pip --no-cache-dir install --user --force-reinstall genai-bench==0.0.2 + + - name: Build and install sgl-router + run: | + source "$HOME/.cargo/env" + echo "Building sgl-router..." + cd sgl-router/bindings/python + pip install maturin + maturin build --release --out dist --features vendored-openssl + pip install --force-reinstall dist/*.whl + + - name: Start disaggregation servers + id: start_servers + run: | + echo "Starting disaggregation servers..." + READY_FILE=".disagg_ready" + rm -f "$READY_FILE" + DISAGG_READY_FILE="$READY_FILE" bash scripts/ci/ci_start_disaggregation_servers.sh & + SERVER_PID=$! + echo "server_pid=$SERVER_PID" >> $GITHUB_OUTPUT + + # Wait until script signals readiness (8/8 healthy) or timeout + TIMEOUT=300 + ELAPSED=0 + while [ $ELAPSED -lt $TIMEOUT ]; do + if [ -f "$READY_FILE" ]; then + echo "✓ All disaggregation servers are healthy (signal detected)" + break + fi + if ! ps -p $SERVER_PID > /dev/null; then + echo "Error: server bootstrap script exited prematurely" + exit 1 + fi + sleep 5 + ELAPSED=$((ELAPSED + 5)) + done + if [ $ELAPSED -ge $TIMEOUT ]; then + echo "❌ Timeout waiting for disaggregation servers to be healthy" + exit 1 + fi + + echo "✓ Servers started (PID: $SERVER_PID)" + + + - name: Test all policies sequentially + timeout-minutes: 30 + run: | + POLICIES=("random" "round_robin" "cache_aware" "power_of_two") + BASE_URL="http://127.0.0.9:8000" + + # Free commonly used ports for router and metrics + echo "Freeing ports 29000 (metrics) and 8000 (API), if in use..." + fuser -k -n tcp 29000 2>/dev/null || true + fuser -k -n tcp 8000 2>/dev/null || true + sleep 1 + + for policy in "${POLICIES[@]}"; do + echo "" + echo "==================================================" + echo "Testing policy: $policy" + echo "==================================================" + + # Free ports before starting router + fuser -k -n tcp 29000 2>/dev/null || true + fuser -k -n tcp 8000 2>/dev/null || true + + # Start router with the current policy + echo "Starting router with policy: $policy..." + RUST_BACKTRACE=1 python3 -m sglang_router.launch_router \ + --pd-disaggregation \ + --policy "$policy" \ + --prefill http://127.0.0.1:30001 9001 \ + --prefill http://127.0.0.2:30002 9002 \ + --prefill http://127.0.0.3:30003 9003 \ + --prefill http://127.0.0.4:30004 9004 \ + --decode http://127.0.0.5:30005 \ + --decode http://127.0.0.6:30006 \ + --decode http://127.0.0.7:30007 \ + --decode http://127.0.0.8:30008 \ + --host 127.0.0.9 \ + --log-level warn \ + --port 8000 & + ROUTER_PID=$! + + # Wait for router to become healthy + echo "Waiting for router to become healthy..." + TIMEOUT=60 + ELAPSED=0 + while [ $ELAPSED -lt $TIMEOUT ]; do + if curl --connect-timeout 5 --silent http://127.0.0.9:8000 > /dev/null 2>&1; then + echo "✓ Router is reachable" + break + fi + if ! ps -p $ROUTER_PID > /dev/null; then + echo "Error: Router process died" + exit 1 + fi + sleep 5 + ELAPSED=$((ELAPSED + 5)) + done + + if [ $ELAPSED -ge $TIMEOUT ]; then + echo "Error: Router health check timeout" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + + # Test API functionality + echo "Testing API completions for $policy..." + response=$(curl -s -X POST "$BASE_URL/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer test-token" \ + -d '{ + "model": "/raid/models/meta-llama/Llama-3.1-8B-Instruct", + "messages": [ + {"role": "user", "content": "Write a Python function to calculate fibonacci numbers recursively"} + ], + "stream": false, + "max_completion_tokens": 100 + }') + + if echo "$response" | jq -e '.choices[0].message.content' > /dev/null 2>&1; then + echo "✓ API test passed for $policy" + else + echo "✗ API test failed for $policy: $response" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + + # Test streaming + echo "Testing streaming API for $policy..." + stream_response=$(timeout 30 curl -s -X POST "$BASE_URL/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer test-token" \ + -d '{ + "model": "/raid/models/meta-llama/Llama-3.1-8B-Instruct", + "messages": [ + {"role": "user", "content": "Count from 1 to 5"} + ], + "stream": true, + "max_completion_tokens": 50 + }') + + if echo "$stream_response" | grep -q "data:"; then + echo "✓ Streaming API test passed for $policy" + else + echo "✗ Streaming API test failed for $policy" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + + # Run genai-bench benchmark + echo "Running genai-bench for $policy..." + genai-bench benchmark \ + --api-backend openai \ + --api-base "http://127.0.0.9:8000" \ + --api-key "dummy-token" \ + --api-model-name "/raid/models/meta-llama/Llama-3.1-8B-Instruct" \ + --model-tokenizer /raid/models/meta-llama/Llama-3.1-8B-Instruct \ + --task text-to-text \ + --num-concurrency 64 \ + --traffic-scenario "D(8000,2000)" \ + --max-requests-per-run 1000 \ + --max-time-per-run 5 \ + --experiment-folder-name "benchmark_${policy}" \ + --experiment-base-dir "." + + # Find the actual experiment folder + actual_folder=$(find . -maxdepth 1 -name "benchmark_${policy}" -type d | head -1) + + if [ -n "$actual_folder" ]; then + # Extract metrics from the Excel summary or JSON files + summary_file="$actual_folder"/*_summary.xlsx + json_files=$(find "$actual_folder" -name "*.json" | grep -v experiment_metadata) + + echo "Genai-bench results saved in: $actual_folder" + + # Extract mean values and validate performance thresholds + echo "📊 Extracting performance metrics for $policy..." + + # Find JSON files excluding experiment metadata + json_files=$(find "$actual_folder" -name "*.json" | grep -v experiment_metadata) + + if [ -n "$json_files" ]; then + # Extract metrics using jq and validate against loose thresholds + for json_file in $json_files; do + echo "Processing: $(basename "$json_file")" + + # Extract mean values for performance validation + ttft_mean=$(jq -r '.aggregated_metrics.stats.ttft.mean' "$json_file") + e2e_latency_mean=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean' "$json_file") + input_throughput_mean=$(jq -r '.aggregated_metrics.stats.input_throughput.mean' "$json_file") + output_throughput_mean=$(jq -r '.aggregated_metrics.stats.output_throughput.mean' "$json_file") + + echo " TTFT mean: ${ttft_mean}s" + echo " E2E Latency mean: ${e2e_latency_mean}s" + echo " Input Throughput mean: ${input_throughput_mean} tokens/s" + echo " Output Throughput mean: ${output_throughput_mean} tokens/s" + + # Set mean thresholds (allowing for reasonable variance) + # These can be adjusted based on your performance requirements + ttft_threshold=4.7 # Max 4.7 seconds for mean TTFT + e2e_latency_threshold=35.0 # Max 35.0 seconds for mean E2E latency + input_throughput_threshold=10000 # Min 02000 tokens/s for mean input throughput + output_throughput_threshold=68 # Min 68 tokens/s for mean output throughput + + + # Validate mean thresholds + validation_passed=true + + if (( $(echo "$ttft_mean > $ttft_threshold" | bc -l) )); then + echo "❌ TTFT validation failed: $ttft_mean > $ttft_threshold" + validation_passed=false + fi + + if (( $(echo "$e2e_latency_mean > $e2e_latency_threshold" | bc -l) )); then + echo "❌ E2E Latency validation failed: $e2e_latency_mean > $e2e_latency_threshold" + validation_passed=false + fi + + if (( $(echo "$input_throughput_mean < $input_throughput_threshold" | bc -l) )); then + echo "❌ Input Throughput validation failed: $input_throughput_mean < $input_throughput_threshold" + validation_passed=false + fi + + if (( $(echo "$output_throughput_mean < $output_throughput_threshold" | bc -l) )); then + echo "❌ Output Throughput validation failed: $output_throughput_mean < $output_throughput_threshold" + validation_passed=false + fi + + if [ "$validation_passed" = true ]; then + echo "✅ Performance validation passed for $policy" + else + echo "❌ Performance validation failed for $policy" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + done + + echo "✓ Genai-bench completed successfully for $policy" + echo "📊 Detailed metrics and plots available in: $actual_folder" + else + echo "✗ Benchmark failed for $policy: No JSON results found" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + else + echo "✗ Benchmark failed for $policy: Experiment folder not found" + kill $ROUTER_PID 2>/dev/null || true + exit 1 + fi + + # Stop router before testing next policy + echo "Stopping router for $policy..." + # First try graceful shutdown + kill $ROUTER_PID 2>/dev/null || true + + # Wait up to 5 seconds for graceful shutdown + for i in {1..5}; do + if ! ps -p $ROUTER_PID > /dev/null 2>&1; then + echo "Router stopped gracefully" + break + fi + sleep 1 + done + + # Force kill if still running + if ps -p $ROUTER_PID > /dev/null 2>&1; then + echo "Force killing router..." + kill -9 $ROUTER_PID 2>/dev/null || true + fi + + # Short delay to ensure port is released + sleep 2 + + echo "✓ Completed testing for $policy" + done + + echo "" + echo "✅ All policies tested successfully!" + + + - name: Upload benchmark results + if: success() + uses: actions/upload-artifact@v4 + with: + name: genai-bench-results-all-policies + path: benchmark_**/ + + - name: Cleanup servers + if: always() + run: | + if [ -n "${{ steps.start_servers.outputs.server_pid }}" ]; then + pkill -P ${{ steps.start_servers.outputs.server_pid }} || true + kill ${{ steps.start_servers.outputs.server_pid }} || true + fi + pkill -f "sglang.launch_server" || true + sleep 5 + remaining=$(ps aux | grep -c "sglang.launch_server" || echo "0") + echo "Cleanup completed. Remaining processes: $remaining" + + summarize-benchmarks: + needs: test-disaggregation + runs-on: ubuntu-latest + if: success() + + steps: + - name: Install jq + run: sudo apt-get update && sudo apt-get install -y jq bc + + - name: Download benchmark results + uses: actions/download-artifact@v4 + with: + name: genai-bench-results-all-policies + + - name: List downloaded contents + run: | + echo "Contents after download:" + ls -la + find . -name "benchmark_*" -type d + echo "JSON files found:" + find . -name "*.json" | head -10 + + - name: Create benchmark summary + run: | + echo "=== DEBUG: Creating benchmark summary ===" + echo "Available benchmark directories:" + find . -name "benchmark_*" -type d + echo "==========================================" + + echo "## PD Router Genai-Bench Results Summary" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "🚀 **Benchmarked with genai-bench for comprehensive LLM serving performance evaluation**" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "| Policy | Status | TTFT (s) | E2E Latency (s) | Input Throughput (tok/s) | Output Throughput (tok/s) |" >> $GITHUB_STEP_SUMMARY + echo "|--------|--------|----------|-----------------|--------------------------|---------------------------|" >> $GITHUB_STEP_SUMMARY + + # First, complete the table with all policies + for policy in random round_robin cache_aware power_of_two; do + # Find genai-bench result folders for this policy (handle zip extraction structure) + result_folder=$(find . -maxdepth 2 -name "benchmark_${policy}" -type d | head -1) + if [ -z "$result_folder" ]; then + # Try alternative patterns in case of different extraction structure + result_folder=$(find . -maxdepth 3 -path "*benchmark_${policy}*" -type d | head -1) + fi + + echo "DEBUG: Policy ${policy} -> Found folder: ${result_folder:-'NOT FOUND'}" + + if [ -n "$result_folder" ] && [ -d "$result_folder" ]; then + # Find JSON file with metrics + json_file=$(find "$result_folder" -name "*.json" -not -name "experiment_metadata.json" | head -1) + + if [ -n "$json_file" ] && [ -f "$json_file" ]; then + # Extract performance metrics + ttft_mean=$(jq -r '.aggregated_metrics.stats.ttft.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + e2e_latency_mean=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + input_throughput_mean=$(jq -r '.aggregated_metrics.stats.input_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + output_throughput_mean=$(jq -r '.aggregated_metrics.stats.output_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + + # Format numbers for display (2 decimal places) + if [ "$ttft_mean" != "N/A" ] && [ "$ttft_mean" != "null" ]; then + ttft_display=$(printf "%.2f" "$ttft_mean" 2>/dev/null || echo "$ttft_mean") + else + ttft_display="N/A" + fi + + if [ "$e2e_latency_mean" != "N/A" ] && [ "$e2e_latency_mean" != "null" ]; then + e2e_display=$(printf "%.2f" "$e2e_latency_mean" 2>/dev/null || echo "$e2e_latency_mean") + else + e2e_display="N/A" + fi + + if [ "$input_throughput_mean" != "N/A" ] && [ "$input_throughput_mean" != "null" ]; then + input_display=$(printf "%.0f" "$input_throughput_mean" 2>/dev/null || echo "$input_throughput_mean") + else + input_display="N/A" + fi + + if [ "$output_throughput_mean" != "N/A" ] && [ "$output_throughput_mean" != "null" ]; then + output_display=$(printf "%.0f" "$output_throughput_mean" 2>/dev/null || echo "$output_throughput_mean") + else + output_display="N/A" + fi + + echo "| ${policy} | ✅ Success | $ttft_display | $e2e_display | $input_display | $output_display |" >> $GITHUB_STEP_SUMMARY + else + echo "| ${policy} | ❌ No Data | N/A | N/A | N/A | N/A |" >> $GITHUB_STEP_SUMMARY + fi + else + echo "| ${policy} | ❌ Failed | N/A | N/A | N/A | N/A |" >> $GITHUB_STEP_SUMMARY + fi + done + + # Add performance validation summary + echo "" >> $GITHUB_STEP_SUMMARY + echo "## 📊 Performance Validation" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "**Thresholds:** TTFT ≤ 2.0s | E2E Latency ≤ 8.0s | Input Throughput ≥ 10,000 tok/s | Output Throughput ≥ 100 tok/s" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + + validation_summary="" + for policy in random round_robin cache_aware power_of_two; do + # Use same robust path finding as above + result_folder=$(find . -maxdepth 2 -name "benchmark_${policy}" -type d | head -1) + if [ -z "$result_folder" ]; then + result_folder=$(find . -maxdepth 3 -path "*benchmark_${policy}*" -type d | head -1) + fi + + if [ -n "$result_folder" ] && [ -d "$result_folder" ]; then + json_file=$(find "$result_folder" -name "*.json" -not -name "experiment_metadata.json" | head -1) + if [ -n "$json_file" ] && [ -f "$json_file" ]; then + # Extract metrics for validation + ttft=$(jq -r '.aggregated_metrics.stats.ttft.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + e2e_latency=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + input_throughput=$(jq -r '.aggregated_metrics.stats.input_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + output_throughput=$(jq -r '.aggregated_metrics.stats.output_throughput.mean // "N/A"' "$json_file" 2>/dev/null || echo "N/A") + + # Check thresholds (using same values as in main workflow) + validation_status="✅" + if [ "$ttft" != "N/A" ] && [ "$ttft" != "null" ]; then + if (( $(echo "$ttft > 4.7" | bc -l 2>/dev/null || echo "0") )); then + validation_status="❌" + fi + fi + if [ "$e2e_latency" != "N/A" ] && [ "$e2e_latency" != "null" ]; then + if (( $(echo "$e2e_latency > 35.0" | bc -l 2>/dev/null || echo "0") )); then + validation_status="❌" + fi + fi + if [ "$input_throughput" != "N/A" ] && [ "$input_throughput" != "null" ]; then + if (( $(echo "$input_throughput < 10000" | bc -l 2>/dev/null || echo "0") )); then + validation_status="❌" + fi + fi + if [ "$output_throughput" != "N/A" ] && [ "$output_throughput" != "null" ]; then + if (( $(echo "$output_throughput < 68" | bc -l 2>/dev/null || echo "0") )); then + validation_status="❌" + fi + fi + + validation_summary="${validation_summary}- **${policy}**: $validation_status\n" + else + validation_summary="${validation_summary}- **${policy}**: ❌ No data\n" + fi + else + validation_summary="${validation_summary}- **${policy}**: ❌ Failed\n" + fi + done + + echo -e "$validation_summary" >> $GITHUB_STEP_SUMMARY + + echo "" >> $GITHUB_STEP_SUMMARY + echo "## 📊 Genai-Bench Features Used" >> $GITHUB_STEP_SUMMARY + echo "- **Token-level Performance**: TTFT, TPOT, End-to-End latency" >> $GITHUB_STEP_SUMMARY + echo "- **Throughput Analysis**: Input/Output/Total token throughput" >> $GITHUB_STEP_SUMMARY + echo "- **Statistical Analysis**: Percentiles, mean, std dev for all metrics" >> $GITHUB_STEP_SUMMARY + echo "- **Visual Reports**: Automated plots and Excel summaries" >> $GITHUB_STEP_SUMMARY + echo "- **SGLang Backend**: Native integration with SGLang serving" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "✅ All policies tested successfully with genai-bench!" >> $GITHUB_STEP_SUMMARY diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-rust.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-rust.yml new file mode 100644 index 00000000..1a67c84f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-rust.yml @@ -0,0 +1,419 @@ +name: PR Test (SMG) + +on: + push: + branches: [ main ] + paths: + - "sgl-router/**" + pull_request: + branches: [ main ] + types: [opened, synchronize, reopened, labeled] + paths: + - "sgl-router/**" + workflow_dispatch: + +concurrency: + group: router-tests-${{ github.ref }} + cancel-in-progress: true + +env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + +jobs: + maturin-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-latest + steps: + - uses: actions/checkout@v4 + with: + path: sglang-repo + + - name: Move sgl-router folder to root + run: | + mv sglang-repo/sgl-router/* . + rm -rf sglang-repo + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.13" + + - name: Install protoc and dependencies + run: | + sudo apt-get update + sudo apt-get install -y wget unzip gcc g++ perl make + cd /tmp + wget https://github.com/protocolbuffers/protobuf/releases/download/v32.0/protoc-32.0-linux-x86_64.zip + sudo unzip protoc-32.0-linux-x86_64.zip -d /usr/local + rm protoc-32.0-linux-x86_64.zip + protoc --version + + - name: Configure sccache + uses: mozilla-actions/sccache-action@v0.0.9 + with: + version: "v0.10.0" + + - name: Test maturin build + uses: PyO3/maturin-action@v1 + with: + working-directory: bindings/python + args: --release --out dist --features vendored-openssl + rust-toolchain: stable + sccache: true + + - name: List built wheel + run: ls -lh bindings/python/dist/ + + - name: Test wheel install + run: | + pip install bindings/python/dist/*.whl + python -c "import sglang_router; print('Python package: OK')" + python -c "from sglang_router.sglang_router_rs import Router; print('Rust extension: OK')" + python -m sglang_router.launch_router --help > /dev/null && echo "Entry point: OK" + router-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/ci_install_rust.sh + + - name: Configure sccache + uses: mozilla-actions/sccache-action@v0.0.9 + with: + version: "v0.10.0" + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + cache-all-crates: true + cache-on-failure: true + + - name: Run lint + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + rustup component add clippy + cargo clippy --all-targets --all-features -- -D warnings + + - name: Run fmt + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + rustup component add --toolchain nightly-x86_64-unknown-linux-gnu rustfmt + rustup toolchain install nightly --profile minimal + cargo +nightly fmt -- --check + + - name: Run Rust tests + timeout-minutes: 20 + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + cargo test + + - name: Check benchmark compilation + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + cargo check --benches + + - name: Quick benchmark sanity check + timeout-minutes: 15 + run: | + source "$HOME/.cargo/env" + cd sgl-router/ + # Run quick benchmarks to ensure they work using Python script + python3 scripts/run_benchmarks.py --quick + + - name: Show sccache stats + if: always() + run: sccache --show-stats + + router-http-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: 4-gpu-a10 + timeout-minutes: 32 + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install rust dependencies + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Configure sccache + uses: mozilla-actions/sccache-action@v0.0.9 + with: + version: "v0.10.0" + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + cache-all-crates: true + cache-on-failure: true + + - name: Install SGLang dependencies + run: | + sudo --preserve-env=PATH bash scripts/ci/ci_install_dependency.sh + + - name: Build python binding + run: | + source "$HOME/.cargo/env" + export RUSTC_WRAPPER=sccache + cd sgl-router/bindings/python + python3 -m pip install --upgrade pip maturin + pip uninstall -y sglang-router + maturin build --profile ci --features vendored-openssl --out dist + pip install dist/*.whl + + - name: Run Python unit tests + run: | + cd sgl-router + source "$HOME/.cargo/env" + python3 -m pip install pytest pytest-cov pytest-xdist + pytest -q py_test/unit --cov=sglang_router --cov-config=bindings/python/.coveragerc --cov-report=term-missing --cov-fail-under=80 + + - name: Run Python integration tests + run: | + cd sgl-router + source "$HOME/.cargo/env" + # Integration tests use FastAPI/uvicorn for mock workers + python3 -m pip install fastapi uvicorn orjson + pytest -q py_test/integration_mock + + - name: Run Python E2E tests + run: | + bash scripts/killall_sglang.sh "nuk_gpus" + cd sgl-router + source "$HOME/.cargo/env" + python3 -m pip --no-cache-dir install --upgrade --ignore-installed blinker + python3 -m pip --no-cache-dir install --upgrade genai-bench==0.0.2 + pytest py_test/e2e_http -s -vv -o log_cli=true --log-cli-level=INFO + + - name: Upload benchmark results + if: success() + uses: actions/upload-artifact@v4 + with: + name: genai-bench-results-all-policies + path: sgl-router/benchmark_**/ + + router-grpc-response-api-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: 4-gpu-a10 + timeout-minutes: 32 + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install rust dependencies + run: | + bash scripts/ci/ci_install_rust.sh + + - name: Configure sccache + uses: mozilla-actions/sccache-action@v0.0.9 + with: + version: "v0.10.0" + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + with: + workspaces: sgl-router + cache-all-crates: true + cache-on-failure: true + + - name: Install SGLang dependencies + run: | + sudo --preserve-env=PATH bash scripts/ci/ci_install_dependency.sh + + - name: Setup Oracle Instant Client + 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 + 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 + 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: Build python binding + run: | + source "$HOME/.cargo/env" + export RUSTC_WRAPPER=sccache + cd sgl-router/bindings/python + python3 -m pip install --upgrade pip maturin + pip uninstall -y sglang-router + maturin build --profile ci --features vendored-openssl --out dist + pip install dist/*.whl + + - name: Run Python E2E response API tests + run: | + python3 -m pip install pytest-rerunfailures + bash scripts/killall_sglang.sh "nuk_gpus" + cd sgl-router + source "$HOME/.cargo/env" + SHOW_ROUTER_LOGS=1 pytest --reruns 3 --reruns-delay 2 py_test/e2e_response_api -s -vv -o log_cli=true --log-cli-level=INFO + + - name: Run Python E2E gRPC tests + run: | + bash scripts/killall_sglang.sh "nuk_gpus" + cd sgl-router + source "$HOME/.cargo/env" + SHOW_ROUTER_LOGS=1 ROUTER_LOCAL_MODEL_PATH="/home/ubuntu/models" pytest py_test/e2e_grpc -s -vv -o log_cli=true --log-cli-level=INFO + + - name: Cleanup Brave MCP Server + if: always() + run: | + docker stop brave-search-server || true + docker rm brave-search-server || true + + - name: Cleanup Oracle Database + if: always() + run: | + docker stop oracle-db || true + docker rm oracle-db || true + + + finish: + needs: [maturin-build-test, router-unit-tests, router-http-tests, router-grpc-response-api-tests] + runs-on: ubuntu-latest + steps: + - name: Finish + run: echo "This is an empty step to ensure that all jobs are completed." + + summarize-benchmarks: + needs: router-http-tests + runs-on: ubuntu-latest + if: success() + + steps: + - name: Install jq + run: sudo apt-get update && sudo apt-get install -y jq bc + + - name: Download benchmark results + uses: actions/download-artifact@v4 + with: + name: genai-bench-results-all-policies + + - name: List downloaded contents + run: | + echo "Contents after download:" + ls -la + find . -name "benchmark_*" -type d + echo "JSON files found:" + find . -name "*.json" | head -10 + + - name: Create benchmark summary + run: | + echo "=== DEBUG: Creating benchmark summary ===" + echo "Available benchmark directories:" + find . -name "benchmark_*" -type d || true + echo "==========================================" + + echo "## Router E2E Genai-Bench Results Summary" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "Results captured from E2E tests for two scenarios: regular router (2 workers, dp=2) and PD router (2 prefill + 2 decode)." >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "| Scenario | Status | TTFT (s) | E2E Latency (s) | Input Throughput (tok/s) | Output Throughput (tok/s) |" >> $GITHUB_STEP_SUMMARY + echo "|----------|--------|----------|-----------------|--------------------------|---------------------------|" >> $GITHUB_STEP_SUMMARY + + scenarios=$'Regular (dp=2, round_robin)|benchmark_round_robin_regular\nPD (2 prefill + 2 decode, round_robin)|benchmark_round_robin_pd' + + echo "$scenarios" | sed 's/^\s*//' | while IFS='|' read -r label pattern; do + [ -z "$label" ] && continue + # Find the result folder (handle different extraction layouts) + result_folder=$(find . -maxdepth 3 \( -name "$pattern" -o -path "*${pattern}*" \) -type d | head -1) + + if [ -n "$result_folder" ] && [ -d "$result_folder" ]; then + json_file=$(find "$result_folder" -name "*.json" -not -name "experiment_metadata.json" | head -1) + + if [ -n "$json_file" ] && [ -f "$json_file" ]; then + ttft_mean=$(jq -r '.aggregated_metrics.stats.ttft.mean' "$json_file") + e2e_latency_mean=$(jq -r '.aggregated_metrics.stats.e2e_latency.mean' "$json_file") + input_throughput_mean=$(jq -r '.aggregated_metrics.stats.input_throughput.mean' "$json_file") + output_throughput_mean=$(jq -r '.aggregated_metrics.stats.output_throughput.mean' "$json_file") + + ttft_display=$(printf "%.2f" "$ttft_mean" 2>/dev/null || echo "$ttft_mean") + e2e_display=$(printf "%.2f" "$e2e_latency_mean" 2>/dev/null || echo "$e2e_latency_mean") + input_display=$(printf "%.0f" "$input_throughput_mean" 2>/dev/null || echo "$input_throughput_mean") + output_display=$(printf "%.0f" "$output_throughput_mean" 2>/dev/null || echo "$output_throughput_mean") + + echo "| ${label} | ✅ Success | $ttft_display | $e2e_display | $input_display | $output_display |" >> $GITHUB_STEP_SUMMARY + + # Optional GPU utilization table if monitor output exists + gpu_json="$result_folder/gpu_utilization.json" + if [ -f "$gpu_json" ]; then + overall_mean=$(jq -r '.overall.mean // 0' "$gpu_json") + printf "\n#### GPU Utilization — %s\n\n" "$label" >> $GITHUB_STEP_SUMMARY + printf "Overall mean: %.2f%%\n\n" "$overall_mean" >> $GITHUB_STEP_SUMMARY + echo "| GPU | Mean (%) | p5 | p10 | p25 | p50 | p75 | p90 | p95 |" >> $GITHUB_STEP_SUMMARY + echo "|-----|----------|----|-----|-----|-----|-----|-----|-----|" >> $GITHUB_STEP_SUMMARY + jq -r ' + .per_gpu + | to_entries[] + | [ .key, + (.value.mean // 0), + (.value.p5 // 0), + (.value.p10 // 0), + (.value.p25 // 0), + (.value.p50 // 0), + (.value.p75 // 0), + (.value.p90 // 0), + (.value.p95 // 0) + ] + | @tsv' "$gpu_json" \ + | while IFS=$'\t' read -r gpu m p5 p10 p25 p50 p75 p90 p95; do + printf "| %s | %.2f | %.2f | %.2f | %.2f | %.2f | %.2f | %.2f | %.2f |\n" "$gpu" "$m" "$p5" "$p10" "$p25" "$p50" "$p75" "$p90" "$p95" >> $GITHUB_STEP_SUMMARY + done + echo "" >> $GITHUB_STEP_SUMMARY + fi + fi + fi + done diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-xeon.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-xeon.yml new file mode 100644 index 00000000..4eb335a2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-xeon.yml @@ -0,0 +1,102 @@ +name: PR Test (Xeon) + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + workflow_dispatch: + +concurrency: + group: pr-test-xeon-${{ github.ref }} + cancel-in-progress: false + +jobs: + # ==================== Check Changes ==================== # + check-changes: + runs-on: ubuntu-latest + outputs: + main_package: ${{ steps.filter.outputs.main_package }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Detect file changes + id: filter + uses: dorny/paths-filter@v3 + with: + filters: | + main_package: + - "python/sglang/!(multimodal_gen)/**" + - "python/*.toml" + - "scripts/ci/**" + - "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 + + - 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-xpu.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-xpu.yml new file mode 100644 index 00000000..e464bd57 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test-xpu.yml @@ -0,0 +1,118 @@ +name: PR Test (XPU) + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + workflow_dispatch: + +concurrency: + group: pr-test-xpu-${{ github.ref }} + cancel-in-progress: true + +jobs: + # ==================== Check Changes ==================== # + check-changes: + runs-on: ubuntu-latest + outputs: + main_package: ${{ steps.filter.outputs.main_package }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Detect file changes + id: filter + uses: dorny/paths-filter@v3 + with: + filters: | + main_package: + - "python/sglang/!(multimodal_gen)/**" + - "python/*.toml" + - "scripts/ci/**" + - "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 + + - 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/huggingface-cli login --token ${HF_TOKEN} ' + docker exec -u root "$cid" /bin/bash -c "ln -sf /home/sdp/miniforge3/envs/py3.10/bin/python3 /usr/bin/python3" + + - name: Run E2E Bfloat16 tests + timeout-minutes: 20 + run: | + cid="${{ steps.start_container.outputs.container_id }}" + docker exec -w /home/sdp/sglang/ "$cid" \ + bash -c "LD_LIBRARY_PATH=/home/sdp/miniforge3/envs/py3.10/lib:$LD_LIBRARY_PATH && cd ./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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test.yml new file mode 100644 index 00000000..beaa89ac --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/pr-test.yml @@ -0,0 +1,1208 @@ +name: PR Test + +on: + schedule: + - cron: '0 */6 * * *' # Run every 6 hours + pull_request: + branches: [main] + workflow_dispatch: + inputs: + version: + description: "FlashInfer version" + required: true + type: choice + default: "release" + options: + - "release" + - "nightly" + target_stage: + description: "Specific stage to run (optional, for quick testing)" + required: false + type: string + default: "" + +concurrency: + group: pr-test-${{ github.ref }} + cancel-in-progress: true + +jobs: + # =============================================== check changes ==================================================== + check-changes: + runs-on: ubuntu-latest + outputs: + main_package: ${{ steps.filter.outputs.main_package || steps.scheduled.outputs.main_package }} + sgl_kernel: ${{ steps.filter.outputs.sgl_kernel || steps.scheduled.outputs.sgl_kernel }} + multimodal_gen: ${{ steps.filter.outputs.multimodal_gen || steps.scheduled.outputs.multimodal_gen }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Detect file changes + id: filter + uses: dorny/paths-filter@v3 + if: github.event_name != 'schedule' + with: + filters: | + main_package: + - "python/sglang/!(multimodal_gen)/**" + - "python/*.toml" + - "scripts/ci/**" + - "test/**" + - ".github/workflows/pr-test.yml" + sgl_kernel: + - "sgl-kernel/**" + multimodal_gen: + - "python/sglang/multimodal_gen/**" + - "python/sglang/cli/**" + - "python/*.toml" + - ".github/workflows/pr-test.yml" + + - name: Set all changes to true for scheduled runs + id: scheduled + if: github.event_name == 'schedule' + run: | + echo "main_package=true" >> $GITHUB_OUTPUT + echo "sgl_kernel=false" >> $GITHUB_OUTPUT + echo "multimodal_gen=true" >> $GITHUB_OUTPUT + + - name: Show filter results in summary (table) + run: | + { + echo "## Change Detection" + echo "" + echo "| Component | Changed |" + echo "|----------------|---------|" + echo "| main_package | ${{ steps.filter.outputs.main_package || steps.scheduled.outputs.main_package }} |" + echo "| sgl_kernel | ${{ steps.filter.outputs.sgl_kernel || steps.scheduled.outputs.sgl_kernel }} |" + echo "| multimodal_gen | ${{ steps.filter.outputs.multimodal_gen || steps.scheduled.outputs.multimodal_gen }} |" + } >> $GITHUB_STEP_SUMMARY + + # =============================================== PR Gate ==================================================== + call-gate: + needs: check-changes + if: | + needs.check-changes.outputs.main_package == 'true' || + needs.check-changes.outputs.sgl_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] + if: needs.check-changes.outputs.sgl_kernel == 'true' + runs-on: x64-kernel-build-node + 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" + + - 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: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }} + path: sgl-kernel/dist/* + + sgl-kernel-build-wheels-arm: + needs: [check-changes, call-gate] + if: needs.check-changes.outputs.sgl_kernel == 'true' + runs-on: arm-kernel-build-node + 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" + + - 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: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }}-aarch64 + path: sgl-kernel/dist/* + + sgl-kernel-unit-test: + needs: [check-changes, call-gate, sgl-kernel-build-wheels] + if: needs.check-changes.outputs.sgl_kernel == 'true' + runs-on: 1-gpu-runner + env: + RUNNER_LABELS: 1-gpu-runner + steps: + - uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + + - 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] + if: needs.check-changes.outputs.sgl_kernel == 'true' + runs-on: 1-gpu-runner + env: + RUNNER_LABELS: 1-gpu-runner + steps: + - uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 30 + run: | + cd test/srt + python3 test_mla_deepseek_v3.py + + sgl-kernel-benchmark-test: + needs: [check-changes, call-gate, sgl-kernel-build-wheels] + if: needs.check-changes.outputs.sgl_kernel == 'true' + runs-on: 1-gpu-runner + env: + CI: true + RUNNER_LABELS: 1-gpu-runner + steps: + - uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/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] + if: needs.check-changes.outputs.sgl_kernel == 'true' + runs-on: 4-gpu-b200 + env: + RUNNER_LABELS: 4-gpu-b200 + steps: + - uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} IS_BLACKWELL=1 bash scripts/ci/ci_install_dependency.sh + + - 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/ci_install_dependency.sh + + # - name: Run kernel unit tests + # timeout-minutes: 30 + # run: | + # cd sgl-kernel + # pytest tests/ + + # =============================================== primary ==================================================== + + stage-a-test-1: + needs: [check-changes, call-gate, sgl-kernel-build-wheels] + if: | + (inputs.target_stage == 'stage-a-test-1') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 1-gpu-runner + env: + RUNNER_LABELS: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 10 + run: | + cd test/ + python3 run_suite.py --hw cuda --suite stage-a-test-1 + # temporarily put backend-independent cpu tests here + python3 run_suite.py --hw cpu --suite default + + + multimodal-gen-test-1-gpu: + needs: [check-changes, call-gate, sgl-kernel-build-wheels] + if: | + (inputs.target_stage == 'multimodal-gen-test-1-gpu') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + needs.check-changes.outputs.multimodal_gen == 'true' + ) + runs-on: 1-gpu-runner + strategy: + fail-fast: false + matrix: + part: [0, 1] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Cleanup + run: | + ls -alh sgl-kernel/dist || true + rm -rf sgl-kernel/dist/* || true + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh diffusion + - name: Run diffusion server tests + timeout-minutes: 60 + run: | + cd python + python3 sglang/multimodal_gen/test/run_suite.py \ + --suite 1-gpu \ + --partition-id ${{ matrix.part }} \ + --total-partitions 2 + + + multimodal-gen-test-2-gpu: + needs: [check-changes, call-gate, sgl-kernel-build-wheels] + if: | + (inputs.target_stage == 'multimodal-gen-test-2-gpu') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + needs.check-changes.outputs.multimodal_gen == 'true' + ) + runs-on: 2-gpu-runner + strategy: + fail-fast: false + matrix: + part: [0, 1] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Cleanup + run: | + ls -alh sgl-kernel/dist || true + rm -rf sgl-kernel/dist/* || true + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh diffusion + + - name: Run diffusion server tests + timeout-minutes: 60 + run: | + cd python + python3 sglang/multimodal_gen/test/run_suite.py \ + --suite 2-gpu \ + --partition-id ${{ matrix.part }} \ + --total-partitions 2 + + quantization-test: + needs: [check-changes, call-gate, stage-a-test-1] + if: | + (inputs.target_stage == 'quantization-test') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + pip install "bitsandbytes>=0.44.0" + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + - name: Run test + timeout-minutes: 30 + run: | + cd test/srt + python3 run_suite.py --suite quantization_test + + unit-test-backend-1-gpu: + needs: [check-changes, call-gate, stage-a-test-1] + if: | + (inputs.target_stage == 'unit-test-backend-1-gpu') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 1-gpu-runner + env: + RUNNER_LABELS: 1-gpu-runner + strategy: + fail-fast: false + max-parallel: 8 + matrix: + part: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 30 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-1-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 15 + + unit-test-backend-2-gpu: + needs: [check-changes, call-gate, unit-test-backend-1-gpu] + if: | + (inputs.target_stage == 'unit-test-backend-2-gpu') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 2-gpu-runner + env: + RUNNER_LABELS: 2-gpu-runner + strategy: + fail-fast: false + matrix: + part: [0, 1] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 30 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-2-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 + + unit-test-backend-4-gpu: + needs: [check-changes, call-gate, unit-test-backend-2-gpu] + if: | + (inputs.target_stage == 'unit-test-backend-4-gpu') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 4-gpu-h100 + env: + RUNNER_LABELS: 4-gpu-h100 + strategy: + fail-fast: false + matrix: + part: [0, 1, 2] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 20 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-4-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 + + unit-test-backend-8-gpu-h200: + needs: [check-changes, call-gate, unit-test-backend-2-gpu] + if: | + (inputs.target_stage == 'unit-test-backend-8-gpu-h200') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 8-gpu-h200 + env: + RUNNER_LABELS: 8-gpu-h200 + strategy: + fail-fast: false + matrix: + part: [0, 1, 2] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 20 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-8-gpu-h200 --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 + + unit-test-backend-8-gpu-h20: + needs: [check-changes, call-gate, unit-test-backend-2-gpu] + if: | + (inputs.target_stage == 'unit-test-backend-8-gpu-h20') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 8-gpu-h20 + env: + SGLANG_CI_RDMA_ALL_DEVICES: "mlx5_1,mlx5_2,mlx5_3,mlx5_4" + RUNNER_LABELS: 8-gpu-h20 + strategy: + fail-fast: false + matrix: + part: [0, 1] + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_deepep.sh + + - name: Run test + timeout-minutes: 20 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-8-gpu-h20 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 + + performance-test-1-gpu-part-1: + needs: [check-changes, call-gate, stage-a-test-1] + if: | + (inputs.target_stage == 'performance-test-1-gpu-part-1') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 1-gpu-runner + env: + RUNNER_LABELS: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + + - name: Benchmark single latency + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_bs1_small + python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_bs1_default + + - name: Benchmark online latency + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_online_latency_default + + - name: Benchmark offline throughput + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_default + + - name: Benchmark offline throughput (Non-streaming, small batch size) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_non_stream_small_batch_size + + - name: Benchmark online latency (EAGLE) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_online_latency_eagle + + - name: Benchmark online latency (LoRA) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_lora_online_latency + python3 -m unittest test_bench_serving.TestBenchServing.test_lora_online_latency_with_concurrent_adapter_updates + + performance-test-1-gpu-part-2: + needs: [check-changes, call-gate, stage-a-test-1] + if: | + (inputs.target_stage == 'performance-test-1-gpu-part-2') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 1-gpu-runner + env: + RUNNER_LABELS: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + + - name: Benchmark offline throughput (w/o RadixAttention) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_without_radix_cache + + - name: Benchmark offline throughput (w/ Triton) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_with_triton_attention_backend + + - name: Benchmark offline throughput (w/ FP8) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_default_fp8 + + - name: Benchmark VLM offline throughput + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_vlm_offline_throughput + + - name: Benchmark VLM online latency + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_vlm_online_latency + + performance-test-1-gpu-part-3: + needs: [check-changes, call-gate, stage-a-test-1] + if: | + (inputs.target_stage == 'performance-test-1-gpu-part-3') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 1-gpu-runner + env: + RUNNER_LABELS: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + + - name: Benchmark Scores online latency and throughput + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_score_api_latency_throughput + + - name: Benchmark Scores online latency and throughput (batch size scaling) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_score_api_batch_scaling + + - name: Benchmark Embeddings online latency and throughput + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_embeddings_api_latency_throughput + + - name: Benchmark Embeddings online latency and throughput (batch size scaling) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_embeddings_api_batch_scaling + + performance-test-2-gpu: + needs: [check-changes, call-gate, unit-test-backend-2-gpu] + if: | + (inputs.target_stage == 'performance-test-2-gpu') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 2-gpu-runner + env: + RUNNER_LABELS: 2-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + + - name: Benchmark single latency (TP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_moe_tp2_bs1 + + - name: Benchmark single latency + torch.compile (TP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_torch_compile_tp2_bs1 + + - name: Benchmark offline throughput (TP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_moe_offline_throughput_default + + - name: Benchmark offline throughput (w/o RadixAttention) (TP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_moe_offline_throughput_without_radix_cache + + - name: Benchmark offline PP decode throughput (PP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_pp_offline_throughput_default_decode + + - name: Benchmark offline PP prefill throughput (PP=2) + timeout-minutes: 10 + run: | + cd test/srt + python3 -m unittest test_bench_serving.TestBenchServing.test_pp_long_context_prefill + + accuracy-test-1-gpu: + needs: [check-changes, call-gate, stage-a-test-1] + if: | + (inputs.target_stage == 'accuracy-test-1-gpu') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 1-gpu-runner + env: + RUNNER_LABELS: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + git clone https://github.com/merrymercy/human-eval.git + cd human-eval + pip install -e . + + - name: Evaluate accuracy + timeout-minutes: 20 + run: | + cd test/srt + python3 test_eval_accuracy_large.py + + accuracy-test-2-gpu: + needs: [check-changes, call-gate, accuracy-test-1-gpu] + if: | + (inputs.target_stage == 'accuracy-test-2-gpu') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 2-gpu-runner + env: + RUNNER_LABELS: 2-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh + git clone https://github.com/merrymercy/human-eval.git + cd human-eval + pip install -e . + + - name: Evaluate accuracy (TP=2) + timeout-minutes: 20 + run: | + cd test/srt + python3 test_moe_eval_accuracy_large.py + + unit-test-deepep-4-gpu: + needs: [check-changes, call-gate, unit-test-backend-2-gpu] + if: | + (inputs.target_stage == 'unit-test-deepep-4-gpu') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 4-gpu-h100 + env: + RUNNER_LABELS: 4-gpu-h100 + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_deepep.sh + + - name: Run test + timeout-minutes: 20 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-4-gpu-deepep + + unit-test-deepep-8-gpu: + needs: [check-changes, call-gate, unit-test-backend-2-gpu] + if: | + (inputs.target_stage == 'unit-test-deepep-8-gpu') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 8-gpu-h200 + env: + RUNNER_LABELS: 8-gpu-h200 + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_deepep.sh + + - name: Run test + timeout-minutes: 20 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-8-gpu-h200-deepep + + unit-test-backend-4-gpu-b200: + needs: [check-changes, call-gate, unit-test-backend-2-gpu] + if: | + (inputs.target_stage == 'unit-test-backend-4-gpu-b200') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 4-gpu-b200 + env: + RUNNER_LABELS: 4-gpu-b200 + strategy: + fail-fast: false + matrix: + part: [0, 1] + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} IS_BLACKWELL=1 bash scripts/ci/ci_install_dependency.sh + + - name: Run test + timeout-minutes: 30 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-4-gpu-b200 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 --timeout-per-file 1800 + + unit-test-backend-4-gpu-gb200: + needs: [check-changes, call-gate, unit-test-backend-2-gpu, sgl-kernel-build-wheels-arm] + if: | + (inputs.target_stage == 'unit-test-backend-4-gpu-gb200') || + ( + always() && + (github.event_name == 'schedule' || (!failure() && !cancelled())) && + ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) + ) + runs-on: 4-gpu-gb200 + env: + RUNNER_LABELS: 4-gpu-gb200 + strategy: + fail-fast: false + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - 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 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} IS_BLACKWELL=1 GRACE_BLACKWELL=1 bash scripts/ci/ci_install_deepep.sh + + - name: Run test + timeout-minutes: 45 + run: | + cd test/srt + python3 run_suite.py --suite per-commit-4-gpu-gb200 --auto-partition-id 0 --auto-partition-size 1 --timeout-per-file 3600 + + 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, + + multimodal-gen-test-1-gpu, + multimodal-gen-test-2-gpu, + + stage-a-test-1, + quantization-test, + unit-test-backend-1-gpu, + unit-test-backend-2-gpu, + unit-test-backend-4-gpu, + unit-test-backend-8-gpu-h20, + unit-test-backend-8-gpu-h200, + performance-test-1-gpu-part-1, + performance-test-1-gpu-part-2, + performance-test-1-gpu-part-3, + performance-test-2-gpu, + accuracy-test-1-gpu, + accuracy-test-2-gpu, + unit-test-deepep-4-gpu, + unit-test-deepep-8-gpu, + unit-test-backend-4-gpu-b200, + unit-test-backend-4-gpu-gb200, + ] + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-amd-nightly.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-amd-nightly.yml new file mode 100644 index 00000000..47508ac2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-amd-nightly.yml @@ -0,0 +1,132 @@ +name: Release Docker Images Nightly (AMD) +on: + workflow_dispatch: + schedule: + - cron: '0 13 * * *' + +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', 'gfx942-rocm700', 'gfx950'] + build_type: ['all'] + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: "Set Date" + run: | + echo "DATE=$(date +%Y%m%d)" >> $GITHUB_ENV + + - 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 + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + echo "Version: ${version}" + + if [ "${{ matrix.gpu_arch }}" = "gfx942" ]; then + rocm_tag="rocm630-mi30x" + elif [ "${{ matrix.gpu_arch }}" = "gfx942-rocm700" ]; 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} + + docker build . -f docker/rocm.Dockerfile --build-arg BUILD_TYPE=${{ matrix.build_type }} --build-arg GPU_ARCH=${{ matrix.gpu_arch }} -t rocm/sgl-dev:${tag}-${{ env.DATE }} --no-cache + docker push rocm/sgl-dev:${tag}-${{ env.DATE }} + + cache: + 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', 'gfx942-rocm700'] + build_type: ['all'] + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: "Set Date" + run: | + echo "DATE=$(date +%Y%m%d)" >> $GITHUB_ENV + + - 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=$(cat python/sglang/version.py | cut -d'"' -f2) + echo "Version: ${version}" + + if [ "${{ matrix.gpu_arch }}" = "gfx942" ]; then + rocm_tag="rocm630-mi30x" + elif [ "${{ matrix.gpu_arch }}" = "gfx942-rocm700" ]; 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}" == rocm630* ]]; then + final_path="/home/runner/sgl-data/docker/image.tar" + elif [[ "${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}" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-amd.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-amd.yml new file mode 100644 index 00000000..8b4fae51 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-amd.yml @@ -0,0 +1,48 @@ +name: Release Docker Images (AMD) +on: + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +jobs: + publish: + if: github.repository == 'sgl-project/sglang' + runs-on: amd-docker-scale + environment: 'prod' + strategy: + matrix: + gpu_arch: ['gfx942', 'gfx942-rocm700', '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: Build and Push + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + echo "Version: ${version}" + + if [ "${{ matrix.gpu_arch }}" = "gfx942" ]; then + rocm_tag="rocm630-mi30x" + elif [ "${{ matrix.gpu_arch }}" = "gfx942-rocm700" ]; 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} + + docker build . -f docker/rocm.Dockerfile --build-arg BUILD_TYPE=${{ matrix.build_type }} --build-arg GPU_ARCH=${{ matrix.gpu_arch }} -t lmsysorg/sglang:${tag} --no-cache + docker push lmsysorg/sglang:${tag} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-cu13.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-cu13.yml new file mode 100644 index 00000000..32763cb7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-cu13.yml @@ -0,0 +1,119 @@ +name: Build and Push CUDA 13 Docker Images + +# release this manually via workflow_dispatch for now +on: + workflow_dispatch: + +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 + tag: dev-x86-cu13-$(date +%Y%m%d) + version: 13.0.1 + - runner: arm-docker-build-node + platform: linux/arm64 + build_type: all + grace_blackwell: 1 + tag: dev-arm64-cu13-$(date +%Y%m%d) + version: 13.0.1 + 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: true + docker-images: true + android: true + dotnet: true + haskell: true + large-packages: true + swap-storage: true + + - 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: | + docker buildx build \ + --platform ${{ matrix.platform }} \ + --push \ + -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 }} \ + --build-arg USE_LATEST_SGLANG=1 \ + -t lmsysorg/sglang:${{ matrix.tag }} \ + --no-cache \ + . + + create-manifests: + runs-on: ubuntu-22.04 + needs: [build-dev] + if: ${{ github.repository == 'sgl-project/sglang' }} + strategy: + matrix: + variant: + - tag: dev-cu13 + x86_tag: dev-x86-cu13 + arm64_tag: dev-arm64-cu13 + steps: + - uses: docker/setup-buildx-action@v3 + + - uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + - run: | + docker buildx imagetools create \ + -t lmsysorg/sglang:${{ matrix.variant.tag }} \ + -t lmsysorg/sglang:nightly-${{ matrix.variant.tag }}-$(date +%Y%m%d)-${GITHUB_SHA:0:8} \ + lmsysorg/sglang:${{ matrix.variant.x86_tag }} \ + lmsysorg/sglang:${{ matrix.variant.arm64_tag }} + + - name: Cleanup Old Nightly Builds + run: | + # Get JWT token for Docker Hub API + 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) + + # Get all tags for the repository + TAGS_RESPONSE=$(curl -s -H "Authorization: JWT $TOKEN" "https://hub.docker.com/v2/repositories/lmsysorg/sglang/tags/?page_size=100") + + # Extract tags that match our pattern and sort by last_updated timestamp (most recent first) + TAGS=$(echo "$TAGS_RESPONSE" | jq -r '.results[] | select(.name | startswith("nightly-${{ matrix.variant.tag }}-")) | "\(.last_updated)|\(.name)"' | sort -r | cut -d'|' -f2) + + # Count total tags and keep only the 14 most recent + 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) + echo "Tags to delete: $TAGS_TO_DELETE" + + # Delete old tags + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-dev.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-dev.yml new file mode 100644 index 00000000..dfe346b2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-dev.yml @@ -0,0 +1,122 @@ +name: Build and Push Development Docker Images + +on: + workflow_dispatch: + schedule: + - cron: "0 0 * * *" + +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 + tag: dev-x86 + version: 12.9.1 + - runner: arm-docker-build-node + platform: linux/arm64 + build_type: all + grace_blackwell: 1 + tag: dev-arm64 + version: 12.9.1 + 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: true + docker-images: true + android: true + dotnet: true + haskell: true + large-packages: true + swap-storage: true + + - 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: | + docker buildx build \ + --platform ${{ matrix.platform }} \ + --push \ + -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 }} \ + --build-arg USE_LATEST_SGLANG=1 \ + --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \ + -t lmsysorg/sglang:${{ matrix.tag }} \ + --no-cache \ + . + + create-manifests: + runs-on: ubuntu-22.04 + needs: [build-dev] + if: ${{ github.repository == 'sgl-project/sglang' }} + strategy: + matrix: + variant: + - tag: dev + x86_tag: dev-x86 + arm64_tag: dev-arm64 + steps: + - uses: docker/setup-buildx-action@v3 + + - uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + - run: | + SHORT_SHA="${{ github.sha }}" + docker buildx imagetools create \ + -t lmsysorg/sglang:${{ matrix.variant.tag }} \ + -t lmsysorg/sglang:nightly-${{ matrix.variant.tag }}-$(date +%Y%m%d)-${SHORT_SHA:0:8} \ + lmsysorg/sglang:${{ matrix.variant.x86_tag }} \ + lmsysorg/sglang:${{ matrix.variant.arm64_tag }} + + - name: Cleanup Old Nightly Builds + run: | + # Get JWT token for Docker Hub API + 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) + + # Get all tags for the repository + TAGS_RESPONSE=$(curl -s -H "Authorization: JWT $TOKEN" "https://hub.docker.com/v2/repositories/lmsysorg/sglang/tags/?page_size=100") + + # Extract tags that match our pattern and sort by last_updated timestamp (most recent first) + TAGS=$(echo "$TAGS_RESPONSE" | jq -r '.results[] | select(.name | startswith("nightly-${{ matrix.variant.tag }}-")) | "\(.last_updated)|\(.name)"' | sort -r | cut -d'|' -f2) + + # Count total tags and keep only the 14 most recent + 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) + echo "Tags to delete: $TAGS_TO_DELETE" + + # Delete old tags + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-gateway.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-gateway.yml new file mode 100644 index 00000000..d1061333 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-gateway.yml @@ -0,0 +1,30 @@ +name: Release SGLang Model Gateway Docker Image +on: + push: + branches: + - main + paths: + - "sgl-router/bindings/python/sglang_router/version.py" + workflow_dispatch: + +jobs: + publish: + if: github.repository == 'sgl-project/sglang' + runs-on: ubuntu-24.04 + 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: Build and Push + run: | + version=$(cat sgl-router/bindings/python/sglang_router/version.py | cut -d'"' -f2) + tag=v${version} + + docker build . -f docker/gateway.Dockerfile -t lmsysorg/sglang-router:${tag} --no-cache + docker push lmsysorg/sglang-router:${tag} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-npu-nightly.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-npu-nightly.yml new file mode 100644 index 00000000..1873c2c3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-npu-nightly.yml @@ -0,0 +1,78 @@ +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.3.rc1"] + 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 }} + + # 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 + # TODO: need add x86 platforms support when memfabric is ready + platforms: linux/arm64 + 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=20251128 + CANN_VERSION=${{ matrix.cann_version }} + DEVICE_TYPE=${{ matrix.device_type }} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-npu.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-npu.yml new file mode 100644 index 00000000..dd054c9c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-npu.yml @@ -0,0 +1,75 @@ +name: Release Docker Images (NPU) +on: + push: + tags-ignore: + - "gateway-*" # Exclude gateway/router tags + - "router-*" # Exclude router tags + pull_request: + branches: + - main + paths: + - ".github/workflows/release-docker-npu.yml" + - "docker/npu.Dockerfile" + workflow_dispatch: + +jobs: + build: + runs-on: ubuntu-22.04-arm + strategy: + matrix: + cann_version: ["8.3.rc1"] + 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 + type=ref,event=tag,suffix=-cann${{ matrix.cann_version }}-${{ matrix.device_type }} + 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 + id: get_version + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + echo "TAG=lmsysorg/sglang:v$version-cann${{ matrix.cann_version }}-${{ matrix.device_type }}" >> $GITHUB_OUTPUT + + - name: Build and push Docker image + id: build-and-push + uses: docker/build-push-action@v6 + with: + context: docker + file: docker/npu.Dockerfile + # TODO: need add x86 platforms support when memfabric is ready + platforms: linux/arm64 + labels: ${{ steps.meta.outputs.labels }} + tags: ${{ steps.meta.outputs.tags || steps.get_version.outputs.TAG }} + push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} + provenance: false + build-args: | + SGLANG_KERNEL_NPU_TAG=20251128 + CANN_VERSION=${{ matrix.cann_version }} + DEVICE_TYPE=${{ matrix.device_type }} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-xeon.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-xeon.yml new file mode 100644 index 00000000..60e24933 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker-xeon.yml @@ -0,0 +1,35 @@ +name: Release Docker Xeon Images +on: + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +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: Build and Push + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + tag=v${version}-xeon + + docker build . -f docker/xeon.Dockerfile -t lmsysorg/sglang:${tag} --no-cache + docker push lmsysorg/sglang:${tag} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker.yml new file mode 100644 index 00000000..59603385 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docker.yml @@ -0,0 +1,138 @@ +name: Release Docker Images +on: + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +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: Build and Push AMD64 + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + tag=v${version}-cu129-amd64 + + docker buildx build \ + --platform linux/amd64 \ + --push \ + -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 }} \ + -t lmsysorg/sglang:${tag} \ + --no-cache \ + . + + 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: Build and Push ARM64 + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + tag=v${version}-cu129-arm64 + + docker buildx build \ + --platform linux/arm64 \ + --push \ + -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 }} \ + -t lmsysorg/sglang:${tag} \ + --no-cache \ + . + + 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: Create multi-arch manifests + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + + # Create versioned manifest + docker buildx imagetools create \ + -t lmsysorg/sglang:v${version} \ + lmsysorg/sglang:v${version}-cu129-amd64 \ + lmsysorg/sglang:v${version}-cu129-arm64 + + # Create latest manifest + docker buildx imagetools create \ + -t lmsysorg/sglang:latest \ + lmsysorg/sglang:v${version}-cu129-amd64 \ + lmsysorg/sglang:v${version}-cu129-arm64 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docs.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docs.yml new file mode 100644 index 00000000..78fafc60 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-docs.yml @@ -0,0 +1,65 @@ +name: Release Documentation + +on: + push: + branches: + - main + paths: + - "docs/**" + - "python/sglang/version.py" + - "python/sglang/**" + workflow_dispatch: + +concurrency: + group: release-docs-${{ github.ref }} + cancel-in-progress: 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: Install dependencies + run: | + bash scripts/ci/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 + python3 wrap_run_llm.py + + cd _build/html + + git clone https://$GITHUB_TOKEN@github.com/sgl-project/sgl-project.github.io.git ../sgl-project.github.io --depth 1 + find ../sgl-project.github.io/ -mindepth 1 -not -path "../sgl-project.github.io/.git*" -not -name CNAME -not -name ".jekyll" -not -name ".nojekyll" -delete + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-fake-tag.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-fake-tag.yml new file mode 100644 index 00000000..d1acc6bf --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-fake-tag.yml @@ -0,0 +1,35 @@ +name: Release Fake Tag +on: + push: + branches: + - main + paths: + - "python/sglang/version.py" + workflow_dispatch: + +permissions: + contents: write + +jobs: + publish: + if: github.repository == 'sgl-project/sglang' + runs-on: ubuntu-latest + environment: 'prod' + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + token: ${{ secrets.GITHUB_TOKEN }} + + - name: Get version + id: get_version + run: | + version=$(cat python/sglang/version.py | cut -d'"' -f2) + echo "TAG=v$version" >> $GITHUB_OUTPUT + + - name: Create and push tag + run: | + git config user.name "sglang-bot" + git config user.email "sglang-bot@users.noreply.github.com" + git tag ${{ steps.get_version.outputs.TAG }} + git push origin ${{ steps.get_version.outputs.TAG }} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-pypi-gateway.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-pypi-gateway.yml new file mode 100644 index 00000000..0f051faa --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-pypi-gateway.yml @@ -0,0 +1,167 @@ +name: Release SGLang Model Gateway to PyPI + +on: + push: + branches: + - main + paths: + - sgl-router/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-router folder to root and delete sglang-repo + run: | + mv sglang-repo/sgl-router/* . + 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-router folder to root and delete sglang-repo + run: | + mv sglang-repo/sgl-router/* . + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-pypi.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-pypi.yml new file mode 100644 index 00000000..354ee21d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-pypi.yml @@ -0,0 +1,31 @@ +name: Release PyPI +on: + push: + branches: + - main + paths: + - "python/sglang/version.py" + 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 + + - name: Upload to pypi + run: | + cd python + cp ../README.md ../LICENSE . + pip install build + python3 -m build + pip install twine + python3 -m twine upload dist/* -u __token__ -p ${{ secrets.PYPI_TOKEN }} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-whl-kernel.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-whl-kernel.yml new file mode 100644 index 00000000..4f0a33ef --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/release-whl-kernel.yml @@ -0,0 +1,197 @@ +name: Release SGLang Kernels + +on: + push: + branches: + - main + paths: + - sgl-kernel/python/sgl_kernel/version.py + workflow_dispatch: + inputs: + tag_name: + 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' + 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" + + - 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: + USE_CCACHE: 0 + + - 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 }} + + - 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 + + - 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.WHL_TOKEN }} + 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.WHL_TOKEN }} + + - 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' + 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" + + - 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: + USE_CCACHE: 0 + + - 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 + + - 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/slash-command-handler.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/slash-command-handler.yml new file mode 100644 index 00000000..9ba79953 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/slash-command-handler.yml @@ -0,0 +1,46 @@ +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 starts with a recognized command + if: > + github.event.issue.pull_request && + (startsWith(github.event.comment.body, '/tag-run-ci-label') || + startsWith(github.event.comment.body, '/rerun-failed-ci') || + startsWith(github.event.comment.body, '/tag-and-rerun-ci') || + startsWith(github.event.comment.body, '/rerun-stage')) + 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: 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/slash_command_handler.py diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/stress-test.yml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/stress-test.yml new file mode 100644 index 00000000..4e6bc741 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.github/workflows/stress-test.yml @@ -0,0 +1,44 @@ +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/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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.gitignore b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.gitignore new file mode 100644 index 00000000..118dd9ae --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.gitignore @@ -0,0 +1,246 @@ +# 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 + +compile_commands.json + +*.iml + +# VSCode +.vscode + +1 + +# Autoenv +.env.leave + +# Rust lib +Cargo.lock + +lmms-eval + +**/.claude/ +**/.serena/ diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.isort.cfg b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.isort.cfg new file mode 100644 index 00000000..83509515 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.isort.cfg @@ -0,0 +1,3 @@ +[settings] +profile=black +known_first_party=sglang diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.pre-commit-config.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.pre-commit-config.yaml new file mode 100644 index 00000000..5ddd480f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/.pre-commit-config.yaml @@ -0,0 +1,77 @@ +default_stages: [pre-commit, pre-push, manual] +exclude: ^python/sglang/multimodal_gen/csrc + +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.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 + - id: debug-statements + - id: no-commit-to-branch + - repo: https://github.com/PyCQA/isort + rev: 5.13.2 + 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.11.7 + hooks: + - id: ruff + args: + - --select=F401,F821 + - --fix + files: ^(benchmark/|docs/|examples/|python/sglang/|sgl-router/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: 24.10.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 + additional_dependencies: ['tomli'] + args: ['--toml', 'python/pyproject.toml', '-L', 'cann,thi,makro,wil,rouge,PRIS'] + exclude: | + (?x)^( + test/srt/test_reasoning_parser\.py| + docs/advanced_features/vlm_query\.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| + sgl-router/src/tokenizer/chat_template\.rs + )$ + - repo: https://github.com/pre-commit/mirrors-clang-format + rev: v18.1.8 + hooks: + - id: clang-format + types_or: [c++, cuda] + args: [--style=file, --verbose] + - repo: https://github.com/kynan/nbstripout + rev: 0.8.1 + hooks: + - id: nbstripout + args: + - '--keep-output' + - '--extra-keys=metadata.kernelspec metadata.language_info.version' diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/PROFILING.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/PROFILING.md new file mode 100644 index 00000000..7e15ec84 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/PROFILING.md @@ -0,0 +1,425 @@ +## 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/client.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/client.sh new file mode 100755 index 00000000..150ea9f1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/client.sh @@ -0,0 +1,27 @@ +#!/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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/install_rpd.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/install_rpd.sh new file mode 100644 index 00000000..d1b04b98 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/install_rpd.sh @@ -0,0 +1,10 @@ +# 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 .. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/loadTracer.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/loadTracer.sh new file mode 100755 index 00000000..8a95a335 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/loadTracer.sh @@ -0,0 +1,43 @@ +#!/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 "$@" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/rpd.patch b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/rpd.patch new file mode 100644 index 00000000..87917654 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/rpd.patch @@ -0,0 +1,12 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/rpd_profile_server_enable.patch b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/rpd_profile_server_enable.patch new file mode 100644 index 00000000..3cd39153 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/rpd_profile_server_enable.patch @@ -0,0 +1,49 @@ +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") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/rpd_profile_server_enable_wCPU_activities.patch b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/rpd_profile_server_enable_wCPU_activities.patch new file mode 100644 index 00000000..5416f4d5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/rpd_profile_server_enable_wCPU_activities.patch @@ -0,0 +1,126 @@ +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, diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/server.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/server.sh new file mode 100755 index 00000000..f877e6c7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/server.sh @@ -0,0 +1,20 @@ +#!/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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/torch_profiler.patch b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/torch_profiler.patch new file mode 100644 index 00000000..40f55740 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/profiling/torch_profiler.patch @@ -0,0 +1,25 @@ +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") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/tuning/TUNING.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/tuning/TUNING.md new file mode 100644 index 00000000..e7b9b204 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/tuning/TUNING.md @@ -0,0 +1,118 @@ +## 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_funtion(): + ... +``` +## 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/tuning/benchmark_moe_rocm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/tuning/benchmark_moe_rocm.py new file mode 100644 index 00000000..af596d21 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/3rdparty/amd/tuning/benchmark_moe_rocm.py @@ -0,0 +1,380 @@ +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)) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/CODE_OF_CONDUCT.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/CODE_OF_CONDUCT.md new file mode 100644 index 00000000..18c91471 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/CODE_OF_CONDUCT.md @@ -0,0 +1,128 @@ +# Contributor Covenant Code of Conduct + +## Our Pledge + +We as members, contributors, and leaders pledge to make participation in our +community a harassment-free experience for everyone, regardless of age, body +size, visible or invisible disability, ethnicity, sex characteristics, gender +identity and expression, level of experience, education, socio-economic status, +nationality, personal appearance, race, religion, or sexual identity +and orientation. + +We pledge to act and interact in ways that contribute to an open, welcoming, +diverse, inclusive, and healthy community. + +## Our Standards + +Examples of behavior that contributes to a positive environment for our +community include: + +* Demonstrating empathy and kindness toward other people +* Being respectful of differing opinions, viewpoints, and experiences +* Giving and gracefully accepting constructive feedback +* Accepting responsibility and apologizing to those affected by our mistakes, + and learning from the experience +* Focusing on what is best not just for us as individuals, but for the + overall community + +Examples of unacceptable behavior include: + +* The use of sexualized language or imagery, and sexual attention or + advances of any kind +* Trolling, insulting or derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information, such as a physical or email + address, without their explicit permission +* Other conduct which could reasonably be considered inappropriate in a + professional setting + +## Enforcement Responsibilities + +Community leaders are responsible for clarifying and enforcing our standards of +acceptable behavior and will take appropriate and fair corrective action in +response to any behavior that they deem inappropriate, threatening, offensive, +or harmful. + +Community leaders have the right and responsibility to remove, edit, or reject +comments, commits, code, wiki edits, issues, and other contributions that are +not aligned to this Code of Conduct, and will communicate reasons for moderation +decisions when appropriate. + +## Scope + +This Code of Conduct applies within all community spaces, and also applies when +an individual is officially representing the community in public spaces. +Examples of representing our community include using an official e-mail address, +posting via an official social media account, or acting as an appointed +representative at an online or offline event. + +## Enforcement + +Instances of abusive, harassing, or otherwise unacceptable behavior may be +reported to the community leaders responsible for enforcement at +. +All complaints will be reviewed and investigated promptly and fairly. + +All community leaders are obligated to respect the privacy and security of the +reporter of any incident. + +## Enforcement Guidelines + +Community leaders will follow these Community Impact Guidelines in determining +the consequences for any action they deem in violation of this Code of Conduct: + +### 1. Correction + +**Community Impact**: Use of inappropriate language or other behavior deemed +unprofessional or unwelcome in the community. + +**Consequence**: A private, written warning from community leaders, providing +clarity around the nature of the violation and an explanation of why the +behavior was inappropriate. A public apology may be requested. + +### 2. Warning + +**Community Impact**: A violation through a single incident or series +of actions. + +**Consequence**: A warning with consequences for continued behavior. No +interaction with the people involved, including unsolicited interaction with +those enforcing the Code of Conduct, for a specified period of time. This +includes avoiding interactions in community spaces as well as external channels +like social media. Violating these terms may lead to a temporary or +permanent ban. + +### 3. Temporary Ban + +**Community Impact**: A serious violation of community standards, including +sustained inappropriate behavior. + +**Consequence**: A temporary ban from any sort of interaction or public +communication with the community for a specified period of time. No public or +private interaction with the people involved, including unsolicited interaction +with those enforcing the Code of Conduct, is allowed during this period. +Violating these terms may lead to a permanent ban. + +### 4. Permanent Ban + +**Community Impact**: Demonstrating a pattern of violation of community +standards, including sustained inappropriate behavior, harassment of an +individual, or aggression toward or disparagement of classes of individuals. + +**Consequence**: A permanent ban from any sort of public interaction within +the community. + +## Attribution + +This Code of Conduct is adapted from the [Contributor Covenant][homepage], +version 2.0, available at +https://www.contributor-covenant.org/version/2/0/code_of_conduct.html. + +Community Impact Guidelines were inspired by [Mozilla's code of conduct +enforcement ladder](https://github.com/mozilla/diversity). + +[homepage]: https://www.contributor-covenant.org + +For answers to common questions about this code of conduct, see the FAQ at +https://www.contributor-covenant.org/faq. Translations are available at +https://www.contributor-covenant.org/translations. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/LICENSE b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/LICENSE new file mode 100644 index 00000000..9c422689 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/LICENSE @@ -0,0 +1,201 @@ + 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/Makefile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/Makefile new file mode 100644 index 00000000..d6ef1942 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/Makefile @@ -0,0 +1,49 @@ +.PHONY: check-deps install-deps format update help + +# Show help for each target +help: + @echo "Available targets:" + @grep -E '^[a-zA-Z0-9_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}' + +check-deps: ## Check and install required Python formatting dependencies + @command -v isort >/dev/null 2>&1 || (echo "Installing isort..." && pip install isort) + @command -v black >/dev/null 2>&1 || (echo "Installing black..." && pip install black) + +install-deps: ## Install Python formatting tools (isort and black) + pip install isort black + +format: check-deps ## Format modified Python files using isort and black + @echo "Formatting modified Python files..." + git diff --name-only --diff-filter=M | grep '\.py$$' | xargs -I {} sh -c 'isort {} && black {}' + +FILES_TO_UPDATE = docker/rocm.Dockerfile \ + python/pyproject.toml \ + python/pyproject_other.toml \ + python/sglang/version.py \ + docs/developer_guide/setup_github_runner.md \ + docs/get_started/install.md \ + docs/platforms/amd_gpu.md \ + docs/platforms/ascend_npu.md \ + docs/platforms/cpu_server.md \ + docs/platforms/xpu.md \ + benchmark/deepseek_v3/README.md + +update: ## Update version numbers across project files. Usage: make update + @if [ -z "$(filter-out $@,$(MAKECMDGOALS))" ]; then \ + echo "Version required. Usage: make update "; \ + exit 1; \ + fi + @OLD_VERSION=$$(grep "version" python/sglang/version.py | cut -d '"' -f2); \ + NEW_VERSION=$(filter-out $@,$(MAKECMDGOALS)); \ + echo "Updating version from $$OLD_VERSION to $$NEW_VERSION"; \ + for file in $(FILES_TO_UPDATE); do \ + if [ "$(shell uname)" = "Darwin" ]; then \ + sed -i '' -e "s/$$OLD_VERSION/$$NEW_VERSION/g" $$file; \ + else \ + sed -i -e "s/$$OLD_VERSION/$$NEW_VERSION/g" $$file; \ + fi \ + done; \ + echo "Version update complete" + +%: + @: diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/README.md new file mode 100644 index 00000000..3ae19493 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/README.md @@ -0,0 +1,84 @@ +
+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**](https://lmsys.org/blog/) +| [**Documentation**](https://docs.sglang.io/) +| [**Roadmap**](https://roadmap.sglang.io/) +| [**Join Slack**](https://slack.sglang.io/) +| [**Weekly Dev Meeting**](https://meet.sglang.io/) +| [**Slides**](https://github.com/sgl-project/sgl-learning-materials?tab=readme-ov-file#slides) | + +## News +- [2025/11] 🔥 SGLang Diffusion accelerates video and image generation ([blog](https://lmsys.org/blog/2025-11-07-sglang-diffusion/)). +- [2025/10] 🔥 SGLang now runs natively on TPU with the SGLang-Jax backend ([blog](https://lmsys.org/blog/2025-10-29-sglang-jax/)). +- [2025/10] PyTorch Conference 2025 SGLang Talk ([slide](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/sglang_pytorch_2025.pdf)). +- [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)). +- [2025/08] SGLang provides day-0 support for OpenAI gpt-oss model ([instructions](https://github.com/sgl-project/sglang/issues/8833)) +- [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/)). + +
+More + +- [2025/10] SGLang x Nvidia SF Meetup on 10/2 ([recap](https://x.com/lmsysorg/status/1975339501934510231)). +- [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/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 vision-language 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 Backend 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. +- **Extensive Model Support**: Supports a wide range of generative 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 integrating 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. +- **Flexible Frontend Language**: Offers an intuitive interface for programming LLM applications, supporting chained generation calls, advanced prompting, control flow, multi-modal inputs, parallelism, and external interactions. +- **Active Community**: SGLang is open-source and supported by a vibrant community with widespread industry adoption, powering over 400,000 GPUs worldwide. + +## 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). diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/assets/logo.png b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/assets/logo.png new file mode 100644 index 00000000..2a8bc258 Binary files /dev/null and b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/assets/logo.png differ diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/assets/logo.svg b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/assets/logo.svg new file mode 100644 index 00000000..4d639392 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/assets/logo.svg @@ -0,0 +1 @@ + diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/assets/logo_square.svg b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/assets/logo_square.svg new file mode 100644 index 00000000..a82fa0ae --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/assets/logo_square.svg @@ -0,0 +1 @@ + diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/bench_attention_sink/bench_attention_sink_triton.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/bench_attention_sink/bench_attention_sink_triton.py new file mode 100644 index 00000000..21bc2a59 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/bench_attention_sink/bench_attention_sink_triton.py @@ -0,0 +1,250 @@ +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!") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/bench_in_batch_prefix/bench_in_batch_prefix.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/bench_in_batch_prefix/bench_in_batch_prefix.py new file mode 100644 index 00000000..28209711 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/bench_in_batch_prefix/bench_in_batch_prefix.py @@ -0,0 +1,130 @@ +# 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") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/benchmark_batch/benchmark_batch.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/benchmark_batch/benchmark_batch.py new file mode 100644 index 00000000..a8592d48 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/benchmark_batch/benchmark_batch.py @@ -0,0 +1,193 @@ +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/benchmark_batch/benchmark_tokenizer.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/benchmark_batch/benchmark_tokenizer.py new file mode 100644 index 00000000..88a5820b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/benchmark_batch/benchmark_tokenizer.py @@ -0,0 +1,126 @@ +import random +import time +from statistics import mean + +from transformers import AutoTokenizer + +# CONFIG +TOKENIZER_DIR = ( + "/shared/public/sharing/fait360brew/training/models/meta-llama/Llama-3.2-3B" +) +NUM_TOKENS = 20000 # Each prompt should contain this many tokens +BATCH_SIZES = [1, 2, 4, 8] # Test different batch sizes +NUM_RUNS = 5 # Number of runs for each batch size to get reliable measurements + + +def generate_random_prompts(num_prompts, num_tokens, tokenizer): + """Generate random prompts with specified token count.""" + vocab_size = tokenizer.vocab_size + all_prompts = [] + + print(f"Generating {num_prompts} random prompts with {num_tokens} tokens each...") + for i in range(num_prompts): + # Generate random token IDs - this directly gives us the exact token count + random_token_ids = [ + random.randint(0, vocab_size - 1) for _ in range(num_tokens) + ] + random_text = tokenizer.decode( + random_token_ids, clean_up_tokenization_spaces=True + ) + + prompt = f"Prompt {i}: {random_text}" + tokens = tokenizer.encode(prompt) + print(f" Prompt {i}: {len(tokens)} tokens") + all_prompts.append(prompt) + + return all_prompts + + +def benchmark_sequential_vs_batch(prompts, batch_size, tokenizer): + """Compare sequential vs batch tokenization for a given batch size.""" + + # Sequential tokenization using encode() + sequential_times = [] + for run in range(NUM_RUNS): + batch_prompts = prompts[:batch_size] # Use same prompts for fair comparison + + start_time = time.perf_counter() + for prompt in batch_prompts: + tokens = tokenizer.encode(prompt) + sequential_time = (time.perf_counter() - start_time) * 1000 + sequential_times.append(sequential_time) + + # Batch tokenization using tokenizer() + batch_times = [] + for run in range(NUM_RUNS): + batch_prompts = prompts[:batch_size] # Use same prompts for fair comparison + + start_time = time.perf_counter() + tokens = tokenizer(batch_prompts) + batch_time = (time.perf_counter() - start_time) * 1000 + batch_times.append(batch_time) + + return { + "batch_size": batch_size, + "avg_sequential_ms": mean(sequential_times), + "avg_batch_ms": mean(batch_times), + "speedup_factor": ( + mean(sequential_times) / mean(batch_times) if mean(batch_times) > 0 else 0 + ), + "sequential_runs": sequential_times, + "batch_runs": batch_times, + } + + +def main(): + print("Tokenizer Benchmark: Sequential vs Batch Processing") + print("-" * 60) + print(f"Tokenizer: {TOKENIZER_DIR}") + print(f"Tokens per prompt: {NUM_TOKENS}") + print(f"Number of runs per batch size: {NUM_RUNS}") + print("-" * 60) + + # Load tokenizer once for all operations + tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR) + + # The largest batch size determines how many prompts we need + max_batch_size = max(BATCH_SIZES) + all_prompts = generate_random_prompts(max_batch_size, NUM_TOKENS, tokenizer) + + results = [] + print("\nRunning benchmark...") + + for batch_size in BATCH_SIZES: + print(f"\nBenchmarking batch size: {batch_size}") + result = benchmark_sequential_vs_batch(all_prompts, batch_size, tokenizer) + results.append(result) + + print(f" Sequential tokenization (encode):") + for i, run_time in enumerate(result["sequential_runs"]): + print(f" Run {i+1}: {run_time:.2f} ms") + print(f" Average: {result['avg_sequential_ms']:.2f} ms") + + print(f" Batch tokenization (tokenizer):") + for i, run_time in enumerate(result["batch_runs"]): + print(f" Run {i+1}: {run_time:.2f} ms") + print(f" Average: {result['avg_batch_ms']:.2f} ms") + + print(f" Speedup factor: {result['speedup_factor']:.2f}x") + + print("\n" + "=" * 60) + print("SUMMARY OF RESULTS") + print("=" * 60) + print( + f"{'Batch Size':<10} {'Sequential (ms)':<18} {'Batch (ms)':<18} {'Speedup':<10}" + ) + print("-" * 60) + + for result in results: + print( + f"{result['batch_size']:<10} {result['avg_sequential_ms']:.2f} ms{' ' * 8} {result['avg_batch_ms']:.2f} ms{' ' * 8} {result['speedup_factor']:.2f}x" + ) + + +if __name__ == "__main__": + random.seed(0) + main() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/benchmark_vllm_060/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/benchmark_vllm_060/README.md new file mode 100644 index 00000000..b480dabf --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/benchmark_vllm_060/README.md @@ -0,0 +1,89 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/405b_sglang.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/405b_sglang.sh new file mode 100644 index 00000000..49185378 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/405b_sglang.sh @@ -0,0 +1,24 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/405b_trt.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/405b_trt.sh new file mode 100644 index 00000000..1950bc92 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/405b_trt.sh @@ -0,0 +1,17 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/405b_vllm.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/405b_vllm.sh new file mode 100644 index 00000000..d1ef048e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/405b_vllm.sh @@ -0,0 +1,24 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/README.md new file mode 100644 index 00000000..7448554e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/README.md @@ -0,0 +1,164 @@ +# 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/config.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/config.md new file mode 100644 index 00000000..3faf6009 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/blog_v0_2/config.md @@ -0,0 +1,100 @@ +### 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 +} +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/README.md new file mode 100644 index 00000000..3704742e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/README.md @@ -0,0 +1,19 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/bench_sglang.py new file mode 100644 index 00000000..b3ce3c99 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/bench_sglang.py @@ -0,0 +1,124 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/convert_parquet_to_json.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/convert_parquet_to_json.py new file mode 100644 index 00000000..e3e69cb3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/convert_parquet_to_json.py @@ -0,0 +1,28 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/parquet_to_json.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/parquet_to_json.sh new file mode 100755 index 00000000..9aaf087f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/boolq/parquet_to_json.sh @@ -0,0 +1,26 @@ +#!/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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/ceval/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/ceval/README.md new file mode 100644 index 00000000..b822e43c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/ceval/README.md @@ -0,0 +1,15 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/ceval/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/ceval/bench_sglang.py new file mode 100644 index 00000000..bcebd55c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/ceval/bench_sglang.py @@ -0,0 +1,138 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/deepseek_v3/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/deepseek_v3/README.md new file mode 100644 index 00000000..59459180 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/deepseek_v3/README.md @@ -0,0 +1,388 @@ +# 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 Model Optimizations in SGLang](https://docs.sglang.io/basic_usage/deepseek.html). + +## 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[all]>=0.5.6" + +# 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 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 http://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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/dspy/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/dspy/README.md new file mode 100644 index 00000000..a2b213aa --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/dspy/README.md @@ -0,0 +1,51 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/dspy/bench_dspy_intro.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/dspy/bench_dspy_intro.py new file mode 100644 index 00000000..2b0936ed --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/dspy/bench_dspy_intro.py @@ -0,0 +1,192 @@ +""" +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/README.md new file mode 100644 index 00000000..393a9ce8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/README.md @@ -0,0 +1,38 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/agent_functions.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/agent_functions.py new file mode 100644 index 00000000..9785fc1e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/agent_functions.py @@ -0,0 +1,300 @@ +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": "}"} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/bench_other.py new file mode 100644 index 00000000..c0b3a340 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/bench_other.py @@ -0,0 +1,80 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/bench_sglang.py new file mode 100644 index 00000000..034b1659 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/generative_agents/bench_sglang.py @@ -0,0 +1,74 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gpt_oss/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gpt_oss/README.md new file mode 100644 index 00000000..4d1b00e9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gpt_oss/README.md @@ -0,0 +1,163 @@ +# 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gsm8k/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gsm8k/README.md new file mode 100644 index 00000000..c110f533 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gsm8k/README.md @@ -0,0 +1,47 @@ +## 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 +``` +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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gsm8k/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gsm8k/bench_other.py new file mode 100644 index 00000000..6dcb9ad7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gsm8k/bench_other.py @@ -0,0 +1,151 @@ +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 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 + 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", + "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) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gsm8k/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gsm8k/bench_sglang.py new file mode 100644 index 00000000..9cdc2cf8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/gsm8k/bench_sglang.py @@ -0,0 +1,148 @@ +import argparse +import ast +import json +import os +import re +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, + 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)) + + # Read data + 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])): + questions.append(get_one_example(lines, i, False)) + 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 += few_shot_examples + question + s += sgl.gen( + "answer", max_tokens=512, stop=["Question", "Assistant:", "<|separator|>"] + ) + + ##################################### + ########## SGL Program End ########## + ##################################### + + # Run requests + tic = time.perf_counter() + states = few_shot_gsm8k.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(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", + "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) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hellaswag/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hellaswag/README.md new file mode 100644 index 00000000..cb7e6536 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hellaswag/README.md @@ -0,0 +1,47 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hellaswag/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hellaswag/bench_other.py new file mode 100644 index 00000000..cde0794b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hellaswag/bench_other.py @@ -0,0 +1,118 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hellaswag/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hellaswag/bench_sglang.py new file mode 100644 index 00000000..2adce99b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hellaswag/bench_sglang.py @@ -0,0 +1,109 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench.sh new file mode 100644 index 00000000..049116b8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench.sh @@ -0,0 +1,59 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench_client.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench_client.py new file mode 100644 index 00000000..0af3c80c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench_client.py @@ -0,0 +1,162 @@ +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) + + 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) + 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench_storage.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench_storage.py new file mode 100644 index 00000000..f0ce171b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench_storage.py @@ -0,0 +1,258 @@ +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench_zerocopy.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench_zerocopy.py new file mode 100644 index 00000000..bfa7bff0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hf3fs/bench_zerocopy.py @@ -0,0 +1,140 @@ +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") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/README.md new file mode 100644 index 00000000..f2aa6ff0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/README.md @@ -0,0 +1,91 @@ +## 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_long_context.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_long_context.py new file mode 100644 index 00000000..a3656cef --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_long_context.py @@ -0,0 +1,102 @@ +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.bench_serving 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"] + init_requests.append( + ( + i, + gen_payload( + self.dataset["contexts"][context_id] + + self.dataset["queries"][i]["question"], + len( + self.tokenizer( + self.dataset["queries"][i]["reference_answer"] + )["input_ids"] + ), + ), + ) + ) + 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_mix.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_mix.py new file mode 100644 index 00000000..cfd25bc4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_mix.py @@ -0,0 +1,567 @@ +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, + get_tokenizer, + remove_prefix, + sample_random_requests, +) + +# 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_mix.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_mix.sh new file mode 100755 index 00000000..5ff6dca9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_mix.sh @@ -0,0 +1,42 @@ +#!/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 & diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_multiturn.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_multiturn.py new file mode 100644 index 00000000..fe154d6b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_multiturn.py @@ -0,0 +1,615 @@ +import argparse +import asyncio +import json +import queue +import random +import threading +import time +from datetime import datetime +from typing import Optional + +import aiohttp +import numpy as np +import requests +from tqdm.asyncio import tqdm + +from sglang.bench_serving import ( + RequestFuncOutput, + get_tokenizer, + remove_prefix, + sample_random_requests, +) + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60) + + +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("--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.", + ) + return parser.parse_args() + + +async def async_request_sglang_generate( + payload, + url, + pbar: Optional[tqdm] = None, +): + """ + 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() + + 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["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 + output.generated_len = len(output.itl) + 1 + 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}") + + if pbar: + pbar.update(1) + return output + + +def gen_payload(prompt, output_len, lora_path=""): + 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": lora_path, + "return_logprob": False, + "logprob_start_len": -1, + } + return payload + + +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): + # Construct the base URL for requests + self.url = f"http://{args.host}:{args.port}/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.candidate_inputs = sample_random_requests( + input_len=args.request_length, + output_len=args.output_length, + num_prompts=args.num_clients, + range_ratio=1.0, + tokenizer=self.tokenizer, + dataset_path=args.dataset_path, + random_sample=not args.disable_random_sample, + ) + self.candidate_inputs = [i.prompt for i in self.candidate_inputs] + + if args.sub_question_input_length != 0: + sub_question_input_length = args.sub_question_input_length + else: + sub_question_input_length = args.request_length + + self.sub_question_inputs = sample_random_requests( + input_len=sub_question_input_length, + output_len=args.output_length, + num_prompts=args.num_clients * max(args.num_rounds - 1, 1), + range_ratio=1.0, + tokenizer=self.tokenizer, + dataset_path=args.dataset_path, + random_sample=not args.disable_random_sample, + ) + + init_requests = [ + ( + i, + gen_payload( + self.candidate_inputs[i], args.output_length, args.lora_path + ), + ) + for i in range(args.num_clients) + ] + self.client_records = { + i: {"round": 0, "history": init_requests[i][1]["text"]} + 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=args.num_clients * args.num_rounds) + self.performance_metrics = { + "ttft": [], + "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(args.num_rounds): + self.performance_metrics[f"round_{i}"] = { + "ttft": [], + "latency": [], + "prompt_len": [], + "cached_tokens": [], + "generated_len": [], + } + self.num_clients = args.num_clients + + self.num_rounds = args.num_rounds + self.max_parallel = args.max_parallel + self.output_length = args.output_length + + async def handle_request(self, item): + try: + client_id, payload = item + response = await async_request_sglang_generate(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: {e}") + + 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 = [] + 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.client_records[client_id]["history"] += response.generated_text + 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["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 + + if self.client_records[client_id]["round"] < self.num_rounds: + # append new request to client's history + self.client_records[client_id][ + "history" + ] += self.sub_question_inputs.pop().prompt + new_req = ( + client_id, + gen_payload( + self.client_records[client_id]["history"], + self.output_length, + args.lora_path, + ), + ) + if self.enable_round_barrier: + next_round_reqs.append(new_req) + if len(next_round_reqs) == self.num_clients: + for req in next_round_reqs: + self.ready_queue.append(req) + next_round_reqs = [] + else: + self.ready_queue.append(new_req) + 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 + 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 + ), + "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 + ], + "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(args.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" 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" 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"] + print( + f" Round {round_num}: Average TTFT = {avg_ttft:.2f}s, " + f"Cache Hit Rate = {cache_hit_rate:.6f} " + f"({request_count} requests)" + ) + 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_serving.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_serving.py new file mode 100644 index 00000000..e38d0d0e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/bench_serving.py @@ -0,0 +1,1029 @@ +# 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.bench_serving 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/data_processing.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/data_processing.py new file mode 100644 index 00000000..1fb3650c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/data_processing.py @@ -0,0 +1,569 @@ +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 + +SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json" + +from sglang.bench_serving import ( + download_and_cache_file, + gen_prompt, + get_gen_prefix_cache_path, +) +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_file(SHAREGPT_URL) + + # 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_file(SHAREGPT_URL) + + # 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, 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/download.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/download.sh new file mode 100755 index 00000000..340bcc3b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/download.sh @@ -0,0 +1,66 @@ +#!/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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/nextqa.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/nextqa.py new file mode 100644 index 00000000..4db6caa1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/nextqa.py @@ -0,0 +1,159 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/perf.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/perf.py new file mode 100644 index 00000000..2349af4b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/hicache/perf.py @@ -0,0 +1,248 @@ +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/README.md new file mode 100644 index 00000000..26acabae --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/README.md @@ -0,0 +1,60 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/bench_other.py new file mode 100644 index 00000000..87051ea8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/bench_other.py @@ -0,0 +1,98 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/bench_sglang.py new file mode 100644 index 00000000..9aab11e4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/bench_sglang.py @@ -0,0 +1,101 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/build_dataset.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/build_dataset.py new file mode 100644 index 00000000..1396e5ed --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_decode_regex/build_dataset.py @@ -0,0 +1,58 @@ +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") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/README.md new file mode 100644 index 00000000..38fb67e8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/README.md @@ -0,0 +1,88 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/bench_other.py new file mode 100644 index 00000000..a64e950d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/bench_other.py @@ -0,0 +1,288 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/bench_sglang.py new file mode 100644 index 00000000..29f635f7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/bench_sglang.py @@ -0,0 +1,143 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/build_dataset.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/build_dataset.py new file mode 100644 index 00000000..1396e5ed --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/build_dataset.py @@ -0,0 +1,58 @@ +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") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/dataset.txt b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/dataset.txt new file mode 100644 index 00000000..c12421e5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_jump_forward/dataset.txt @@ -0,0 +1,50 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_schema/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_schema/README.md new file mode 100644 index 00000000..e6790e1f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_schema/README.md @@ -0,0 +1,15 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_schema/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_schema/bench_sglang.py new file mode 100644 index 00000000..8de68df3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/json_schema/bench_sglang.py @@ -0,0 +1,146 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/all_reduce/benchmark_aiter.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/all_reduce/benchmark_aiter.py new file mode 100644 index 00000000..bca45620 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/all_reduce/benchmark_aiter.py @@ -0,0 +1,330 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/all_reduce/benchmark_mscclpp.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/all_reduce/benchmark_mscclpp.py new file mode 100644 index 00000000..eebbd00c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/all_reduce/benchmark_mscclpp.py @@ -0,0 +1,224 @@ +"""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") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/all_reduce/benchmark_torch_symm_mem.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/all_reduce/benchmark_torch_symm_mem.py new file mode 100644 index 00000000..030fd5bb --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/all_reduce/benchmark_torch_symm_mem.py @@ -0,0 +1,251 @@ +"""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") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py new file mode 100644 index 00000000..b61240b0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py @@ -0,0 +1,403 @@ +import itertools +import math + +import cudnn +import torch +import torch.utils.benchmark as benchmark +from flashinfer import BatchDecodeWithPagedKVCacheWrapper + +from sglang.srt.layers.attention.triton_ops.decode_attention import decode_attention_fwd +from sglang.srt.utils import should_use_tensor_core + + +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, + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepep/deepep_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepep/deepep_utils.py new file mode 100644 index 00000000..169529ef --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepep/deepep_utils.py @@ -0,0 +1,218 @@ +# 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepep/tuning_deepep.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepep/tuning_deepep.py new file mode 100644 index 00000000..db08a8f1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepep/tuning_deepep.py @@ -0,0 +1,476 @@ +# 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 = ( + 4096, + 7168, + min(num_nodes, 4), + 8, + (256 // 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("--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 + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/README.md new file mode 100644 index 00000000..fba775f2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/README.md @@ -0,0 +1,19 @@ +## 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py new file mode 100644 index 00000000..bd02e2ae --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py @@ -0,0 +1,401 @@ +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.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 = triton.testing.do_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 = triton.testing.do_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 = triton.testing.do_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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py new file mode 100644 index 00000000..de14bd90 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py @@ -0,0 +1,329 @@ +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.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 = triton.testing.do_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 = triton.testing.do_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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py new file mode 100644 index 00000000..b2cea070 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py @@ -0,0 +1,487 @@ +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.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 = triton.testing.do_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 = triton.testing.do_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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/elementwise/benchmark_concat_mla.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/elementwise/benchmark_concat_mla.py new file mode 100644 index 00000000..c4d7bb1c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/elementwise/benchmark_concat_mla.py @@ -0,0 +1,198 @@ +import torch +import triton +import triton.language as tl +from sgl_kernel import concat_mla_k as concat_mla_k_cuda + +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 = triton.testing.do_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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/flashinfer_allreduce_fusion/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/flashinfer_allreduce_fusion/README.md new file mode 100644 index 00000000..e651604c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/flashinfer_allreduce_fusion/README.md @@ -0,0 +1,102 @@ +# 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py new file mode 100644 index 00000000..4aebf62b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py @@ -0,0 +1,1304 @@ +# 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 sgl_kernel 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/README.md new file mode 100644 index 00000000..f11c6541 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/README.md @@ -0,0 +1,210 @@ +## 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py new file mode 100644 index 00000000..b418855a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py @@ -0,0 +1,239 @@ +# 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.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, select_experts + + +def fused_moe_triton_api( + x, + w1, + w2, + input_gating, + topk, +): + topk_op = TopK( + top_k=topk, + renormalize=False, + use_grouped_topk=False, + ) + topk_op.use_triton_kernels = True + 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 = triton.testing.do_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() + + 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=args.ep_size, + pipeline_model_parallel_size=args.tp_size, + ) + + 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py new file mode 100644 index 00000000..2b4faa24 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py @@ -0,0 +1,305 @@ +# 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.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 = triton.testing.do_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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py new file mode 100644 index 00000000..206ee2a8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py @@ -0,0 +1,264 @@ +# 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.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 = triton.testing.do_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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/common_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/common_utils.py new file mode 100644 index 00000000..40f8697f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/common_utils.py @@ -0,0 +1,260 @@ +import json +from typing import Dict, List, TypedDict + +import torch +from transformers import AutoConfig + +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 + + +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 = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + + 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 + + architecture = config.architectures[0] + + # 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() + + 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", + ]: + E = config.num_experts // ep_size + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + elif architecture in [ + "DeepseekV2ForCausalLM", + "DeepseekV3ForCausalLM", + "Glm4MoeForCausalLM", + ]: + E = (config.n_routed_experts // ep_size) + ( + 0 + if disable_shared_experts_fusion + or architecture not in ["DeepseekV3ForCausalLM", "Glm4MoeForCausalLM"] + 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 + 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": config.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, + 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, + ) + + # 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, + 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, + ] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_client.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_client.py new file mode 100644 index 00000000..68cbfa73 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_client.py @@ -0,0 +1,71 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py new file mode 100644 index 00000000..aef7ed8f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py @@ -0,0 +1,458 @@ +# 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.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, + 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, + ) + 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_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, + 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) -> 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]) + + 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, + 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 + ) + # 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 + op_config = get_moe_configs( + num_experts, + shard_intermediate_size // 2, + 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, + 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, + 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, + 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): + print(args) + + 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" + 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) 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, + 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, + 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, + 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"], + 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py new file mode 100644 index 00000000..afee7994 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py @@ -0,0 +1,694 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py +import argparse +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 sgl_kernel import silu_and_mul + +from sglang.srt.layers.moe.fused_moe_triton import override_config +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.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, + topk_ids_dir: str, + block_shape: List[int] = None, + 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, + ) + 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_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) + + def prepare(i: int): + input_gating = gating_output[i] + topk_ids = torch.load(f"{topk_ids_dir}/topk_ids_layer{i%58+3}_idx{i//58}.pt") + new_topk_output = select_experts(hidden_states, input_gating, topk_config) + topk_output.topk_weights.copy_(new_topk_output.topk_weights) + tokens, _topk = topk_output.topk_ids.shape + topk_output.topk_ids.copy_(topk_ids[:tokens, :_topk]) + topk_output.router_logits.copy_(new_topk_output.router_logits) + + moe_use_tma = False + + def run(): + moe_runner_config = MoeRunnerConfig( + inplace=True, + ) + topk_weights, topk_ids, _ = topk_output + + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, config["BLOCK_SIZE_M"], num_experts + ) + M = hidden_states.shape[0] + E, N, _ = w1.shape + + topk = topk_ids.shape[1] + 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]), + ) + + compute_type = ( + tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16 + ) + apply_router_weight_on_input = moe_runner_config.apply_router_weight_on_input + + with override_config(config): + 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(10 if not ncu_enable else 1): + invoke_fused_moe_kernel( + hidden_states, + w1, + None, + intermediate_cache1, + None, + w1_scale, + None, + topk_weights, + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + apply_router_weight_on_input, + topk_ids.shape[1], + config, + compute_type=compute_type, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=False, + use_int8_w8a16=False, + use_int4_w4a16=False, + per_channel_quant=False, + block_shape=block_shape, + b_use_tma=moe_use_tma, + c_sorted=moe_use_tma, + filter_expert=False, + ) + end_event.record() + end_event.synchronize() + time_cost0 = start_event.elapsed_time(end_event) + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start_event.record() + + silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2) + for _ in range(10 if not ncu_enable else 1): + invoke_fused_moe_kernel( + intermediate_cache2, + w2, + None, + intermediate_cache3, + a2_scale, + w2_scale, + None, + topk_weights, + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + not apply_router_weight_on_input, + 1, + config, + compute_type=compute_type, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=False, + use_int8_w8a16=False, + use_int4_w4a16=False, + per_channel_quant=False, + block_shape=block_shape, + a_use_tma=moe_use_tma, + b_use_tma=moe_use_tma, + filter_expert=False, + ) + end_event.record() + end_event.synchronize() + time_cost1 = start_event.elapsed_time(end_event) + return time_cost0, time_cost1 + + # JIT compilation & warmup + if not ncu_enable: + moe_use_tma = False + run() + moe_use_tma = True + run() + latencies: List[float] = [] + latencies1: List[float] = [] + latencies_tma: List[float] = [] + latencies1_tma: List[float] = [] + + for i in range(num_iters): + prepare(i) + torch.cuda.synchronize() + moe_use_tma = False + t0, t1 = run() + torch.cuda.synchronize() + latencies.append(t0) + latencies1.append(t1) + + moe_use_tma = True + t0, t1 = run() + torch.cuda.synchronize() + latencies_tma.append(t0) + latencies1_tma.append(t1) + + avg = sum(latencies) / (num_iters * 10) * 1000 # us + avg_tma = sum(latencies_tma) / (num_iters * 10) * 1000 # us + avg1 = sum(latencies1) / (num_iters * 10) * 1000 # us + avg1_tma = sum(latencies1_tma) / (num_iters * 10) * 1000 # us + + return avg, avg_tma, avg1, avg1_tma + + +class BestConfigTrace: + def __init__(self, name): + self.name = name + self.config = None + self.time_cost = float("inf") + self.time_cost_all = None # kernel0 without tma,, kernel0 with tma, kernel1 without tma, kernel1 with tma + + def update(self, config, time_cost, time_cost_all): + if time_cost < self.time_cost: + print( + f"New best config for {self.name}: {config}, {time_cost=}, {time_cost_all=}, org: {self.config}, {self.time_cost_all}", + flush=True, + ) + self.config = config + self.time_cost = time_cost + self.time_cost_all = time_cost_all + + @property + def total_time(self): + return self.time_cost_all[0] + min(self.time_cost_all[2], self.time_cost_all[3]) + + def config_dict(self, down_moe=False): + if not down_moe: + return self.config + else: + return { + **self.config, + "USE_TMA": self.time_cost_all[2] > self.time_cost_all[3], + } + + +class BenchmarkWorker: + + def __init__(self, seed: int) -> 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]) + + 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, + block_shape: List[int], + cfg: Dict[str, int], + topk_ids_dir: str, + ) -> 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 + ) + # 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 + 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, + topk_ids_dir, + block_shape, + ) + 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, + block_shape: List[int], + search_space: List[Dict[str, int]], + topk_ids_dir: str, + ) -> Dict[str, int]: + trace0 = BestConfigTrace("kernel0") + trace1 = BestConfigTrace("kernel1") + trace2 = BestConfigTrace("kernel all") + + 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, + topk_ids_dir, + block_shape, + num_iters=10, + ) + except triton.runtime.autotuner.OutOfResources: + # Some configurations may be invalid and fail to compile. + continue + kt0 = kt0_no_tma + kt1 = min(kt1_no_tma, kt1_tma) + trace0.update( + config, + kt0, + (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma), + ) + trace1.update( + config, + kt1, + (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma), + ) + trace2.update( + config, + kt0 + kt1, + (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma), + ) + + now = datetime.now() + print(f"{now.ctime()}] Completed tuning for batch_size={num_tokens}") + assert trace0.config is not None + assert trace1.config is not None + print( + f"{num_tokens=}, {trace0.config=}, {trace0.time_cost_all=}, {trace1.config=}, {trace1.time_cost_all=}" + ) + if trace0.config["BLOCK_SIZE_M"] != trace1.config["BLOCK_SIZE_M"]: + best_trace = trace0 if trace0.total_time < trace1.total_time else trace1 + best_trace = ( + best_trace if best_trace.total_time < trace2.total_time else trace2 + ) + return ( + best_trace.config_dict(), + best_trace.config_dict(True), + best_trace.time_cost_all, + best_trace.time_cost_all, + ) + return ( + trace0.config_dict(), + trace1.config_dict(True), + trace0.time_cost_all, + trace1.time_cost_all, + ) + + +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, + 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, + ) + + # 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) + + 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" + + 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 len(batch_sizes) == 1: + worker = BenchmarkWorker(args.seed) + 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, + block_shape, + search_space, + topk_ids_dir, + ) + 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, + block_shape, + cfg, + topk_ids_dir, + ) + 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) + 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, + 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, + block_shape, + search_space, + topk_ids_dir, + ) + 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, + 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, + 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"], + 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) + args = parser.parse_args() + + main(args) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_text.json b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_text.json new file mode 100644 index 00000000..80242160 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/fused_moe_triton/tuning_text.json @@ -0,0 +1 @@ +{"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"} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/quantization/bench_fp4_quant.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/quantization/bench_fp4_quant.py new file mode 100644 index 00000000..afc12dd8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/quantization/bench_fp4_quant.py @@ -0,0 +1,136 @@ +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.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 = triton.testing.do_bench_cudagraph( + 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 = triton.testing.do_bench_cudagraph( + lambda: scaled_fp4_grouped_quantize( + silu_and_mul(x), + masks, + glb_scales, + ), + quantiles=quantiles, + ) + if provider == "cuda_fused_fp4": + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/quantization/bench_int8_quant.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/quantization/bench_int8_quant.py new file mode 100644 index 00000000..94b79569 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/quantization/bench_int8_quant.py @@ -0,0 +1,94 @@ +import argparse + +import torch +import triton +from vllm._custom_ops import scaled_int8_quant as vllm_scaled_int8_quant + +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 = triton.testing.do_bench( + lambda: vllm_scaled_int8_quant(x, symmetric=True), + quantiles=quantiles, + ) + if provider == "triton": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: per_token_quant_int8(x), + quantiles=quantiles, + ) + if provider == "torch.compile": + ms, min_ms, max_ms = triton.testing.do_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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/quantization/tuning_block_wise_kernel.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/quantization/tuning_block_wise_kernel.py new file mode 100644 index 00000000..1b51e54b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/quantization/tuning_block_wise_kernel.py @@ -0,0 +1,474 @@ +# 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) + + 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, + ) + + 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", +) -> 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}...") + + with open(config_file_path, "w") as f: + json.dump(configs, f, indent=4) + f.write("\n") + + +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"] + + 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) + + 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 + + weight_shapes = get_weight_shapes(args.tp_size) + + batches_per_gpu = distribute_batch_sizes(batch_sizes, num_gpus) + + 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, + } + ) + + ctx = mp.get_context("spawn") + 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) + 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() + + main(args) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py new file mode 100644 index 00000000..3e17205e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py @@ -0,0 +1,169 @@ +import os + +import torch +import triton +import triton.language as tl + + +@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 = triton.testing.do_bench( + lambda: get_last_loc_torch(req_to_token, req_pool_indices, pre_lens), + quantiles=quantiles, + ) + elif provider == "triton": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: get_last_loc_triton(req_to_token, req_pool_indices, pre_lens), + quantiles=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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py new file mode 100644 index 00000000..1ce43c8b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py @@ -0,0 +1,342 @@ +import itertools +import os + +import torch +import triton +import triton.language as tl + + +@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 = triton.testing.do_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=quantiles, + ) + elif provider == "triton": + ms, min_ms, max_ms = triton.testing.do_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=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 = triton.testing.do_bench( + run_optimized, quantiles=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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py new file mode 100644 index 00000000..98144d47 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py @@ -0,0 +1,283 @@ +import itertools + +import torch +import torch.nn.functional as F +import triton.testing as tt + +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 = tt.do_bench(lambda: _run_triton(inputs), warmup=warmup, rep=rep) + elif provider == "torch": + ms = tt.do_bench(lambda: _run_torch_ref(inputs), warmup=warmup, rep=rep) + else: + raise ValueError(provider) + + return ms + + +if __name__ == "__main__": + bench.run(print_data=True, show_plots=False) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/line_retrieval/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/line_retrieval/README.md new file mode 100644 index 00000000..3c81a63c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/line_retrieval/README.md @@ -0,0 +1,37 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/line_retrieval/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/line_retrieval/bench_sglang.py new file mode 100644 index 00000000..e974e7dd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/line_retrieval/bench_sglang.py @@ -0,0 +1,149 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/line_retrieval/gen_data.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/line_retrieval/gen_data.py new file mode 100644 index 00000000..a01d40a0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/line_retrieval/gen_data.py @@ -0,0 +1,139 @@ +""" +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/README.md new file mode 100644 index 00000000..a71f339d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/README.md @@ -0,0 +1,61 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_hf_llava_bench.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_hf_llava_bench.sh new file mode 100755 index 00000000..a51a715b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_hf_llava_bench.sh @@ -0,0 +1,9 @@ +#!/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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_hf_mme.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_hf_mme.sh new file mode 100755 index 00000000..6ed332fe --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_hf_mme.sh @@ -0,0 +1,9 @@ +#!/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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_sglang.py new file mode 100644 index 00000000..b9e8c140 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_sglang.py @@ -0,0 +1,96 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_sglang_mme.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_sglang_mme.sh new file mode 100644 index 00000000..8c8c09cc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/bench_sglang_mme.sh @@ -0,0 +1,2 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/download_images.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/download_images.py new file mode 100644 index 00000000..701190a0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llava_bench/download_images.py @@ -0,0 +1,20 @@ +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.") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llm_judge/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llm_judge/README.md new file mode 100644 index 00000000..164b19a0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llm_judge/README.md @@ -0,0 +1,33 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llm_judge/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llm_judge/bench_other.py new file mode 100644 index 00000000..8e602906 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llm_judge/bench_other.py @@ -0,0 +1,151 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llm_judge/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llm_judge/bench_sglang.py new file mode 100644 index 00000000..97e6c397 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/llm_judge/bench_sglang.py @@ -0,0 +1,97 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/README.md new file mode 100644 index 00000000..37fceee1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/README.md @@ -0,0 +1,33 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/bench_other.py new file mode 100644 index 00000000..0ad38a01 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/bench_other.py @@ -0,0 +1,89 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/bench_sglang.py new file mode 100644 index 00000000..8394cfc2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/bench_sglang.py @@ -0,0 +1,81 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/build_dataset.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/build_dataset.py new file mode 100644 index 00000000..7b15d395 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/long_json_decode/build_dataset.py @@ -0,0 +1,27 @@ +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") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/lora/launch_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/lora/launch_server.py new file mode 100644 index 00000000..5dcf66ad --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/lora/launch_server.py @@ -0,0 +1,82 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/lora/lora_bench.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/lora/lora_bench.py new file mode 100644 index 00000000..4f380c70 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/lora/lora_bench.py @@ -0,0 +1,475 @@ +# 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, + get_tokenizer, + remove_prefix, + sample_random_requests, +) + +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/README.md new file mode 100644 index 00000000..16de20cd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/README.md @@ -0,0 +1,59 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/bench_other.py new file mode 100644 index 00000000..f1b166c2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/bench_other.py @@ -0,0 +1,173 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/bench_sglang.py new file mode 100644 index 00000000..23057be4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/bench_sglang.py @@ -0,0 +1,181 @@ +import argparse +import json +import os +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, +) + +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 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="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) + main(args) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/download_data.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/download_data.sh new file mode 100644 index 00000000..fb224165 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmlu/download_data.sh @@ -0,0 +1,2 @@ +wget https://people.eecs.berkeley.edu/~hendrycks/data.tar +tar xf data.tar diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/README.md new file mode 100644 index 00000000..61fea8bc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/README.md @@ -0,0 +1,49 @@ +## 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/bench_hf.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/bench_hf.py new file mode 100644 index 00000000..c841f444 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/bench_hf.py @@ -0,0 +1,169 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/bench_sglang.py new file mode 100644 index 00000000..9a0bf452 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/bench_sglang.py @@ -0,0 +1,216 @@ +""" +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 re +import sys +import time +import traceback +from dataclasses import dataclass, field +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, 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"] + extra_body = None if lora_path is None else {"lora_path": lora_path} + response = await client.chat.completions.create( + model="default", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": prefix}, + {"type": "image_url", "image_url": {"url": image_path}}, + {"type": "text", "text": suffix}, + ], + } + ], + temperature=0, + max_completion_tokens=sampling_params["max_new_tokens"], + max_tokens=sampling_params["max_new_tokens"], + extra_body=extra_body, + ) + return sample, response.choices[0].message.content + + +async def process_sample_with_semaphore( + semaphore: asyncio.Semaphore, + client: Any, + sample: dict, + sampling_params: dict, + 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, 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) + 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, 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, 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() + 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/data_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/data_utils.py new file mode 100644 index 00000000..8c36768d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/data_utils.py @@ -0,0 +1,215 @@ +"""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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/eval_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/eval_utils.py new file mode 100644 index 00000000..955a3bfa --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/eval_utils.py @@ -0,0 +1,683 @@ +"""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: int = 30 + response_answer_regex: str = "(.*)" + lora_path: 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( + "--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`.", + ) + + @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): + max_new_tokens = eval_args.max_new_tokens + temperature = 0.001 + + extra_request_body = {} + if eval_args.extra_request_body: + extra_request_body = json.loads(eval_args.extra_request_body) + + return { + "temperature": temperature, + "max_new_tokens": max_new_tokens, + **extra_request_body, + } + + +# ----------- 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}") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/prompt_format.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/prompt_format.yaml new file mode 100644 index 00000000..1a0f7211 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mmmu/prompt_format.yaml @@ -0,0 +1,15 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/README.md new file mode 100644 index 00000000..fc37caee --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/README.md @@ -0,0 +1,48 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/bench_other.py new file mode 100644 index 00000000..5e579e9a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/bench_other.py @@ -0,0 +1,111 @@ +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 + + +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): + questions = load_questions(args.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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/bench_sglang.py new file mode 100644 index 00000000..0d0545b3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/bench_sglang.py @@ -0,0 +1,99 @@ +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, +) + + +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): + # Construct prompts + questions = load_questions(args.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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/bench_sglang_eagle.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/bench_sglang_eagle.py new file mode 100644 index 00000000..7ad3dcb1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/mtbench/bench_sglang_eagle.py @@ -0,0 +1,137 @@ +""" +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, +) + + +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, + "prompt": [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): + # Construct prompts + questions = load_questions(args.question_file)[: args.num_questions] + arguments = [ + {"question_1": q["prompt"][0], "question_2": q["prompt"][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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_chain_reasoning/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_chain_reasoning/README.md new file mode 100644 index 00000000..4c9f740f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_chain_reasoning/README.md @@ -0,0 +1,49 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_chain_reasoning/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_chain_reasoning/bench_other.py new file mode 100644 index 00000000..f361496a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_chain_reasoning/bench_other.py @@ -0,0 +1,186 @@ +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 = 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_chain_reasoning/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_chain_reasoning/bench_sglang.py new file mode 100644 index 00000000..1d3129db --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_chain_reasoning/bench_sglang.py @@ -0,0 +1,140 @@ +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 = 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/README.md new file mode 100644 index 00000000..09f49c78 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/README.md @@ -0,0 +1,47 @@ +## 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) +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/bench_other.py new file mode 100644 index 00000000..627837c5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/bench_other.py @@ -0,0 +1,114 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/bench_sglang.py new file mode 100644 index 00000000..0b4b0dbc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/bench_sglang.py @@ -0,0 +1,93 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/build_dataset.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/build_dataset.py new file mode 100644 index 00000000..27df9474 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_document_qa/build_dataset.py @@ -0,0 +1,70 @@ +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" + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/README.md new file mode 100644 index 00000000..0fb5b21f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/README.md @@ -0,0 +1,66 @@ +### 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/bench_other.py new file mode 100644 index 00000000..9189af5b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/bench_other.py @@ -0,0 +1,93 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/bench_sglang.py new file mode 100644 index 00000000..1051bf19 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/bench_sglang.py @@ -0,0 +1,79 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/data_gen.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/data_gen.py new file mode 100644 index 00000000..043c07a7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/data_gen.py @@ -0,0 +1,29 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/long_prompt_multi_turn.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/long_prompt_multi_turn.py new file mode 100644 index 00000000..88eba70c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/multi_turn_chat/long_prompt_multi_turn.py @@ -0,0 +1,129 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/prefill_only/bench_embeddings.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/prefill_only/bench_embeddings.py new file mode 100644 index 00000000..74d8a582 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/prefill_only/bench_embeddings.py @@ -0,0 +1,159 @@ +""" +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()) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/prefill_only/bench_score.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/prefill_only/bench_score.py new file mode 100644 index 00000000..117335ea --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/prefill_only/bench_score.py @@ -0,0 +1,192 @@ +""" +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()) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/prefill_only/util.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/prefill_only/util.py new file mode 100644 index 00000000..3b385591 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/prefill_only/util.py @@ -0,0 +1,813 @@ +""" +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" 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, + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/react/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/react/README.md new file mode 100644 index 00000000..51bcaa44 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/react/README.md @@ -0,0 +1,34 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/react/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/react/bench_other.py new file mode 100644 index 00000000..08666662 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/react/bench_other.py @@ -0,0 +1,202 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/react/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/react/bench_sglang.py new file mode 100644 index 00000000..331638e9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/react/bench_sglang.py @@ -0,0 +1,153 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/README.md new file mode 100644 index 00000000..f9d26e0c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/README.md @@ -0,0 +1,77 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/answer_extraction.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/answer_extraction.py new file mode 100644 index 00000000..45ce59af --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/answer_extraction.py @@ -0,0 +1,269 @@ +# 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+" + ans = re.findall(pattern, pred_str.replace(",", "")) + if len(ans) >= 1: + ans = ans[-1] + else: + ans = "" + if ans: + pred.append(ans) + + # multiple line + _pred = [] + for ans in pred: + ans = ans.strip().split("\n")[0] + ans = ans.lstrip(":") + ans = ans.rstrip(".") + ans = ans.rstrip("/") + ans = strip_string(ans) + _pred.append(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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/bench_sglang.py new file mode 100644 index 00000000..ccbff9d1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/bench_sglang.py @@ -0,0 +1,135 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/eval_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/eval_utils.py new file mode 100644 index 00000000..ab736954 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/reasoning_benchmark/eval_utils.py @@ -0,0 +1,206 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/.gitignore b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/.gitignore new file mode 100644 index 00000000..3322de9b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/.gitignore @@ -0,0 +1 @@ +!topic.jsonl diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/README.md new file mode 100644 index 00000000..97ebe1d8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/README.md @@ -0,0 +1,33 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/bench_other.py new file mode 100644 index 00000000..2630081b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/bench_other.py @@ -0,0 +1,133 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/bench_sglang.py new file mode 100644 index 00000000..86c476f9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/bench_sglang.py @@ -0,0 +1,100 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/lmql_funcs.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/lmql_funcs.py new file mode 100644 index 00000000..7790bbe9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/lmql_funcs.py @@ -0,0 +1,50 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/topic.jsonl b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/topic.jsonl new file mode 100644 index 00000000..c8ac6d12 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tip_suggestion/topic.jsonl @@ -0,0 +1,50 @@ +{"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} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/README.md new file mode 100644 index 00000000..bf5ab163 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/README.md @@ -0,0 +1,51 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/bench_other.py new file mode 100644 index 00000000..0ef8c636 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/bench_other.py @@ -0,0 +1,222 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/bench_sglang.py new file mode 100644 index 00000000..bcdb6e54 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/bench_sglang.py @@ -0,0 +1,171 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/lmql_funcs.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/lmql_funcs.py new file mode 100644 index 00000000..c783cdbe --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_deep/lmql_funcs.py @@ -0,0 +1,82 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_v0/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_v0/README.md new file mode 100644 index 00000000..821bb20d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_v0/README.md @@ -0,0 +1,43 @@ +## 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_v0/bench_other.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_v0/bench_other.py new file mode 100644 index 00000000..703ecd7f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_v0/bench_other.py @@ -0,0 +1,179 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_v0/bench_sglang.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_v0/bench_sglang.py new file mode 100644 index 00000000..6d7575f3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/benchmark/tree_of_thought_v0/bench_sglang.py @@ -0,0 +1,159 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/Dockerfile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/Dockerfile new file mode 100644 index 00000000..7121a4d5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/Dockerfile @@ -0,0 +1,326 @@ +ARG CUDA_VERSION=12.9.1 +FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu22.04 AS base + +ARG TARGETARCH +ARG BUILD_TYPE=all +ARG BRANCH_TYPE=remote +ARG GRACE_BLACKWELL=0 + +ARG GRACE_BLACKWELL_DEEPEP_BRANCH=gb200_blog_part_2 +ARG DEEPEP_COMMIT=9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee +ARG BUILD_AND_DOWNLOAD_PARALLEL=8 +ARG SGL_KERNEL_VERSION=0.3.18.post2 +ARG SGL_VERSION=0.5.6 +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.5.3 +ARG NVSHMEM_VERSION=3.4.5 + +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 it is 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 + +RUN --mount=type=cache,target=/var/cache/apt apt update && apt install wget -y && apt install software-properties-common -y \ + && add-apt-repository ppa:deadsnakes/ppa -y \ + && apt install python3.12-full python3.12-dev python3.10-venv -y \ + && 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 https://bootstrap.pypa.io/get-pip.py \ + && python3 get-pip.py \ + # Fix for `apt-add-repository` + && cd /usr/lib/python3/dist-packages/ \ + && ln -s apt_pkg.cpython-310-*-linux-gnu.so apt_pkg.so + +# Set timezone and install all packages +RUN --mount=type=cache,target=/var/cache/apt echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ + && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ + && apt-get update && apt-get install -y --no-install-recommends \ + tzdata \ + software-properties-common netcat-openbsd kmod unzip openssh-server \ + curl wget lsof zsh ccache tmux htop git-lfs tree \ + build-essential cmake perl \ + libopenmpi-dev libnuma1 libnuma-dev numactl \ + libibverbs-dev libibverbs1 libibumad3 \ + librdmacm1 libnl-3-200 libnl-route-3-200 libnl-route-3-dev libnl-3-dev \ + ibverbs-providers infiniband-diags perftest \ + 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 \ + patchelf \ + nvidia-dkms-550 \ + 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 it is 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 \ + && wget -q 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 + +FROM scratch AS local_src +COPY . /src + +FROM base AS build-image +# Install SGLang +# Until torch 2.9 and cu13 are stable we manually update torch if you are on CUDA 13 +WORKDIR /sgl-workspace +ARG BRANCH_TYPE +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; \ + 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}/sgl_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 sgl-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}/sgl_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 + +# We use Tom's DeepEP fork for GB200 for now; the 1fd57b0276311d035d16176bb0076426166e52f3 commit is https://github.com/fzyzcjy/DeepEP/tree/gb200_blog_part_2 +RUN set -eux; \ + if [ "${CUDA_VERSION%%.*}" = "12" ]; then \ + pip install nvidia-nvshmem-cu12==${NVSHMEM_VERSION} ; \ + elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \ + pip install nvidia-nvshmem-cu13==${NVSHMEM_VERSION} ; \ + else \ + echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ; \ + fi && \ + 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 && \ + cd .. ; \ + else \ + wget -q https://${GITHUB_ARTIFACTORY}/deepseek-ai/DeepEP/archive/${DEEPEP_COMMIT}.zip && \ + unzip ${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 && \ + cd .. ; \ + fi + +# Install DeepEP +# CTK13 requires the cccl include +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|12.9.1|13.0.1) \ + 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 . + +# Python tools +RUN --mount=type=cache,target=/root/.cache/pip python3 -m pip install \ + datamodel_code_generator \ + mooncake-transfer-engine==0.3.7.post2 \ + pre-commit \ + pytest \ + black \ + isort \ + icdiff \ + uv \ + wheel \ + scikit-build-core \ + nixl \ + py-spy + +# Some patching packages +# TODO: Remove this when torch version covers these packages +# Move cutlass-dsl to pyproject.toml after drivers on CI runners are updated +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 nvidia-nvshmem-cu12==${NVSHMEM_VERSION} --force-reinstall --no-deps; \ + python3 -m pip install nvidia-cutlass-dsl==4.3.0 --force-reinstall --no-deps; \ +elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \ + python3 -m pip install nvidia-nccl-cu13==2.28.3 --force-reinstall --no-deps ; \ + python3 -m pip install nvidia-cublas==13.1.0.3 --force-reinstall --no-deps ; \ + python3 -m pip install nixl-cu13 ; \ + python3 -m pip install nvidia-nvshmem-cu13==${NVSHMEM_VERSION} --force-reinstall --no-deps; \ + python3 -m pip install nvidia-cutlass-dsl==4.3.1 --force-reinstall --no-deps; \ +fi + +# Install development tools and utilities +RUN --mount=type=cache,target=/var/cache/apt apt-get update && apt-get install -y \ + gdb \ + ninja-build \ + vim \ + tmux \ + htop \ + wget \ + curl \ + locales \ + lsof \ + git \ + git-lfs \ + zsh \ + tree \ + silversearcher-ag \ + cloc \ + unzip \ + pkg-config \ + libssl-dev \ + bear \ + ccache \ + less \ + gnuplot \ + && apt install -y rdma-core infiniband-diags openssh-server perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 \ + && rm -rf /var/lib/apt/lists/* \ + && apt-get clean + +RUN --mount=type=cache,target=/var/cache/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 nsight-systems-cli -y + +# Set up locale +RUN 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 minimal Python 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 + +# Install diff-so-fancy +RUN curl -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 + +# Install clang-format +RUN curl -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 + +# Install clangd +RUN curl -L https://${GITHUB_ARTIFACTORY}/clangd/clangd/releases/download/18.1.3/clangd-linux-18.1.3.zip -o clangd.zip \ + && unzip 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 + +# Install CMake +RUN CMAKE_VERSION=3.31.1 \ + && ARCH=$(uname -m) \ + && CMAKE_INSTALLER="cmake-${CMAKE_VERSION}-linux-${ARCH}" \ + && wget -q "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" + +# Build and install sgl-router (Rust toolchain removed after build to save space) +RUN --mount=type=cache,target=/root/.cache/pip curl --proto '=https' --tlsv1.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-router/bindings/python \ + && ulimit -n 65536 && maturin build --release --features vendored-openssl --out dist \ + && python3 -m pip install --force-reinstall dist/*.whl \ + && rm -rf /root/.cargo /root/.rustup target dist ~/.cargo \ + && sed -i '/\.cargo\/env/d' /root/.profile /root/.bashrc /root/.zshenv 2>/dev/null || true + + +# 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 -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 + +# Configure Vim and tmux +COPY docker/configs/.vimrc /root/.vimrc +COPY docker/configs/.tmux.conf /root/.tmux.conf + +# Configure Git +COPY docker/configs/.gitconfig /tmp/.gitconfig +RUN cat /tmp/.gitconfig >> /root/.gitconfig && rm /tmp/.gitconfig + +# Configure zsh +COPY docker/configs/.zshrc /root/.zshrc + +RUN set -euxo ; \ + curl --proto '=https' --tlsv1.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 + +# Set workspace directory +WORKDIR /sgl-workspace/sglang diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/compose.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/compose.yaml new file mode 100644 index 00000000..f7ff1fbd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/compose.yaml @@ -0,0 +1,35 @@ +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 + device_ids: ["0"] + capabilities: [gpu] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.gitconfig b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.gitconfig new file mode 100644 index 00000000..8150e40d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.gitconfig @@ -0,0 +1,30 @@ +[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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.tmux.conf b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.tmux.conf new file mode 100644 index 00000000..89f20064 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.tmux.conf @@ -0,0 +1,27 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.vimrc b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.vimrc new file mode 100644 index 00000000..d4414000 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.vimrc @@ -0,0 +1,45 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.zshrc b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.zshrc new file mode 100644 index 00000000..5c7113e0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/.zshrc @@ -0,0 +1,27 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/yank b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/yank new file mode 100755 index 00000000..c9de641b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/configs/yank @@ -0,0 +1,12 @@ +#!/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" ||: diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/diffusion.Dockerfile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/diffusion.Dockerfile new file mode 100644 index 00000000..d8af45b7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/diffusion.Dockerfile @@ -0,0 +1,104 @@ +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"] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/gateway.Dockerfile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/gateway.Dockerfile new file mode 100644 index 00000000..e63bf0db --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/gateway.Dockerfile @@ -0,0 +1,79 @@ +######################## 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 echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ + && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ + && 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-router + +# 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-router/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"] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/k8s-sglang-distributed-sts.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/k8s-sglang-distributed-sts.yaml new file mode 100644 index 00000000..4252363c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/k8s-sglang-distributed-sts.yaml @@ -0,0 +1,103 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/k8s-sglang-service.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/k8s-sglang-service.yaml new file mode 100644 index 00000000..866d50be --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/k8s-sglang-service.yaml @@ -0,0 +1,117 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/npu.Dockerfile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/npu.Dockerfile new file mode 100644 index 00000000..54261e70 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/npu.Dockerfile @@ -0,0 +1,100 @@ +ARG CANN_VERSION=8.3.rc1 +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 PIP_INDEX_URL="https://pypi.org/simple/" +ARG APTMIRROR="" +ARG PYTORCH_VERSION="2.8.0" +ARG TORCHVISION_VERSION="0.23.0" +ARG PTA_URL="https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/sglang/torch_npu/torch_npu-2.8.0.post2.dev20251113-cp311-cp311-manylinux_2_28_aarch64.whl" +ARG TRITON_ASCEND_URL="https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/sglang/triton_ascend/triton_ascend-3.2.0.dev2025112116-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl" +ARG BISHENG_NAME="Ascend-BiSheng-toolkit_aarch64_20251121.run" +ARG BISHENG_URL="https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/sglang/triton_ascend/${BISHENG_NAME}" +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 + +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 \ + 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} mf-adapter==1.0.0 +### Install SGLang Model Gateway +RUN ${PIP_INSTALL} sglang-router + + +### Install PyTorch and PTA +RUN (${PIP_INSTALL} torch==${PYTORCH_VERSION} torchvision==${TORCHVISION_VERSION} --index-url https://download.pytorch.org/whl/cpu) \ + && (${PIP_INSTALL} ${PTA_URL}) + + +# TODO: install from pypi released triton-ascend +RUN (${PIP_INSTALL} pybind11) \ + && (${PIP_INSTALL} ${TRITON_ASCEND_URL}) + +# Install SGLang +RUN git clone https://github.com/sgl-project/sglang --branch $SGLANG_TAG && \ + (cd sglang/python && rm -rf pyproject.toml && mv pyproject_other.toml pyproject.toml && ${PIP_INSTALL} -v .[srt_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 && git clone --branch $SGLANG_KERNEL_NPU_TAG https://github.com/sgl-project/sgl-kernel-npu.git \ + && export LD_LIBRARY_PATH=${ASCEND_CANN_PATH}/latest/runtime/lib64/stub:$LD_LIBRARY_PATH && \ + source ${ASCEND_CANN_PATH}/set_env.sh && \ + cd sgl-kernel-npu && \ + bash build.sh \ + && ${PIP_INSTALL} output/deep_ep*.whl output/sgl_kernel_npu*.whl \ + && cd .. && rm -rf sgl-kernel-npu \ + && cd "$(python3 -m pip show deep-ep | awk '/^Location:/ {print $2}')" && ln -s deep_ep/deep_ep_cpp*.so + +# Install CustomOps +RUN wget https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/ops/CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run && \ + chmod a+x ./CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run && \ + ./CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run --quiet --install-path=/usr/local/Ascend/ascend-toolkit/latest/opp && \ + wget https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/ops/custom_ops-1.0.$DEVICE_TYPE-cp311-cp311-linux_aarch64.whl && \ + ${PIP_INSTALL} ./custom_ops-1.0.$DEVICE_TYPE-cp311-cp311-linux_aarch64.whl + +# Install Bisheng +RUN wget -O "${BISHENG_NAME}" "${BISHENG_URL}" && chmod a+x "${BISHENG_NAME}" && "./${BISHENG_NAME}" --install && rm "${BISHENG_NAME}" + +CMD ["/bin/bash"] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/rocm.Dockerfile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/rocm.Dockerfile new file mode 100644 index 00000000..b7fcf770 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/rocm.Dockerfile @@ -0,0 +1,318 @@ +# Usage (to build SGLang ROCm docker image): +# docker build --build-arg SGL_BRANCH=v0.5.6 --build-arg GPU_ARCH=gfx942 -t v0.5.6-rocm630-mi30x -f rocm.Dockerfile . +# docker build --build-arg SGL_BRANCH=v0.5.6 --build-arg GPU_ARCH=gfx942-rocm700 -t v0.5.6-rocm700-mi30x -f rocm.Dockerfile . +# docker build --build-arg SGL_BRANCH=v0.5.6 --build-arg GPU_ARCH=gfx950 -t v0.5.6-rocm700-mi35x -f rocm.Dockerfile . + + +# Default base images +ARG BASE_IMAGE_942="rocm/sgl-dev:vllm20250114" +ARG BASE_IMAGE_942_ROCM700="rocm/sgl-dev:rocm7-vllm-20250904" +ARG BASE_IMAGE_950="rocm/sgl-dev:rocm7-vllm-20250904" + +# This is necessary for scope purpose +ARG GPU_ARCH=gfx950 + +# =============================== +# Base image 942 with rocm630 and args +FROM $BASE_IMAGE_942 AS gfx942 +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.4" +ENV NO_DEPS_FLAG="" + +# =============================== +# Base image 942 and args +FROM $BASE_IMAGE_942_ROCM700 AS gfx942-rocm700 +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.7.post1" +ENV NO_DEPS_FLAG="" + +# =============================== +# 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="0" +ENV BUILD_MOONCAKE="1" +ENV AITER_COMMIT="v0.1.7.post2" +ENV NO_DEPS_FLAG="" +# =============================== +# Chosen arch and args +FROM ${GPU_ARCH} + +# This is necessary for scope purpose, again +ARG GPU_ARCH=gfx950 +ENV GPU_ARCH_LIST=${GPU_ARCH%-*} + +ARG SGL_REPO="https://github.com/sgl-project/sglang.git" +ARG SGL_DEFAULT="main" +ARG SGL_BRANCH=${SGL_DEFAULT} + +ARG TRITON_REPO="https://github.com/ROCm/triton.git" +ARG TRITON_COMMIT="improve_fa_decode_3.0.0" + +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/HaiShaw/tilelang.git" +ARG TILELANG_BRANCH="dsv32-mi35x" +ARG TILELANG_COMMIT="ae938cf885743f165a19656d1122ad42bb0e30b8" + +ARG FHT_REPO="https://github.com/jeffdaily/fast-hadamard-transform.git" +ARG FHT_BRANCH="rocm" +ARG FHT_COMMIT="46efb7d776d38638fc39f3c803eaee3dd7016bd1" +USER root + +# 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)" + +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 +RUN pip uninstall -y aiter +RUN git clone ${AITER_REPO} \ + && cd aiter \ + && git checkout ${AITER_COMMIT} \ + && git submodule update --init --recursive +RUN cd aiter \ + && echo "[AITER] GPU_ARCH=${GPU_ARCH}" \ + && 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 develop"; \ + elif [ "$BUILD_AITER_ALL" = "1" ]; then \ + sh -c "PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST python setup.py develop"; \ + else \ + sh -c "GPU_ARCHS=$GPU_ARCH_LIST python setup.py develop"; \ + fi + +# ----------------------- +# Triton +RUN if [ "$BUILD_TRITON" = "1" ]; then \ + pip uninstall -y triton \ + && git clone ${TRITON_REPO} \ + && cd triton \ + && git checkout ${TRITON_COMMIT} \ + && cd python \ + && python setup.py install; \ + fi + +# ----------------------- +# Build vLLM +ARG VLLM_REPO="https://github.com/ROCm/vllm.git" +ARG VLLM_BRANCH="9f6b92db47c3444b7a7d67451ba0c3a2d6af4c2c" +RUN if [ "$BUILD_VLLM" = "1" ]; then \ + git clone ${VLLM_REPO} \ + && cd vllm \ + && git checkout ${VLLM_BRANCH} \ + && python -m pip install -r requirements/rocm.txt \ + && python setup.py clean --all \ + && python setup.py develop; \ + fi + +# ----------------------- +# 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 + +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 \ + python -m pip --no-cache-dir install -e "python[srt_hip]" ${NO_DEPS_FLAG}; \ + else \ + python -m pip --no-cache-dir install -e "python[all_hip]" ${NO_DEPS_FLAG}; \ + 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-router +ENV PATH="/root/.cargo/bin:${PATH}" +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ + && rustc --version && cargo --version + +# Build and install sgl-router +RUN python3 -m pip install --no-cache-dir setuptools-rust \ + && cd /sgl-workspace/sglang/sgl-router/bindings/python \ + && cargo build --release \ + && python3 -m pip install --no-cache-dir . \ + && 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; \ + # Build TileLang only for gfx950 + if [ "${GPU_ARCH:-}" != "gfx950" ]; then \ + echo "[TileLang] Skipping (GPU_ARCH=${GPU_ARCH:-unset})"; \ + exit 0; \ + fi; \ + 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 \ + gcc libtinfo-dev zlib1g-dev libedit-dev libxml2-dev \ + cmake ninja-build pkg-config libstdc++6 \ + && rm -rf /var/lib/apt/lists/*; \ + \ + # 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) + python3 -m pip install --upgrade "setuptools>=77.0.3,<80" wheel cmake ninja && \ + python3 -m 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.sh -o /tmp/llvm.sh; \ + chmod +x /tmp/llvm.sh; \ + /tmp/llvm.sh 18; \ + 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 + python3 -m pip install --no-cache-dir "cython>=0.29.36,<3.0"; \ + \ + # Clone + pin TileLang (bundled TVM), then build + git clone --recursive --branch "${TILELANG_BRANCH}" "${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="-DLLVM_CONFIG=${LLVM_CONFIG} ${CMAKE_ARGS:-}" && \ + bash ./install_rocm.sh' + +# ----------------------- +# 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 + +# ----------------------- +# Performance environment variable. + +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 VLLM_FP8_PADDING=1 +ENV VLLM_FP8_ACT_PADDING=1 +ENV VLLM_FP8_WEIGHT_PADDING=1 +ENV VLLM_FP8_REDUCE_CONV=1 +ENV TORCHINDUCTOR_MAX_AUTOTUNE=1 +ENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1 + +CMD ["/bin/bash"] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/sagemaker.Dockerfile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/sagemaker.Dockerfile new file mode 100644 index 00000000..5fbff509 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/sagemaker.Dockerfile @@ -0,0 +1,6 @@ +FROM lmsysorg/sglang:latest + +COPY serve /usr/bin/serve +RUN chmod 777 /usr/bin/serve + +ENTRYPOINT [ "/usr/bin/serve" ] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/serve b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/serve new file mode 100755 index 00000000..9f464bf4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/serve @@ -0,0 +1,34 @@ +#!/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[@]}" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/xeon.Dockerfile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/xeon.Dockerfile new file mode 100644 index 00000000..da2f4a0f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/xeon.Dockerfile @@ -0,0 +1,56 @@ +FROM ubuntu:24.04 +SHELL ["/bin/bash", "-c"] + +ARG SGLANG_REPO=https://github.com/sgl-project/sglang.git +ARG VER_SGLANG=main + +ARG VER_TORCH=2.9.0 +ARG VER_TORCHVISION=0.24.0 +ARG VER_TRITON=3.5.0 + +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 = "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 . && \ + uv pip install torch==${VER_TORCH} torchvision==${VER_TORCHVISION} triton==${VER_TRITON} --force-reinstall && \ + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/xpu.Dockerfile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/xpu.Dockerfile new file mode 100644 index 00000000..5aa57b3d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docker/xpu.Dockerfile @@ -0,0 +1,73 @@ +# 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.2.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.9.0+xpu torchao torchvision torchaudio pytorch-triton-xpu==3.5.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 --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"] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/Makefile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/Makefile new file mode 100644 index 00000000..6b8792c4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/Makefile @@ -0,0 +1,58 @@ +# 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 + +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/README.md new file mode 100644 index 00000000..2cad1335 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/README.md @@ -0,0 +1,55 @@ +# SGLang Documentation + +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 + +```bash +apt-get update && apt-get install -y 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/css/custom_log.css b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/css/custom_log.css new file mode 100644 index 00000000..61f65d01 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/css/custom_log.css @@ -0,0 +1,29 @@ +.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; +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/css/readthedocs.css b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/css/readthedocs.css new file mode 100644 index 00000000..aca6649b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/css/readthedocs.css @@ -0,0 +1,9 @@ +table.autosummary td { + width: 50% +} + +img.align-center { + display: block; + margin-left: auto; + margin-right: auto; +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/image/logo.ico b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/image/logo.ico new file mode 100644 index 00000000..7a0b5b2f Binary files /dev/null and b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/image/logo.ico differ diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/image/logo.png b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/image/logo.png new file mode 100644 index 00000000..2a8bc258 Binary files /dev/null and b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/_static/image/logo.png differ diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/attention_backend.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/attention_backend.md new file mode 100644 index 00000000..2deb9c6e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/attention_backend.md @@ -0,0 +1,269 @@ +# 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. +``` + +## 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](https://github.com/sgl-project/sglang/blob/main/docs/basic_usage/deepseek.md#multi-head-latent-attention-mla) and the original [DeepSeek MLA paper](https://arxiv.org/pdf/2405.04434). + +### MHA Backends + +| **Backend** | **Page Size > 1 (native)** | **FP8 KV Cache** | **Spec topk=1** | **Spec topk>1** | **Sliding Window** | **MultiModal** | +|---------------------------------|-----------------------------|------------------|-----------------|-----------------|--------------------|----------------| +| **FlashInfer** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| **FA3 (FlashAttention 3)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| **FA4 (FlashAttention 4)** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **Triton** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | +| **Torch Native (SDPA)** | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | +| **FlexAttention (PyTorch)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **TRTLLM MHA** | 16, 32 or 64 | ✅ | ✅ | ❌ | ✅ | ❌ | +| **Dual Chunk FlashAttention** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **AITER (ROCm)** | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | +| **Wave (ROCm)** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **Ascend (NPU)** | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | +| **Intel XPU** | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | + +### MLA Backends + +| **Backend** | **Native Page Sizes** | **FP8 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** | 128 | ❌ | ❌ | ❌ | ❌ | +| **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. +``` + +```{warning} +FlashMLA FP8 KV cache is currently not working. See upstream issue [#8856](https://github.com/sgl-project/sglang/pull/8856). Use non-FP8 KV or another backend when FP8 KV cache is required. +``` + +```{note} +- FlashAttention 4 is prefill-only for now. +- NSA is specifically designed for [DeepSeek V3.2 DSA](https://lmsys.org/blog/2025-09-29-deepseek-V32/). +``` + +```{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: 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), FA4 (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}. +- FA4: page_size = 128. + +### 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. +- `flex_attention` is not supported with speculative decoding. +- For MLA backends, `trtllm_mla` supports `topk > 1`; `flashmla` and `flashinfer_mla` support only `topk = 1`. +- 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. +``` + +## 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 +``` + +- Torch Native +```bash +python3 -m sglang.launch_server \ + --model meta-llama/Meta-Llama-3.1-8B-Instruct \ + --attention-backend torch_native +``` + +- 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 +``` + +- 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 +``` + +- Cutlass MLA +```bash +python3 -m sglang.launch_server \ + --tp 8 \ + --model deepseek-ai/DeepSeek-R1 \ + --attention-backend cutlass_mla \ + --trust-remote-code +``` + +- FlashAttention 4 (MHA & MLA) +```bash +python3 -m sglang.launch_server \ + --tp 8 \ + --model deepseek-ai/DeepSeek-R1 \ + --prefill-attention-backend fa4 \ + --trust-remote-code +``` + +## 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. + +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/checkpoint_engine.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/checkpoint_engine.md new file mode 100644 index 00000000..5e39a7ee --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/checkpoint_engine.md @@ -0,0 +1,254 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/deterministic_inference.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/deterministic_inference.md new file mode 100644 index 00000000..b5b6b521 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/deterministic_inference.md @@ -0,0 +1,154 @@ +# 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). diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/dp_for_multi_modal_encoder.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/dp_for_multi_modal_encoder.md new file mode 100644 index 00000000..ad6132c6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/dp_for_multi_modal_encoder.md @@ -0,0 +1,69 @@ +# 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 deocoder 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. + +## Pros and Cons for DP Multi-Modal Encoder + +- Unfavorable compute/communication ratio for small ViTs +ViTs used in multimodal stacks are typically modest in parameter count and activation sizes. TP introduces per-layer all-reduce collectives (attention/MLP) whose latency and synchronization overhead outweigh the speedup of splitting relatively small GEMMs. With DP, each GPU runs a full ViT locally—no inference-time collectives—so latency is dominated by compute, not wire time. + +- Graph-capture gaps amplify TP overhead +In production, the vision path often has dynamic shapes (pre/post-processing, variable resolution, patching) that break CUDA Graphs and limit torch.compile fusion. Without capture, we need to pay extra kernel-launch and framework overhead; TP then multiplies that cost with additional NCCL synchronizations. Keeping ViT in DP avoids layering collective latency on top of non-captured kernels. + +- Better interconnect hygiene for the true bottleneck (the LLM) +The LLM’s prefill and decode phases benefit materially from TP on fast links. Offloading ViT to DP eliminates “chatty” small collectives on the same fabric, reducing congestion and jitter for the LLM’s large, bandwidth-hungry all-reduces. + +- Shorter and steadier critical path → lower TTFT +TTFT ≈ T(image encode via ViT) + T(LLM prefill) + T(softmax/sample) +DP has several advantages: +(a) batch and prefetch ViT encodes independently, +(b) overlap them with other requests’ LLM decodes on separate streams, +(c) hand off compact visual embeddings to the TP LLM with minimal queuing. + +- For vision encoders that use hardware-unoptimized Conv3D operations, +batch-level DP can provide another 40% improvement compared to regular TP. + +- Nevertheless, since the weights of the multi-modal encoder are replicated across each TP rank, +there will be a minor increase in memory consumption and may cause OOM if you can barely fit the model already. + +## Command Example +You can enable batch-level DP by setting `mm-enable-dp-encoder`, for example: +``` +SGLANG_MM_FEATURE_CACHE_MB=4096 \ +SGLANG_USE_CUDA_IPC_TRANSPORT=1 \ +SGLANG_VLM_CACHE_SIZE_MB=512 \ +python3 -m sglang.launch_server --host 127.0.0.1 \ + --mem-fraction-static 0.7 \ + --port 30000 \ + --trust-remote-code \ + --dtype auto \ + --max-running-requests 4 \ + --chunked-prefill-size 8192 \ + --attention-backend flashinfer \ + --tp 4 \ + --enable-multimodal \ + --chat-template internvl-2-5 \ + --model OpenGVLab/InternVL2_5-8B \ + --disable-radix-cache \ + --mm-enable-dp-encoder +``` +!!! important + Batch-level multi-modal DP is not to be confused with API request-level DP + (which is instead controlled by `data_parallel_size`). + +## Known supported models +- Qwen2.5-VL () +- Qwen3-VL () +- InternVL () diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/expert_parallelism.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/expert_parallelism.md new file mode 100644 index 00000000..3d9075ff --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/expert_parallelism.md @@ -0,0 +1,141 @@ +# 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. | + +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). 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 and Mooncake 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. | SM100+ with TRT-LLM. | +| `flashinfer_cutlass` | FlashInfer combined with CUTLASS for high-performance grouped GEMMs in MoE layers, handling FP4/FP8 quantization efficiently. | SM100+ 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 true --load-balance-method 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). diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/forward_hooks.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/forward_hooks.md new file mode 100644 index 00000000..4dc170b3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/forward_hooks.md @@ -0,0 +1,297 @@ +## 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`). diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hicache.rst b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hicache.rst new file mode 100644 index 00000000..b2bd08b7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hicache.rst @@ -0,0 +1,8 @@ +Hierarchical KV Caching (HiCache) +================================= + +.. toctree:: + :maxdepth: 1 + + hicache_best_practices.md + hicache_design.md diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hicache_best_practices.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hicache_best_practices.md new file mode 100644 index 00000000..cb1baa01 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hicache_best_practices.md @@ -0,0 +1,196 @@ +# 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.) +``` + +## 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 + +### 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!* diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hicache_design.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hicache_design.md new file mode 100644 index 00000000..b775cef5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hicache_design.md @@ -0,0 +1,155 @@ +# 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`**: 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}' ` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hyperparameter_tuning.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hyperparameter_tuning.md new file mode 100644 index 00000000..d9461e19 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/hyperparameter_tuning.md @@ -0,0 +1,77 @@ +# 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 router](../advanced_features/router.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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/lora.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/lora.ipynb new file mode 100644 index 00000000..da25e988 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/lora.ipynb @@ -0,0 +1,595 @@ +{ + "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", + "* `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", + " \"\"\"\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 1 \\\n", + " --log-level warning \\\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "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", + " \"\"\"\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_lora_4_alpha_16 \\\n", + " --max-loras-per-batch 2 \\\n", + " --log-level warning \\\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "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_lora_4_alpha_16\" # 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", + " \"\"\"\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", + "\n", + "url = f\"http://127.0.0.1:{port}\"\n", + "wait_for_server(url)" + ] + }, + { + "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": "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": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "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", + " \"\"\"\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_lora_4_alpha_16\",\"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", + "\n", + "url = f\"http://127.0.0.1:{port}\"\n", + "wait_for_server(url)" + ] + }, + { + "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`: Default basic Triton-based backend.\n", + "- `csgmv`: 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.\n", + "Currently it is at preview phase, we expect to make it our the default LoRA backend in future release. Before that, you can adopt it by manually setting the `--lora-backend` server config." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\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", + " \"\"\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/observability.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/observability.md new file mode 100644 index 00000000..9c5d2e17 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/observability.md @@ -0,0 +1,35 @@ +# 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`. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/pd_disaggregation.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/pd_disaggregation.md new file mode 100644 index 00000000..ec437ecb --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/pd_disaggregation.md @@ -0,0 +1,325 @@ +# 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 Router documentation](router.md#mode-3-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=True +export MC_FORCE_MNNVL=True +``` + +#### 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` | + +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 +``` + +## ASCEND + +### Usage + +Use ascend backend with [mf_adapter(download link)](https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com:443/sglang/mf_adapter-1.0.0-cp311-cp311-linux_aarch64.whl?AccessKeyId=HPUAXT4YM0U8JNTERLST&Expires=1783151861&Signature=3j10QDUjqk70enaq8lostYV2bEA%3D) and ASCEND_MF_STORE_URL being set + +```bash +pip install mf_adapter-1.0.0-cp311-cp311-linux_aarch64.whl --force-reinstall +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 +``` + + +### 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/quantization.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/quantization.md new file mode 100644 index 00000000..18ef0e8a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/quantization.md @@ -0,0 +1,384 @@ +# 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), [ModelCloud](https://huggingface.co/collections/ModelCloud/vortex-673743382af0a52b2a8b9fe2) +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. + +## 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/TensorRT-Model-Optimizer) + +NVIDIA Model Optimizer (ModelOpt) provides advanced quantization techniques optimized for NVIDIA hardware. SGLang includes a streamlined workflow for quantizing models with ModelOpt and automatically exporting them for deployment. + +##### Installation + +First, install ModelOpt. You can either install it directly or as an optional SGLang dependency: + +```bash +# Option 1: Install ModelOpt directly +pip install nvidia-modelopt + +# Option 2: Install SGLang with ModelOpt support (recommended) +pip install sglang[modelopt] +``` + +##### Quantization and Export Workflow + +SGLang provides an example script that demonstrates the complete ModelOpt quantization and export workflow: + +```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 +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: + +```python +import sglang as sgl + +# Deploy exported ModelOpt quantized model +llm = sgl.Engine( + model_path="./quantized_tinyllama_fp8", + 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.outputs[0].text}") +``` + +##### 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: + +```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 +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"]`. + +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 +``` + +## Reference + +- [GPTQModel](https://github.com/ModelCloud/GPTQModel) +- [LLM Compressor](https://github.com/vllm-project/llm-compressor/) +- [NVIDIA Model Optimizer (ModelOpt)](https://github.com/NVIDIA/TensorRT-Model-Optimizer) +- [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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/router.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/router.md new file mode 100644 index 00000000..388b86cd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/router.md @@ -0,0 +1,470 @@ +# SGLang Model Gateway (formerly SGLang Router) + +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 router 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 & Privacy](#storage--privacy) +3. [Deployment Modes](#deployment-modes) + - [Co-launch Router + Workers](#co-launch-router--workers) + - [Separate Launch (HTTP)](#separate-launch-http) + - [gRPC Launch](#grpc-launch) + - [Prefill/Decode Disaggregation](#prefilldecode-disaggregation) + - [OpenAI Backend Proxy](#openai-backend-proxy) +4. [Worker Lifecycle & Dynamic Scaling](#worker-lifecycle--dynamic-scaling) +5. [Reliability & Flow Control](#reliability--flow-control) +6. [Load Balancing Policies](#load-balancing-policies) +7. [Service Discovery (Kubernetes)](#service-discovery-kubernetes) +8. [Security & Authentication](#security--authentication) +9. [History & Data Connectors](#history--data-connectors) +10. [MCP & Advanced Tooling](#mcp--advanced-tooling) +11. [API Surface](#api-surface) +12. [Configuration Reference](#configuration-reference) +13. [Observability](#observability) +14. [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). +- **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. +- **Observability**: Prometheus metrics, structured tracing, request ID propagation, and detailed job queue stats. + +--- + +## 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/{url}`) 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. + +### Data Plane +- **HTTP routers** (regular & PD) implement `/generate`, `/v1/chat/completions`, `/v1/completions`, `/v1/responses`, `/v1/embeddings`, `/v1/rerank`, 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. +- **OpenAI router** proxies OpenAI-compatible endpoints to external vendors (OpenAI, xAI, etc.) while keeping chat history and multi-turn orchestration local. + +### Storage & Privacy +- Conversation and response history is stored at the router tier (memory, none, or Oracle ATP). 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. + +--- + +## Deployment Modes + +### Co-launch Router + Workers +Launch the router and a fleet of SGLang workers in one process (ideal for single-node or quick starts). The CLI accepts two namespaces of arguments: +- **Worker arguments** (no prefix) configure the SGLang runtime (`--model`, `--tp-size`, `--dp-size`, `--grpc-mode`, etc.). +- **Router arguments** are prefixed with `--router-` and map directly to `launch_router` flags (`--router-policy`, `--router-model-path`, `--router-log-level`, ...). + +```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: +```bash +python3 -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-health-success-threshold 2 \ + --router-health-check-timeout-secs 6000 \ + --router-health-check-interval-secs 60 \ + --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 +``` + +> gRPC router supports both single-stage and PD serving. Provide `--tokenizer-path` or `--model-path` (HF repo or local directory) plus optional `--chat-template`. + +### 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 \ + --policy cache_aware \ + --prefill-policy cache_aware \ + --decode-policy power_of_two +``` + +### OpenAI Backend Proxy +Proxy OpenAI-compatible endpoints (OpenAI, xAI, etc.) 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. + +--- + +## Worker Lifecycle & Dynamic Scaling + +Add or remove workers at runtime using the REST APIs. Jobs are queued and tracked for eventual consistency. + +```bash +# Add a worker (HTTP or gRPC) +curl -X POST http://localhost:30000/workers \ + -H "Content-Type: application/json" \ + -d '{"url":"grpc://0.0.0.0:31000","worker_type":"regular"}' + +# Inspect registry +curl http://localhost:30000/workers + +# Remove a worker +curl -X DELETE http://localhost:30000/workers/grpc://0.0.0.0:31000 +``` + +Legacy endpoints (`/add_worker`, `/remove_worker`, `/list_workers`) remain available but will be deprecated. `/workers/{url}` returns both registry data and queued job status. + +--- + +## Reliability & Flow Control + +### 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 +``` + +### Circuit Breaker +```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 +``` + +### Rate Limiting & 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 (up to `queue-size`). A `429` is returned when the queue is full; `408` is returned when `queue-timeout-secs` expires. + +--- + +## 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 (requires Load Monitor). | `--policy power_of_two` | +| `cache_aware` | Default policy; combines cache locality with load balancing, falling back to shortest queue. | `--policy cache_aware` + tuning flags | + +Key tuning flags: +```bash +--cache-threshold 0.5 \ +--balance-abs-threshold 32 \ +--balance-rel-threshold 1.5 \ +--eviction-interval-secs 120 \ +--max-tree-size 67108864 +``` + +--- + +## 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 deployments can specify `--prefill-selector` and `--decode-selector` plus the `sglang.ai/bootstrap-port` annotation for prefill bootstrap ports. Ensure RBAC grants `get/list/watch` on pods. + +--- + +## Security & Authentication + +- **Router API key (`--api-key`)**: clients must supply `Authorization: Bearer `. +- **Worker API keys**: when adding workers dynamically, include `api_key` in the payload; workers listed via CLI inherit the router key. +- **Full-stack auth**: start router with `--api-key`, then add workers with their own keys: + ```bash + curl -H "Authorization: Bearer router-key" \ + -X POST http://localhost:30000/workers \ + -H "Content-Type: application/json" \ + -d '{"url":"http://worker:8000","api_key":"worker-key"}' + ``` +- **Privacy**: All conversation history, `/v1/responses` state, and MCP sessions stay inside the router. Nothing is persisted at remote model vendors unless explicitly proxied. + +--- + +## History & Data Connectors + +| Backend | Description | Usage | +|---------|-------------|-------| +| `memory` (default) | In-memory storage for quick prototyping. | `--history-backend memory` | +| `none` | No persistence; APIs operate but store nothing. | `--history-backend none` | +| `oracle` | Oracle Autonomous Database-backed storage (pooled connections). | `--history-backend oracle` | +| `postgres` | PostgreSQL Database-backed storage (pooled connections). | `--history-backend postgres` | + +Oracle configuration (choose DSN *or* TNS alias): +Install the Oracle Instant Client and set `LD_LIBRARY_PATH` accordingly. +Choose **one** connection method: +```bash +# Option 1: Full connection descriptor +export ATP_DSN="(description=(address=(protocol=tcps)(port=1522)(host=adb.region.oraclecloud.com))(connect_data=(service_name=service_name)))" + +# Option 2: TNS alias (requires wallet) +export ATP_TNS_ALIAS="sglroutertestatp_high" +export ATP_WALLET_PATH="/path/to/wallet" +``` +Provide database credentials and optional pool sizing: +```bash +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 +``` + +> History backends currently apply to OpenAI router mode. gRPC parity for `/v1/responses` is on the roadmap. + +--- + +## MCP & Advanced Tooling + +- Native MCP client supports **STDIO**, **HTTP**, **SSE**, and **Streamable** transports—no external config files required. +- Tool-call parsers cover JSON, Pythonic, XML, and custom schemas with streaming/non-streaming execution loops. +- Reasoning parsers ship for DeepSeek-R1, Qwen3, Step-3, GLM4, Llama families, Kimi K2, GPT-OSS, Mistral, and more (`src/reasoning_parser`). +- Tokenizer factory accepts HuggingFace IDs, local directories, and explicit `tokenizer.json` files with chat template overrides (`src/tokenizer`). + +Use CLI flags to select parsers: +```bash +--reasoning-parser deepseek-r1 \ +--tool-call-parser json \ +--chat-template /path/to/template.json +``` + +--- + +## API Surface + +| Method | Path | Description | +|-----------------------|------------------------------------------|------------------------------------------------| +| `POST` | `/generate` | SGLang generate API. | +| `POST` | `/v1/chat/completions` | OpenAI-compatible chat (streaming/tool calls). | +| `POST` | `/v1/completions` | OpenAI-compatible text completions. | +| `POST` | `/v1/responses` | Create background responses (agentic loops). | +| `GET` | `/v1/responses/{id}` | Retrieve stored responses. | +| `POST` | `/v1/embeddings` | Forward embedding requests. | +| `POST` | `/v1/rerank` | Ranking endpoint (`/rerank` synonym). | +| `POST` | `/v1/conversations` | Create conversation metadata. | +| `GET`/`POST`/`DELETE` | `/v1/conversations/{id}` | Get/update/delete conversation. | +| `GET`/`POST` | `/v1/conversations/{id}/items` | List or append conversation items. | +| `GET`/`DELETE` | `/v1/conversations/{id}/items/{item_id}` | Inspect/delete conversation item. | +| `GET` | `/workers` | List registered workers with health/load. | +| `POST` | `/workers` | Queue worker registration. | +| `DELETE` | `/workers/{url}` | Queue worker removal. | +| `POST` | `/flush_cache` | Flush worker caches (HTTP workers). | +| `GET` | `/get_loads` | Retrieve worker load snapshot. | +| `GET` | `/liveness` / `/readiness` / `/health` | Health probes. | + +--- + +## 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 (`random`, `round_robin`, `cache_aware`, `power_of_two`). | +| `--max-concurrent-requests` | int | -1 | Concurrency limit (-1 disables rate limiting). | +| `--request-timeout-secs` | int | 600 | Request timeout. | +| `--max-payload-size` | int | 256MB | Maximum request payload. | + +### Cache-Aware Tuning + +| Parameter | Type | Default | Description | +|----------------------------|-------|----------|-----------------------------| +| `--cache-threshold` | float | 0.3 | Minimum prefix match ratio. | +| `--balance-abs-threshold` | int | 64 | Absolute load threshold. | +| `--balance-rel-threshold` | float | 1.5 | Relative load ratio. | +| `--eviction-interval-secs` | int | 120 | Cache eviction cadence. | +| `--max-tree-size` | int | 67108864 | Max nodes in cache tree. | + +### Fault Tolerance + +| Parameter | Type | Default | Description | +|------------------------------|-------|---------|----------------------------------| +| `--retry-max-retries` | int | 5 | Max retries. | +| `--retry-initial-backoff-ms` | int | 50 | Initial backoff (ms). | +| `--retry-max-backoff-ms` | int | 30000 | Max backoff (ms). | +| `--retry-backoff-multiplier` | float | 1.5 | Backoff multiplier. | +| `--retry-jitter-factor` | float | 0.2 | Retry jitter (0.0-1.0). | +| `--disable-retries` | flag | False | Disable retries. | +| `--cb-failure-threshold` | int | 5 | Failures before opening circuit. | +| `--cb-success-threshold` | int | 2 | Successes to close circuit. | +| `--cb-timeout-duration-secs` | int | 30 | Cooldown period. | +| `--cb-window-duration-secs` | int | 60 | Window size. | +| `--disable-circuit-breaker` | flag | False | Disable circuit breaker. | + +### 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. | +| `--worker-startup-check-interval` | int | 30 | Polling interval. | + +### Kubernetes Discovery + +| Parameter | Type | Description | +|--------------------------------------------|------|--------------------------------------------------------------------| +| `--service-discovery` | flag | Enable discovery. | +| `--selector key=value ...` | list | Label selectors (regular mode). | +| `--prefill-selector` / `--decode-selector` | list | Label selectors for PD mode. | +| `--service-discovery-namespace` | str | Namespace to watch. | +| `--service-discovery-port` | int | Worker port (default 80). | +| `--bootstrap-port-annotation` | str | Prefill bootstrap annotation (default `sglang.ai/bootstrap-port`). | + +--- + +## Observability + +Enable Prometheus metrics: +```bash +python -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 http://worker2:8001 \ + --prometheus-host 0.0.0.0 \ + --prometheus-port 29000 +``` + +Key metrics: + +| Metric | Type | Description | +|--------|------|-------------| +| `sgl_router_requests_total` | Counter | Total requests by endpoint/method. | +| `sgl_router_processed_requests_total` | Counter | Requests processed per worker. | +| `sgl_router_active_workers` | Gauge | Healthy worker count. | +| `sgl_router_running_requests` | Gauge | In-flight requests per worker. | +| `sgl_router_cache_hits_total` / `misses_total` | Counter | Cache-aware routing hits/misses. | +| `sgl_router_generate_duration_seconds` | Histogram | Request latency distribution. | + +Enable request ID propagation: +```bash +python -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 \ + --request-id-headers x-request-id x-trace-id +``` + +--- + +## Troubleshooting + +1. **Workers never ready** + Increase `--worker-startup-timeout-secs` or ensure health probes respond before router startup. + +2. **Load imbalance / hot workers** + Inspect `sgl_router_processed_requests_total` and tune cache-aware thresholds (`--balance-*`, `--cache-threshold`). + +3. **Circuit breaker flapping** + Increase `--cb-failure-threshold` or extend the timeout/window durations. Consider temporarily disabling retries. + +4. **Queue overflow (429)** + Increase `--queue-size` or reduce client concurrency. Ensure `--max-concurrent-requests` matches downstream capacity. + +5. **Memory growth** + Reduce `--max-tree-size` or lower `--eviction-interval-secs` for more aggressive cache pruning. + +6. **Debugging** + ```bash + python -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 \ + --log-level debug \ + --log-dir ./router_logs + ``` + +--- + +SGLang Model Gateway continues to evolve alongside the SGLang runtime. Keep CLI flags, integrations, and documentation aligned when adopting new features or contributing improvements. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/separate_reasoning.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/separate_reasoning.ipynb new file mode 100644 index 00000000..fa24e63b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/separate_reasoning.ipynb @@ -0,0 +1,377 @@ +{ + "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 models](https://huggingface.co/moonshotai/models) | `◁think▷` … `◁/think▷` | `kimi` | Uses special thinking delimiters |\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:**\n", + "- Kimi: Uses special `◁think▷` and `◁/think▷` tags\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}\")" + ] + }, + { + "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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/server_arguments.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/server_arguments.md new file mode 100644 index 00000000..891cbacf --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/server_arguments.md @@ -0,0 +1,467 @@ +# 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 Router](../advanced_features/router.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 `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). +- To enable torchao quantization, add `--torchao-config int4wo-128`. It supports other [quantization strategies (INT8/FP8)](https://github.com/sgl-project/sglang/blob/v0.3.6/python/sglang/srt/server_args.py#L671) as well. +- 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_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). +- 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` + + ```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. | `auto` | `auto`, `pt`, `safetensors`, `npcache`, `dummy`, `sharded_state`, `gguf`, `bitsandbytes`, `layered`, `remote`, `remote_instance` | +| `--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 | +| `--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 | + +## 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_fp4`, `petit_nvfp4`, `w8a8_int8`, `w8a8_fp8`, `moe_wna16`, `qoq`, `w4afp8`, `mxfp4` | +| `--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] | +| `--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 | +| `--kv-cache-dtype` | Data type for kv cache storage. "auto" will use model data type. "fp8_e5m2" and "fp8_e4m3" is supported for CUDA 11.8+. | `auto` | `auto`, `fp8_e5m2`, `fp8_e4m3` | +| `--enable-fp32-lm-head` | If set, the LM head outputs (logits) are in FP32. | `False` | bool flag (set to enable) | + +## 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 | +| `--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` | +| `--enable-priority-scheduling` | Enable priority scheduling. Requests with higher priority integer values will be scheduled first by default. | `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 | +| `--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] | +| `--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) | + +## 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 | +| `--elastic-ep-backend` | Select the collective communication backend for elastic EP. Currently supports 'mooncake'. | None | N/A | +| `--mooncake-ib-device` | The InfiniBand devices for Mooncake Backend, accepts multiple comma-separated devices. Default is None, which triggers automatic device detection when Mooncake Backend is enabled. | None | N/A | +| `--tensor-parallel-size`
`--tp-size` | The tensor parallelism size. | `1` | Type: int | +| `--pipeline-parallel-size`
`--pp-size` | The pipeline parallelism size. | `1` | Type: int | +| `--pp-max-micro-batch-size` | The maximum micro batch size in pipeline parallelism. | `None` | 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 | +| `--stream-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 | +| `--dist-timeout` | Set timeout for torch.distributed initialization. | `None` | Type: int | +| `--download-dir` | Model download directory for huggingface. | `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) | +| `--mm-process-config` | A JSON string for multimodal preprocessing configuration. It can contain keys: `image`, `video`, `audio`. | `{}` | + +## 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` | +| `--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 | +| `--crash-on-nan` | Crash the server on nan logprobs. | `False` | 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) | +| `--oltp-traces-endpoint` | Config opentelemetry collector endpoint if --enable-trace is set. format: : | `localhost:4317` | 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 | +| `--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 buliltin chat template name or the path of the chat template file. This is only used for OpenAI-compatible API server. | `None` | Type: str | +| `--completion-template` | The buliltin 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, gpt-oss, kimi_k2, llama3, mistral, pythonic, qwen, qwen25, qwen3_coder, step3]. | `None` | `deepseekv3`, `deepseekv31`, `glm`, `glm45`, `gpt-oss`, `kimi_k2`, `llama3`, `mistral`, `pythonic`, `qwen`, `qwen25`, `qwen3_coder`, `step3` | +| `--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` | +| `--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 | + +## 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 Minimum Token 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. | `round_robin` | `round_robin`, `shortest_queue`, `minimum_tokens` | +| `--load-watch-interval` | The interval of load watching in seconds. | `0.1` | Type: float | +| `--prefill-round-robin-balance` | Prefill is round robin balanced. This is used to promise decode server can get the correct dp rank. | `False` | bool flag (set to enable) | + +## 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) | +| `--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. | `triton` | `triton`, `csgmv` | +| `--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 backend +| 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`, `triton_attn`, `ascend_attn`, `aiter_attn` | +| `--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` | + +## 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-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-moe-runner-backend` | MOE backend for EAGLE speculative decoding, see --moe-runner-backend for options. Same as moe runner backend if unset. | None | + +## 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 | + +## 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` | +| `--moe-runner-backend` | Choose the runner backend for MoE. | `auto` | `auto`, `deep_gemm`, `triton`, `triton_kernel`, `flashinfer_trtllm`, `flashinfer_cutlass`, `flashinfer_mxfp4`, `flashinfer_cutedsl` | +| `--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) | +| `--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 | + +## 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` | +| `--mamba-full-memory-ratio` | The ratio of mamba state memory to full kv cache memory. | `0.2` | Type: float | + +## 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 | + +## 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` | +| `--radix-eviction-policy` | The eviction policy of radix trees. 'lru' stands for Least Recently Used, 'lfu' stands for Least Frequently Used. | `lru` | `lru`, `lfu` | +| `--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` | +| `--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 containing extra configuration for the storage backend. | `None` | Type: str | + +## LMCache +| Argument | Description | Defaults | Options | +| --- | --- | --- | --- | +| `--enable-lmcache` | Using LMCache as an alternative hierarchical cache solution | `False` | bool flag (set to enable) | + +## 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 | + +## 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-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-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) | +| `--enable-piecewise-cuda-graph` | Optimize the model with piecewise cuda graph for extend/prefill only. Experimental feature. | `False` | bool flag (set to enable) | +| `--piecewise-cuda-graph-tokens` | Set the list of tokens when using piecewise cuda graph. | `None` | Type: JSON list | +| `--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) | +| `--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) | +| `--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-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-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` | Context parallelism used in the long sequence prefill phase of DeepSeek v3.2 | `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 | + +## Debug tensor dumps +| Argument | Description | Defaults | Options | +| --- | --- | --- | --- | +| `--debug-tensor-dump-output-folder` | The output folder for dumping tensors. | `None` | Type: str | +| `--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 | +| `--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 | + +## 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-decode-tp` | Decode tp size. If not set, it matches the tp size of the current engine. This is only set on the prefill server. | `None` | Type: int | +| `--disaggregation-decode-dp` | Decode dp size. If not set, it matches the dp size of the current engine. This is only set on the prefill server. | `None` | Type: int | +| `--disaggregation-prefill-pp` | Prefill pp size. If not set, it is default to 1. This is only set on the decode server. | `1` | 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 | + +## 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 | + +## 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 | + +## For deterministic inference +| Argument | Description | Defaults | Options | +| --- | --- | --- | --- | +| `--enable-deterministic-inference` | Enable deterministic inference mode with batch invariant ops. | `False` | bool flag (set to enable) | + +## 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 | +| `--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 | + +## 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 | diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/speculative_decoding.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/speculative_decoding.ipynb new file mode 100644 index 00000000..aa62b897 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/speculative_decoding.ipynb @@ -0,0 +1,370 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Speculative Decoding\n", + "\n", + "SGLang now provides an EAGLE-based (EAGLE-2/EAGLE-3) speculative decoding option. Our implementation aims to maximize speed and efficiency and is considered to be among the fastest in open-source LLM engines.\n", + "\n", + "### Performance Highlights\n", + "\n", + "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.\n", + "For further details please see the [EAGLE3 paper](https://arxiv.org/pdf/2503.01840).\n", + "\n", + "| Method | Throughput (tokens/s) |\n", + "|--------|----------------|\n", + "| SGLang (w/o speculative, 1x H100) | 158.34 tokens/s |\n", + "| SGLang + EAGLE-2 (1x H100) | 244.10 tokens/s |\n", + "| SGLang + EAGLE-3 (1x H100) | 373.25 tokens/s |" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## EAGLE Decoding\n", + "\n", + "To enable EAGLE speculative decoding the following parameters are relevant:\n", + "* `speculative_draft_model_path`: Specifies draft model. This parameter is required.\n", + "* `speculative_num_steps`: Depth of autoregressive drafting. Increases speculation range but risks rejection cascades. Default is 5.\n", + "* `speculative_eagle_topk`: Branching factor per step. Improves candidate diversity, will lead to higher acceptance rate, but more lead to higher memory/compute consumption. Default is 4.\n", + "* `speculative_num_draft_tokens`: Maximum parallel verification capacity. Allows deeper tree evaluation but will lead to higher GPU memory usage. Default is 8.\n", + "\n", + "These parameters are the same for EAGLE-2 and EAGLE-3.\n", + "\n", + "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).\n", + "\n", + "In the documentation below, we set `--cuda-graph-max-bs` to be a small value for faster engine startup. For your own workloads, please tune the above parameters together with `--cuda-graph-max-bs`, `--max-running-requests`, `--mem-fraction-static` for the best performance. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EAGLE-2 decoding\n", + "\n", + "You can enable EAGLE-2 decoding by setting `--speculative-algorithm EAGLE` and choosing an appropriate model." + ] + }, + { + "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", + "import openai" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Llama-2-7b-chat-hf --speculative-algorithm EAGLE \\\n", + " --speculative-draft-model-path lmsys/sglang-EAGLE-llama2-chat-7B --speculative-num-steps 3 \\\n", + " --speculative-eagle-topk 4 --speculative-num-draft-tokens 16 --cuda-graph-max-bs 8 --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "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=\"meta-llama/Llama-2-7b-chat-hf\",\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": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EAGLE-2 Decoding with `torch.compile`\n", + "\n", + "You can also enable `torch.compile` for further optimizations and optionally set `--torch-compile-max-bs`:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Llama-2-7b-chat-hf --speculative-algorithm EAGLE \\\n", + " --speculative-draft-model-path lmsys/sglang-EAGLE-llama2-chat-7B --speculative-num-steps 5 \\\n", + " --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --mem-fraction 0.6 \\\n", + " --enable-torch-compile --torch-compile-max-bs 2 --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "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=\"meta-llama/Llama-2-7b-chat-hf\",\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": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EAGLE-2 Decoding via Frequency-Ranked Speculative Sampling\n", + "\n", + "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, checkout [the paper](https://arxiv.org/pdf/arXiv:2502.14856).\n", + "\n", + "In our implementation, set `--speculative-token-map` to enable the optimization. You can get the high-frequency token in FR-Spec from [this model](https://huggingface.co/thunlp/LLaMA3-Instruct-8B-FR-Spec). Or you can obtain high-frequency token by directly downloading these token from [this repo](https://github.com/thunlp/FR-Spec/tree/main?tab=readme-ov-file#prepare-fr-spec-vocabulary-subset).\n", + "\n", + "Thanks for the contribution from [Weilin Zhao](https://github.com/Achazwl) and [Zhousx](https://github.com/Zhou-sx). " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3-8B-Instruct --speculative-algorithm EAGLE \\\n", + " --speculative-draft-model-path lmsys/sglang-EAGLE-LLaMA3-Instruct-8B --speculative-num-steps 5 \\\n", + " --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --speculative-token-map thunlp/LLaMA3-Instruct-8B-FR-Spec/freq_32768.pt \\\n", + " --mem-fraction 0.7 --cuda-graph-max-bs 2 --dtype float16 --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "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=\"meta-llama/Meta-Llama-3-8B-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": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EAGLE-3 Decoding\n", + "\n", + "You can enable EAGLE-3 decoding by setting `--speculative-algorithm EAGLE3` and choosing an appropriate model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Llama-3.1-8B-Instruct --speculative-algorithm EAGLE3 \\\n", + " --speculative-draft-model-path jamesliu1/sglang-EAGLE3-Llama-3.1-Instruct-8B --speculative-num-steps 5 \\\n", + " --speculative-eagle-topk 8 --speculative-num-draft-tokens 32 --mem-fraction 0.6 \\\n", + " --cuda-graph-max-bs 2 --dtype float16 --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "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=\"meta-llama/Meta-Llama-3.1-8B-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": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multi Token Prediction\n", + "\n", + "We support [MTP(Multi-Token Prediction)](https://arxiv.org/pdf/2404.19737) in SGLang by using speculative decoding. We use Xiaomi/MiMo-7B-RL model as example here (deepseek mtp usage refer to [deepseek doc](../basic_usage/deepseek.md#multi-token-prediction))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + " python3 -m sglang.launch_server --model-path XiaomiMiMo/MiMo-7B-RL --host 0.0.0.0 --trust-remote-code \\\n", + " --speculative-algorithm EAGLE --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \\\n", + " --mem-fraction 0.5 --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "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\": \"XiaomiMiMo/MiMo-7B-RL\",\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": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## References\n", + "\n", + "EAGLE process is as follows:\n", + "\n", + "- 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})$. \n", + "- 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.\n", + "- EAGLE-2 additionally uses the draft model to evaluate how probable certain branches in the draft tree are, dynamically stopping the expansion of unlikely branches. After the expansion phase, reranking is employed to select only the top `speculative_num_draft_tokens` final nodes as draft tokens.\n", + "- EAGLE-3 removes the feature prediction objective, incorporates low and mid-layer features, and is trained in an on-policy manner.\n", + "\n", + "This enhances drafting accuracy by operating on the features instead of tokens for more regular inputs and passing the tokens from the next timestep additionally to minimize randomness effects from sampling. Furthermore the dynamic adjustment of the draft tree and selection of reranked final nodes increases acceptance rate of draft tokens further. For more details see [EAGLE-2](https://arxiv.org/abs/2406.16858) and [EAGLE-3](https://arxiv.org/abs/2503.01840) paper.\n", + "\n", + "\n", + "For guidance how to train your own EAGLE model please see the [EAGLE repo](https://github.com/SafeAILab/EAGLE/tree/main?tab=readme-ov-file#train)." + ] + } + ], + "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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/structured_outputs.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/structured_outputs.ipynb new file mode 100644 index 00000000..b0ec5e6c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/structured_outputs.ipynb @@ -0,0 +1,998 @@ +{ + "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}\")\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", + "\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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/structured_outputs_for_reasoning_models.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/structured_outputs_for_reasoning_models.ipynb new file mode 100644 index 00000000..9cdcc29e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/structured_outputs_for_reasoning_models.ipynb @@ -0,0 +1,830 @@ +{ + "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}\")\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\": [\"\")[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", + " \"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", + " \"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", + " \"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", + " \"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(f\"http://localhost:{port_tool_choice}\")\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}\")\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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/vlm_query.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/vlm_query.ipynb new file mode 100644 index 00000000..c753f2fd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/advanced_features/vlm_query.ipynb @@ -0,0 +1,325 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Query Vision Language Model" + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## Querying Qwen-VL" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import nest_asyncio\n", + "\n", + "nest_asyncio.apply() # Run this first.\n", + "\n", + "model_path = \"Qwen/Qwen2.5-VL-3B-Instruct\"\n", + "chat_template = \"qwen2-vl\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "# Lets create a prompt.\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", + "image = Image.open(\n", + " BytesIO(\n", + " requests.get(\n", + " \"https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true\"\n", + " ).content\n", + " )\n", + ")\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(conv.get_prompt())\n", + "image" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, + "source": [ + "### Query via the offline Engine API" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "from sglang import Engine\n", + "\n", + "llm = Engine(\n", + " model_path=model_path, chat_template=chat_template, mem_fraction_static=0.8\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "out = llm.generate(prompt=conv.get_prompt(), image_data=[image])\n", + "print(out[\"text\"])" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "### Query via the offline Engine API, but send precomputed embeddings" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "# Compute the image embeddings using Huggingface.\n", + "\n", + "from transformers import AutoProcessor\n", + "from transformers import Qwen2_5_VLForConditionalGeneration\n", + "\n", + "processor = AutoProcessor.from_pretrained(model_path, use_fast=True)\n", + "vision = (\n", + " Qwen2_5_VLForConditionalGeneration.from_pretrained(model_path).eval().visual.cuda()\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "processed_prompt = processor(\n", + " images=[image], text=conv.get_prompt(), return_tensors=\"pt\"\n", + ")\n", + "input_ids = processed_prompt[\"input_ids\"][0].detach().cpu().tolist()\n", + "precomputed_embeddings = vision(\n", + " processed_prompt[\"pixel_values\"].cuda(), processed_prompt[\"image_grid_thw\"].cuda()\n", + ")\n", + "\n", + "mm_item = dict(\n", + " modality=\"IMAGE\",\n", + " image_grid_thw=processed_prompt[\"image_grid_thw\"],\n", + " precomputed_embeddings=precomputed_embeddings,\n", + ")\n", + "out = llm.generate(input_ids=input_ids, image_data=[mm_item])\n", + "print(out[\"text\"])" + ] + }, + { + "cell_type": "markdown", + "id": "10", + "metadata": {}, + "source": [ + "## Querying Llama 4 (Vision)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "import nest_asyncio\n", + "\n", + "nest_asyncio.apply() # Run this first.\n", + "\n", + "model_path = \"meta-llama/Llama-4-Scout-17B-16E-Instruct\"\n", + "chat_template = \"llama-4\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "# Lets create a prompt.\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", + "image = Image.open(\n", + " BytesIO(\n", + " requests.get(\n", + " \"https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true\"\n", + " ).content\n", + " )\n", + ")\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(conv.get_prompt())\n", + "print(f\"Image size: {image.size}\")\n", + "\n", + "image" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "### Query via the offline Engine API" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.test_utils import is_in_ci\n", + "\n", + "if not is_in_ci():\n", + " from sglang import Engine\n", + "\n", + " llm = Engine(\n", + " model_path=model_path,\n", + " trust_remote_code=True,\n", + " enable_multimodal=True,\n", + " mem_fraction_static=0.8,\n", + " tp_size=4,\n", + " attention_backend=\"fa3\",\n", + " context_length=65536,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "if not is_in_ci():\n", + " out = llm.generate(prompt=conv.get_prompt(), image_data=[image])\n", + " print(out[\"text\"])" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "### Query via the offline Engine API, but send precomputed embeddings" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "if not is_in_ci():\n", + " # Compute the image embeddings using Huggingface.\n", + "\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", + " vision = model.vision_model.cuda()\n", + " multi_modal_projector = model.multi_modal_projector.cuda()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "if not is_in_ci():\n", + " processed_prompt = processor(\n", + " images=[image], text=conv.get_prompt(), return_tensors=\"pt\"\n", + " )\n", + " print(f'{processed_prompt[\"pixel_values\"].shape=}')\n", + " input_ids = processed_prompt[\"input_ids\"][0].detach().cpu().tolist()\n", + "\n", + " image_outputs = vision(\n", + " processed_prompt[\"pixel_values\"].to(\"cuda\"), output_hidden_states=False\n", + " )\n", + " image_features = image_outputs.last_hidden_state\n", + " vision_flat = image_features.view(-1, image_features.size(-1))\n", + " precomputed_embeddings = multi_modal_projector(vision_flat)\n", + "\n", + " mm_item = dict(modality=\"IMAGE\", precomputed_embeddings=precomputed_embeddings)\n", + " out = llm.generate(input_ids=input_ids, image_data=[mm_item])\n", + " print(out[\"text\"])" + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all", + "custom_cell_magics": "kql", + "encoding": "# -*- coding: utf-8 -*-" + }, + "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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/deepseek_v3.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/deepseek_v3.md new file mode 100644 index 00000000..285a318f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/deepseek_v3.md @@ -0,0 +1,280 @@ +# 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**
*(recommended)* | 8 x H200 | +| | 8 x MI300X | +| | 2 x 8 x H100/800/20 | +| | Xeon 6980P CPU | +| **Full precision BF16** | 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 (AWQ)** | 8 x H100/800/20 | +| | 8 x A100/A800 | +| **Quantized weights (int8)** | 16 x A100/800 | +| | 32 x L40S | +| | Xeon 6980P CPU | +| | 2 x Atlas 800I A3 | + + + +Detailed commands for reference: + +- [8 x H200](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#using-docker-recommended) +- [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-h208-nodes) +- [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) +- [2 x Atlas 800I A3 (int8)](../platforms/ascend_npu.md#running-deepseek-v3) + +### 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). +**Note that Deepseek V3 is already in FP8**, so we should not run it with any quantization arguments like `--quantization fp8 --kv-cache-dtype fp8_e5m2`. + +### Running examples on Multi-node + +- [Serving with two H20*8 nodes](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-two-h208-nodes). + +- [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. For MLA models on Blackwell architecture (e.g., B200), the default backend is FlashInfer. To use the optimized TRTLLM MLA backend for prefill and decode operations, explicitly specify `--attention-backend trtllm_mla`. + +**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. However, it is not recommended for low-latency, small-batch use cases. +- 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. + +**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-2-h208) 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 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`. + +Before serving the DeepSeek model, precompile the DeepGEMM kernels using: +```bash +python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code +``` +The precompilation process typically takes around 10 minutes to complete. + +### 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 arguments `--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 deepseek-ai/DeepSeek-V3-0324 \ + --speculative-algorithm EAGLE \ + --speculative-num-steps 1 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 2 \ + --trust-remote-code \ + --tp 8 +``` +- 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. +- FlashAttention3, FlashMLA, and Triton backend fully supports MTP usage. For FlashInfer backend (`--attention-backend flashinfer`) with speculative decoding,`--speculative-eagle-topk` parameter should be set to `1`. MTP support for the CutlassMLA and TRTLLM MLA backends are still under development. +- To enable DeepSeek MTP for large batch sizes (>32), there are some parameters should be changed (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 is set [here](https://github.com/sgl-project/sglang/blob/49420741746c8f3e80e0eb17e7d012bfaf25793a/python/sglang/srt/model_executor/cuda_graph_runner.py#L126). You can include more batch sizes into it. + + +### 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 an 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": "Hows 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 an 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": "Hows 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 Notes: +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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/deepseek_v32.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/deepseek_v32.md new file mode 100644 index 00000000..229c7c5a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/deepseek_v32.md @@ -0,0 +1,234 @@ +# DeepSeek V3.2 Usage + +[DeepSeek-V3.2-Exp](https://huggingface.co/deepseek-ai/DeepSeek-V3.2-Exp) 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). + +## Installation + +### Docker + +```bash +# H200/B200 +docker pull lmsysorg/sglang:latest + +# MI350/MI355 +docker pull lmsysorg/sglang:dsv32-rocm + +# 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 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 +``` + +### 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. + - `alter`: Alter kernel on AMD HPUs. Can only be used as decode kernel. +- 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. + +## 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: +```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 +``` +- 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. + + +## 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. + +## 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 120000 --repeat 8 --thinking-mode deepseek-v3 +``` + +The mean accuracy over 8 runs shows 0.797, which matches the number 79.9 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'] +``` + +### 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 +``` + +Modify the [`jinja chat_template`](https://huggingface.co/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/tokenizer_config.json#L34) by replacing + +``` +{% set thinking = false %} +``` +with +``` +{% set thinking = true %} +``` +and save it to `chat_template_thinking.jinja`. + +Launch the SGLang server with the modified chat-template file: +``` +python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --dp 8 --enable-dp-attention --chat-template chat_template_thinking.jinja +``` + +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" +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) \ + ++max_concurrent_requests=512 \ + ++server.api_key=dummy \ + ++inference.tokens_to_generate=64000 +``` + +Test results: + + +| evaluation_mode | num_entries | avg_tokens | gen_seconds | symbolic_correct | no_answer | +|--------------------|-------------|------------|-------------|-----------------------|-----------| +| pass@1[avg-of-4] | 30 | 14410 | 1758 | 85.83% ± 4.19% | 0.00% | +| majority@4 | 30 | 14410 | 1758 | 90.00% | 0.00% | +| pass@4 | 30 | 14410 | 1758 | 93.33% | 0.00% | + +Note that the result of problem#3 with id `aime25-2` is marked as false by nemo-skills but is actually correct because nemo-skills fails to match predicted_answer `016` with expected_answer `16`. If we add 1/30 = 3.33% to the results, the pass@1[avg-of-4] result matches with reference which is 89.3. + + +## DSA long sequence context parallel optimization(experimental) + +Accuracy benchmark on long context can be tested on GPQA-diamond dataset with long output tokens and thinking enabled: + +Example usage: +```bash +# Launch 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 --max-running-requests 32 +``` +### Context-parallel Tips +`CP_size` reuses `atten_tp_size`, which is equal to `TP_size` / `DP_size`. +Some features are still not supported at present. +- **Multi-batch prefill**: Currently, only single-request processing is supported during the prefill process. +- **disaggregation**: P/D disaggregation. +- **Cross-machine support**: - Currently only tested on a single machine (TP=8,EP=8). +- **Other Args**: Currently only supports moe_dense_tp_size=1, kv_cache_dtype = "bf16", moe_a2a_backend = "deepep", +- **DP_size**: `CP_size` reuses `atten_tp_size`, which is equal to `TP_size` / `DP_size`. For the cp function to work correctly, `TP_size` must be divisible by `DP_size`, and TP_size / DP_size > 1 (to ensure CP_size > 1). +- **Detailed design reference**: https://github.com/sgl-project/sglang/pull/12065 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/gpt_oss.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/gpt_oss.md new file mode 100644 index 00000000..d1af32f5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/gpt_oss.md @@ -0,0 +1,129 @@ +# 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. + +### 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 helfpul 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 helfpul 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. + +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/llama4.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/llama4.md new file mode 100644 index 00000000..1a2338a3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/llama4.md @@ -0,0 +1,92 @@ +# 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**: Add `--hybrid-kvcache-ratio` for hybrid kv cache. Details can be seen in [this PR](https://github.com/sgl-project/sglang/pull/6563) + + +### 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). diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/native_api.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/native_api.ipynb new file mode 100644 index 00000000..028e646d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/native_api.ipynb @@ -0,0 +1,601 @@ +{ + "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", + "- `/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}\")" + ] + }, + { + "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." + ] + }, + { + "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", + "}" + ] + }, + { + "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", + " \"\"\"\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", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "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", + " \"\"\"\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", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "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": [ + "## 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", + " \"\"\"\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", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "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}\")" + ] + }, + { + "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", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path qwen/qwen2.5-0.5b-instruct\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/offline_engine_api.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/offline_engine_api.ipynb new file mode 100644 index 00000000..9c03e90a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/offline_engine_api.ipynb @@ -0,0 +1,235 @@ +{ + "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\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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api.rst b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api.rst new file mode 100644 index 00000000..370abe99 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api.rst @@ -0,0 +1,9 @@ +OpenAI-Compatible APIs +====================== + +.. toctree:: + :maxdepth: 1 + + openai_api_completions.ipynb + openai_api_vision.ipynb + openai_api_embeddings.ipynb diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api_completions.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api_completions.ipynb new file mode 100644 index 00000000..d498f13e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api_completions.ipynb @@ -0,0 +1,534 @@ +{ + "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}\")\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": "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": [ + "## 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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api_embeddings.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api_embeddings.ipynb new file mode 100644 index 00000000..26e95a4e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api_embeddings.ipynb @@ -0,0 +1,195 @@ +{ + "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/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", + " \"\"\"\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", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "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/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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api_vision.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api_vision.ipynb new file mode 100644 index 00000000..1db599dc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/openai_api_vision.ipynb @@ -0,0 +1,254 @@ +{ + "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/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", + "vision_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path Qwen/Qwen2.5-VL-7B-Instruct --log-level warning\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "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\": \"https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true\"\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\": {\n", + " \"url\": \"https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true\"\n", + " },\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\": {\n", + " \"url\": \"https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true\"\n", + " },\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\": \"https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true\",\n", + " },\n", + " },\n", + " {\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {\n", + " \"url\": \"https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png\",\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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/popular_model_usage.rst b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/popular_model_usage.rst new file mode 100644 index 00000000..b8e75f21 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/popular_model_usage.rst @@ -0,0 +1,12 @@ +Popular Model Usage (DeepSeek, GPT-OSS, Llama, Qwen, and more) +=============================================================== + +.. toctree:: + :maxdepth: 1 + + deepseek_v3.md + deepseek_v32.md + gpt_oss.md + qwen3.md + qwen3_vl.md + llama4.md diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/qwen3.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/qwen3.md new file mode 100644 index 00000000..5e6773c0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/qwen3.md @@ -0,0 +1,33 @@ +# 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`. + +### 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). diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/qwen3_vl.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/qwen3_vl.md new file mode 100644 index 00000000..f05e7832 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/qwen3_vl.md @@ -0,0 +1,130 @@ +# 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/sampling_params.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/sampling_params.md new file mode 100644 index 00000000..a97a7368 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/sampling_params.md @@ -0,0 +1,346 @@ +# 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. Can be an image instance, file name, URL, or base64 encoded string. Can be a single image, list of images, or list of lists of images. | +| 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. | + +## 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 structal 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}") +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/send_request.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/send_request.ipynb new file mode 100644 index 00000000..6e457a02 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/basic_usage/send_request.ipynb @@ -0,0 +1,253 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Sending Requests\n", + "This notebook provides a quick-start guide to use SGLang in chat completions after installation.\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", + " \"\"\"\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", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/conf.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/conf.py new file mode 100644 index 00000000..d6ca64d8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/conf.py @@ -0,0 +1,205 @@ +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 + + +""" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/deploy.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/deploy.py new file mode 100644 index 00000000..75b7ea7f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/deploy.py @@ -0,0 +1,22 @@ +# 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" +) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/bench_serving.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/bench_serving.md new file mode 100644 index 00000000..b2f8568e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/bench_serving.md @@ -0,0 +1,355 @@ +# 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 mode-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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/benchmark_and_profiling.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/benchmark_and_profiling.md new file mode 100644 index 00000000..728bcba3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/benchmark_and_profiling.md @@ -0,0 +1,456 @@ +# Benchmark and Profiling + +## Benchmark + +- Benchmark the latency of running a single static batch without a server. The arguments are the same as for `launch_server.py`. + Note that this is a simplified test script without a dynamic batching server, so it may run out of memory for a batch size that a real server can handle. A real server truncates the prefill into several batches, while this simplified script does not. + - Without a server (do not need to launch a server) + ```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 + ``` + - With a server (please use `sglang.launch_server` to launch a server first and run the following command.) + ```bash + python -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 + ``` + + +- Benchmark offline processing. This script will start an offline engine and run the benchmark. + + ```bash + python3 -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10 + ``` + +- Benchmark online serving. Please use `sglang.launch_server` to launch a server first and run the following command. + + ```bash + python3 -m sglang.bench_serving --backend sglang --num-prompt 10 + ``` + +## 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 +``` + +Please make sure that the `SGLANG_TORCH_PROFILER_DIR` should be set at both server and client side, otherwise the trace file cannot be generated correctly . A secure way will be setting `SGLANG_TORCH_PROFILER_DIR` in the `.*rc` file of shell (e.g. `~/.bashrc` for bash shells). + +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). diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/contribution_guide.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/contribution_guide.md new file mode 100644 index 00000000..6a5056ce --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/contribution_guide.md @@ -0,0 +1,140 @@ +# 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. +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) + +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`. + +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`). + +If you don’t have permission, please ask maintainers to trigger CI for you. + +### CI rate limits + +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 package, 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 sgl-kernel (e.g., [#9220](https://github.com/sgl-project/sglang/pull/9220/files)). + - Once merged, this will trigger an automatic release of the sgl-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 sgl-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.ai). + +Thank you for your interest in SGLang. Happy coding! diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/development_guide_using_docker.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/development_guide_using_docker.md new file mode 100644 index 00000000..e3894790 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/development_guide_using_docker.md @@ -0,0 +1,108 @@ +# 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: Monting 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/release_process.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/release_process.md new file mode 100644 index 00000000..e817a256 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/release_process.md @@ -0,0 +1,18 @@ +# 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/setup_github_runner.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/setup_github_runner.md new file mode 100644 index 00000000..3ca9627f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/developer_guide/setup_github_runner.md @@ -0,0 +1,51 @@ +# Set Up Self-Hosted Runners for GitHub Action + +## 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.0rc1-rocm630 /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.0rc1-rocm630 /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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/get_started/install.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/get_started/install.md new file mode 100644 index 00000000..a4911c41 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/get_started/install.md @@ -0,0 +1,185 @@ +# 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" --prerelease=allow +``` + +**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.6 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 +``` + +You can also find the nightly docker images [here](https://hub.docker.com/r/lmsysorg/sglang/tags?name=nightly). + +## 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) + +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`. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/index.rst b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/index.rst new file mode 100644 index 00000000..d605d9cc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/index.rst @@ -0,0 +1,106 @@ +SGLang Documentation +==================== + +SGLang is a high-performance serving framework for large language models and vision-language 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 Backend 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. +- **Extensive Model Support**: Supports a wide range of generative 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 integrating 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. +- **Flexible Frontend Language**: Offers an intuitive interface for programming LLM applications, supporting chained generation calls, advanced prompting, control flow, multi-modal inputs, parallelism, and external interactions. +- **Active Community**: SGLang is open-source and supported by a vibrant community with widespread industry adoption, powering over 400,000 GPUs worldwide. + +.. 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/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/expert_parallelism.md + advanced_features/lora.ipynb + advanced_features/pd_disaggregation.md + advanced_features/hicache.rst + advanced_features/pd_multiplexing.md + advanced_features/vlm_query.ipynb + advanced_features/router.md + advanced_features/deterministic_inference.md + advanced_features/observability.md + advanced_features/checkpoint_engine.md + +.. toctree:: + :maxdepth: 1 + :caption: Supported Models + + supported_models/generative_models.md + supported_models/multimodal_language_models.md + supported_models/embedding_models.md + supported_models/reward_models.md + supported_models/rerank_models.md + supported_models/support_new_models.md + supported_models/transformers_fallback.md + supported_models/modelscope.md + +.. toctree:: + :maxdepth: 1 + :caption: Hardware Platforms + + platforms/amd_gpu.md + platforms/cpu_server.md + platforms/tpu.md + platforms/nvidia_jetson.md + platforms/ascend_npu.md + platforms/xpu.md + +.. toctree:: + :maxdepth: 1 + :caption: Developer Guide + + developer_guide/contribution_guide.md + developer_guide/development_guide_using_docker.md + developer_guide/benchmark_and_profiling.md + developer_guide/bench_serving.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/learn_more.md + references/post_training_integration.md + +.. toctree:: + :maxdepth: 1 + :caption: Security Acknowledgement + + security/acknowledgements.md diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/amd_gpu.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/amd_gpu.md new file mode 100644 index 00000000..95e92dcc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/amd_gpu.md @@ -0,0 +1,159 @@ +# AMD GPUs + +This document describes how 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.6 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 +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. + +## 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/ascend_npu.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/ascend_npu.md new file mode 100644 index 00000000..cd251ba3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/ascend_npu.md @@ -0,0 +1,190 @@ +# Ascend NPUs + +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. + +## 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 +``` + +## Installing SGLang + +### 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 +``` + +#### MemFabric Adaptor + +_TODO: MemFabric is still a working project yet open sourced til end of year 2025. We will release it as prebuilt wheel package for now._ + +MemFabric Adaptor is a drop-in replacement of Mooncake Transfer Engine that enables KV cache transfer on Ascend NPU clusters. + +```shell +pip install mf-adapter==1.0.0 +``` + +#### Pytorch and Pytorch Framework Adaptor on Ascend + +```shell +PYTORCH_VERSION="2.8.0" +TORCHVISION_VERSION="0.23.0" +pip install torch==$PYTORCH_VERSION torchvision==$TORCHVISION_VERSION --index-url https://download.pytorch.org/whl/cpu + +PTA_VERSION="2.8.0" +pip install torch-npu==$PTA_VERSION +``` + +#### Triton on Ascend + +_Notice:_ We recommend installing triton-ascend from source due to its rapid development, the version on PYPI can't keep up for know. This problem will be solved on Sep. 2025, afterwards `pip install` would be the one and only installing method. + +Please follow Triton-on-Ascend's [installation guide from source](https://gitee.com/ascend/triton-ascend#2%E6%BA%90%E4%BB%A3%E7%A0%81%E5%AE%89%E8%A3%85-triton-ascend) to install the latest `triton-ascend` package. + +#### DeepEP-compatible Library + +We are also providing 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 -b v0.5.6 https://github.com/sgl-project/sglang.git +cd sglang + +pip install --upgrade pip +rm -vf python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml +pip install -e python[srt_npu] +``` + +### Method 2: Using 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 +# Clone the SGLang repository +git clone https://github.com/sgl-project/sglang.git +cd sglang/docker + +# Build the docker image +docker build -t -f npu.Dockerfile . + +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/' + +drun --env "HF_TOKEN=" \ + \ + python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --attention-backend ascend --host 0.0.0.0 --port 30000 +``` + +## Examples + +### Running DeepSeek-V3 + +Running DeepSeek with PD disaggregation on 2 x Atlas 800I A3. +Model weights could be found [here](https://modelers.cn/models/State_Cloud/Deepseek-R1-bf16-hfd-w8a8). + +Prefill: + +```shell +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export ASCEND_MF_STORE_URL="tcp://:" + +drun \ + python3 -m sglang.launch_server --model-path State_Cloud/DeepSeek-R1-bf16-hfd-w8a8 \ + --trust-remote-code \ + --attention-backend ascend \ + --mem-fraction-static 0.8 \ + --quantization w8a8_int8 \ + --tp-size 16 \ + --dp-size 1 \ + --nnodes 1 \ + --node-rank 0 \ + --disaggregation-mode prefill \ + --disaggregation-bootstrap-port 6657 \ + --disaggregation-transfer-backend ascend \ + --dist-init-addr :6688 \ + --host \ + --port 8000 +``` + +Decode: + +```shell +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export ASCEND_MF_STORE_URL="tcp://:" +export HCCL_BUFFSIZE=200 +export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=24 +export SGLANG_NPU_USE_MLAPO=1 + +drun \ + python3 -m sglang.launch_server --model-path State_Cloud/DeepSeek-R1-bf16-hfd-w8a8 \ + --trust-remote-code \ + --attention-backend ascend \ + --mem-fraction-static 0.8 \ + --quantization w8a8_int8 \ + --enable-deepep-moe \ + --deepep-mode low_latency \ + --tp-size 16 \ + --dp-size 1 \ + --ep-size 16 \ + --nnodes 1 \ + --node-rank 0 \ + --disaggregation-mode decode \ + --disaggregation-transfer-backend ascend \ + --dist-init-addr :6688 \ + --host \ + --port 8001 +``` + +Mini_LB: + +```shell +drun \ + python -m sglang.srt.disaggregation.launch_lb \ + --prefill http://:8000 \ + --decode http://:8001 \ + --host 127.0.0.1 --port 5000 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/cpu_server.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/cpu_server.md new file mode 100644 index 00000000..71be9f6f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/cpu_server.md @@ -0,0 +1,216 @@ +# 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'd prefer to install SGLang in a bare metal environment, +the command list is as below. +It is worth noting that the environment variable `SGLANG_USE_CPU_ENGINE=1` +is required to enable SGLang service with CPU engine. + +```bash +# Create and activate a conda environment +conda create -n sgl-cpu python=3.12 -y +conda activate sgl-cpu + +# Set PyTorch CPU as primary pip install channel to avoid installing the larger CUDA-enabled version and prevent potential runtime issues. +pip config set global.index-url https://download.pytorch.org/whl/cpu +pip config set global.extra-index-url https://pypi.org/simple + +# Check if some conda related environment variables have been set +env | grep -i conda +# The following environment variable settings are required +# if they have not been set properly +export CONDA_EXE=$(which conda) +export CONDA_ROOT=${CONDA_EXE}/../.. +export CONDA_PREFIX=${CONDA_ROOT}/envs/sgl-cpu +export PATH=${PATH}:${CONDA_ROOT}/bin:${CONDA_ROOT}/condabin + +# 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 +pip install --upgrade pip setuptools +conda install -y libsqlite==3.48.0 gperftools tbb libnuma numactl +pip install . +pip install torch==2.9.0 torchvision==0.24.0 triton==3.5.0 --force-reinstall + +# Build the CPU backend kernels +cd ../sgl-kernel +cp pyproject_cpu.toml pyproject.toml +pip install . + +# Other required environment variables +# Recommend to set these in ~/.bashrc in order not to set every time in a new terminal +export SGLANG_USE_CPU_ENGINE=1 +export LD_PRELOAD=${LD_PRELOAD}:${CONDA_PREFIX}/lib/libiomp5.so:${CONDA_PREFIX}/lib/libtcmalloc.so:${CONDA_PREFIX}/lib/libtbbmalloc.so.2 +``` + +## 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. + Users can specify TP to be no more than the total available SNCs in current system. + + 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. + + To specify the cores to be used, we need to explicitly set the environment variable `SGLANG_CPU_OMP_THREADS_BIND`. + For example, if we want to run the SGLang service 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" + ``` + + 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. Currently the maximum applicable batch size + for optimizing with `torch.compile` is 16. + +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. + +```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. + +## Example: Running DeepSeek-V3.1-Terminus + +An example command to launch service for 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 \ + --mem-fraction-static 0.8 \ + --enable-torch-compile \ + --torch-compile-max-bs 4 \ + --tp 6 +``` + +Similarly, an example command to launch service for 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 \ + --mem-fraction-static 0.8 \ + --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. + +Then you can test with `bench_serving` command or construct your own command or script +following [the benchmarking example](#benchmarking-with-requests). diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/nvidia_jetson.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/nvidia_jetson.md new file mode 100644 index 00000000..ba3b68ae --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/nvidia_jetson.md @@ -0,0 +1,80 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/tpu.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/tpu.md new file mode 100644 index 00000000..925287c3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/tpu.md @@ -0,0 +1,3 @@ +# TPU + +SGLang supports TPU inference via the SGLang-Jax backend. Please go to https://github.com/sgl-project/sglang-jax. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/xpu.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/xpu.md new file mode 100644 index 00000000..88fa1552 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/platforms/xpu.md @@ -0,0 +1,92 @@ +# 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.9.0+xpu torchao torchvision torchaudio pytorch-triton-xpu==3.5.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 . +``` + +### 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/custom_chat_template.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/custom_chat_template.md new file mode 100644 index 00000000..f22ee8be --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/custom_chat_template.md @@ -0,0 +1,51 @@ +# 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/environment_variables.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/environment_variables.md new file mode 100644 index 00000000..d8a1e67e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/environment_variables.md @@ -0,0 +1,127 @@ +# 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_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) | + +## 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 | `true` | +| `SGLANG_SET_CPU_AFFINITY` | Enable CPU affinity setting (often set to `1` in Docker builds) | `0` | +| `SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN` | Allows the scheduler to overwrite longer context length requests (often set to `1` in Docker builds) | `0` | +| `SGLANG_IS_FLASHINFER_AVAILABLE` | Control FlashInfer availability check | `true` | +| `SGLANG_SKIP_P2P_CHECK` | Skip P2P (peer-to-peer) access check | `false` | +| `SGL_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_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_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` | + + +## DeepGEMM Configuration (Advanced Optimization) + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_ENABLE_JIT_DEEPGEMM` | Enable Just-In-Time compilation of DeepGEMM kernels | `"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` | +| `SGL_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` | +| `SGL_DG_USE_NVRTC` | Use NVRTC (instead of Triton) for JIT compilation (Experimental) | `"0"` | +| `SGL_USE_DEEPGEMM_BMM` | Use DeepGEMM for Batched Matrix Multiplication (BMM) operations | `"false"` | + +## DeepEP Configuration + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_DEEPEP_BF16_DISPATCH` | Use Bfloat16 for dispatch | `"false"` | +| `SGLANG_MOE_NVFP4_DISPATCH` | Use nvfp4 for moe dispatch | `"false"` | + +## 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 | +| `SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK` | Disable checks for memory imbalance across Tensor Parallel ranks | Not set (defaults to enabled check) | + +## Model-Specific Options + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_USE_AITER` | Use AITER optimize implementation | `false` | +| `SGLANG_INT4_WEIGHT` | Enable INT4 weight quantization | `false` | +| `SGLANG_MOE_PADDING` | Enable MoE padding (sets padding size to 128 if value is `1`, often set to `1` in Docker builds) | `0` | +| `SGLANG_FORCE_FP8_MARLIN` | Force using FP8 MARLIN kernels even if other FP8 kernels are available | `false` | +| `SGLANG_ENABLE_FLASHINFER_FP8_GEMM` | Use flashinfer kernels when running blockwise fp8 GEMM on Blackwell GPUs | `false` | +| `SGLANG_FLASHINFER_FP4_GEMM_BACKEND` | Select backend for `mm_fp4` on Blackwell GPUS | `` | +| `SGLANG_SUPPORT_CUTLASS_BLOCK_FP8` | Use Cutlass kernels when running blockwise fp8 GEMM on Hopper or Blackwell GPUs | `false` | +| `SGLANG_CUTLASS_MOE` (deprecated) | Use Cutlass FP8 MoE kernel on Blackwell GPUs (deprecated, use --moe-runner-backend=cutlass) | `false` | + + +## Distributed Computing + +| Environment Variable | Description | Default Value | +| --- | --- | --- | +| `SGLANG_BLOCK_NONZERO_RANK_CHILDREN` | Control blocking of non-zero rank children processes | `1` | +| `SGL_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 | `0` | +| `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` | +| `SGLANG_CI_SMALL_KV_SIZE` | Use small KV cache size in CI | Not set | + +## 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 | `true` | + +## 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` | diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/faq.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/faq.md new file mode 100644 index 00000000..ffa1a7c5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/faq.md @@ -0,0 +1,39 @@ +# 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/ diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/frontend/choices_methods.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/frontend/choices_methods.md new file mode 100644 index 00000000..30a0a181 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/frontend/choices_methods.md @@ -0,0 +1,77 @@ +# 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, + ) + ) +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/frontend/frontend_index.rst b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/frontend/frontend_index.rst new file mode 100644 index 00000000..62544cba --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/frontend/frontend_index.rst @@ -0,0 +1,9 @@ +Frontend Language +================= + +.. toctree:: + :maxdepth: 1 + :caption: Frontend Language + + frontend_tutorial.ipynb + choices_methods.md diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/frontend/frontend_tutorial.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/frontend/frontend_tutorial.ipynb new file mode 100644 index 00000000..166f8cac --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/frontend/frontend_tutorial.ipynb @@ -0,0 +1,456 @@ +{ + "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}\")\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/generative_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}\")\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://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true\"\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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/learn_more.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/learn_more.md new file mode 100644 index 00000000..f0d6ffb8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/learn_more.md @@ -0,0 +1,9 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/mindspore_models.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/mindspore_models.md new file mode 100644 index 00000000..80dd3b7f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/mindspore_models.md @@ -0,0 +1,164 @@ +# 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.RC1. + +## 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`, which needs to be installed separately. + +```shell +git clone https://github.com/chz34/sgl-mindspore.git +cd sgl-mindspore +pip install -e . +``` + +You will need to install the following packages. + +```shell +pip install "mindspore==2.7.1" +pip install "torch==2.8" +pip install "torch_npu==2.8" +pip install triton_ascend +``` + +```shell +cp python/pyproject_other.toml python/pyproject.toml +pip install -e "python[all_npu]" +``` + +## 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.01, "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/) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/deploy_on_k8s.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/deploy_on_k8s.md new file mode 100644 index 00000000..cfc099f5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/deploy_on_k8s.md @@ -0,0 +1,337 @@ +# 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). diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/d-svc.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/d-svc.yaml new file mode 100644 index 00000000..27f98009 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/d-svc.yaml @@ -0,0 +1,12 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/d.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/d.yaml new file mode 100644 index 00000000..dbb51b51 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/d.yaml @@ -0,0 +1,290 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/lb.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/lb.yaml new file mode 100644 index 00000000..4ca69096 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/lb.yaml @@ -0,0 +1,56 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/p-svc.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/p-svc.yaml new file mode 100644 index 00000000..6826a13d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/p-svc.yaml @@ -0,0 +1,12 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/p.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/p.yaml new file mode 100644 index 00000000..1c5b5870 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws-examples/p.yaml @@ -0,0 +1,304 @@ +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: SGL_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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws_pd_deploy.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws_pd_deploy.md new file mode 100644 index 00000000..b3508968 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/lws_pd/lws_pd_deploy.md @@ -0,0 +1,783 @@ +# 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: SGL_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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/multi_node.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/multi_node.md new file mode 100644 index 00000000..e6e5b534 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/multi_node.md @@ -0,0 +1,100 @@ +# 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). diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/multi_node_index.rst b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/multi_node_index.rst new file mode 100644 index 00000000..78636869 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/multi_node_index.rst @@ -0,0 +1,14 @@ +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 `_ diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/rbg_pd/deepseekv32_pd.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/rbg_pd/deepseekv32_pd.md new file mode 100644 index 00000000..d4dcf73a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/multi_node_deployment/rbg_pd/deepseekv32_pd.md @@ -0,0 +1,569 @@ +# 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" + - --prefill-round-robin-balance + - --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" + - --prefill-round-robin-balance + - --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). diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/post_training_integration.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/post_training_integration.md new file mode 100644 index 00000000..5e82f837 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/post_training_integration.md @@ -0,0 +1,31 @@ +# 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 maxium-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 parternes, 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/production_metrics.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/production_metrics.md new file mode 100644 index 00000000..85a6ff8a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/production_metrics.md @@ -0,0 +1,231 @@ +# 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/production_request_trace.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/production_request_trace.md new file mode 100644 index 00000000..2d19570c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/production_request_trace.md @@ -0,0 +1,160 @@ +# 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 mini lb + 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. + +## How to add Tracing for slices you're interested in? +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. + +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. Mark the beginning and end of a request + ``` + trace_req_start(rid, bootstrap_room) + trace_req_finish(rid) + ``` + These two APIs must be called within the same process, for example, in the tokenizer. + +3. Add tracing for slice + + * Add slice tracing normally: + ```python + trace_slice_start("slice A", rid) + trace_slice_end("slice A", rid) + ``` + + - Use the "anonymous" flag to not specify a slice name at the start of the slice, allowing the slice name to be determined by trace_slice_end. +
Note: Anonymous slices must not be nested. + ```python + trace_slice_start("", rid, anonymous = True) + trace_slice_end("slice A", rid) + ``` + + - In trace_slice_end, use auto_next_anon to automatically create the next anonymous slice, which can reduce the number of instrumentation points needed. + ```python + trace_slice_start("", rid, anonymous = True) + trace_slice_end("slice A", rid, auto_next_anon = True) + trace_slice_end("slice B", rid, auto_next_anon = True) + trace_slice_end("slice C", rid, auto_next_anon = True) + trace_slice_end("slice D", rid) + ``` + - The end of the last slice in a thread must be marked with thread_finish_flag=True; otherwise, the thread's span will not be properly generated. + ```python + trace_slice_end("slice D", rid, thread_finish_flag = True) + ``` + +4. When the request execution flow transfers to another thread, the trace context needs to be explicitly propagated. + - sender: Execute the following code before sending the request to another thread via ZMQ + ```python + trace_context = trace_get_proc_propagate_context(rid) + req.trace_context = trace_context + ``` + - receiver: Execute the following code after receiving the request via ZMQ + ```python + trace_set_proc_propagate_context(rid, req.trace_context) + ``` + +5. When the request execution flow transfers to another node(PD disaggregation), the trace context needs to be explicitly propagated. + - sender: Execute the following code before sending the request to node thread via http + ```python + trace_context = trace_get_remote_propagate_context(bootstrap_room_list) + headers = {"trace_context": trace_context} + session.post(url, headers=headers) + ``` + - receiver: Execute the following code after receiving the request via http + ```python + trace_set_remote_propagate_context(request.headers['trace_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 two-level trace context structure and a four-level span structure: `SglangTraceReqContext`, `SglangTraceThreadContext`. Their relationship is as follows: +``` +SglangTraceReqContext (req_id="req-123") +├── SglangTraceThreadContext(thread_label="scheduler", tp_rank=0) +| +└── SglangTraceThreadContext(thread_label="scheduler", tp_rank=1) +``` + +Each traced request maintains a global `SglangTraceReqContext`. For every thread processing the request, a corresponding `SglangTraceThreadContext` is recorded and composed within the `SglangTraceReqContext`. Within each thread, every currently traced slice (possibly nested) is stored in a list. + +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. + +When the request execution flow transfers to a new thread, the trace context needs to be explicitly propagated. In the framework, this is represented by `SglangTracePropagateContext`, which contains the context of the request span and the previous slice span. + + +We designed a four-level span structure, consisting of `bootstrap_room_span`, `req_root_span`, `thread_span`, and `slice_span`. Among them, `req_root_span` and `thread_span` correspond to `SglangTraceReqContext` and `SglangTraceThreadContext`, respectively, and `slice_span` is stored within the `SglangTraceThreadContext`. The `bootstrap_room_span` is designed to accommodate the separation of PD-disaggregation. On different nodes, we may want to add certain attributes to the `req_root_span`. However, if the `req_root_span` is shared across all nodes, the Prefill and Decode nodes would not be allowed to add attributes due to the constraints imposed by OpenTelemetry's design. + +``` +bootstrap room span +├── router req root span +| └── router thread span +| └── slice span +├── prefill req root span +| ├── tokenizer thread span +| | └── slice span +| └── scheduler thread span +| └── slice span +└── decode req root span + ├── tokenizer thread span + | └── slice span + └── scheduler thread span + └── slice span +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/torch_compile_cache.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/torch_compile_cache.md new file mode 100644 index 00000000..f2bb257f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/references/torch_compile_cache.md @@ -0,0 +1,13 @@ +# 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`. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/requirements.txt b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/requirements.txt new file mode 100644 index 00000000..5d730967 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/requirements.txt @@ -0,0 +1,20 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/serve.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/serve.sh new file mode 100644 index 00000000..049f767c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/serve.sh @@ -0,0 +1,3 @@ +# Clean and serve documentation with auto-build +make clean +make serve diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/classify_models.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/classify_models.md new file mode 100644 index 00000000..c6d18f9a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/classify_models.md @@ -0,0 +1,162 @@ +# 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 Router**: Handles routing and request/response models in `sgl-router/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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/embedding_models.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/embedding_models.md new file mode 100644 index 00000000..906466ac --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/embedding_models.md @@ -0,0 +1,126 @@ +# 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 | diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/generative_models.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/generative_models.md new file mode 100644 index 00000000..671fbaaf --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/generative_models.md @@ -0,0 +1,64 @@ +# 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)| +| **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, 3MoE, 3Next, 2.5, 2 series) | `Qwen/Qwen3-0.6B`, `Qwen/Qwen3-30B-A3B` `Qwen/Qwen3-Next-80B-A3B-Instruct ` | 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-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** | `minimax/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** (Cohere) | `CohereForAI/c4ai-command-r-v01` | 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. | +| **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. | diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/modelscope.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/modelscope.md new file mode 100644 index 00000000..4740c277 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/modelscope.md @@ -0,0 +1,28 @@ +# 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/multimodal_language_models.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/multimodal_language_models.md new file mode 100644 index 00000000..8aa642d5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/multimodal_language_models.md @@ -0,0 +1,111 @@ +# 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. | | +| **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` | +| **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. | +| **JetVLM** | | JetVLM is an vision-language model designed for high-performance multimodal understanding and generation tasks built upon Jet-Nemotron. | Coming soon | + +## 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` | For video, the processor is configured to sample at 2 FPS, at a max of 128 frames, as per model training. | +| **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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/rerank_models.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/rerank_models.md new file mode 100644 index 00000000..b6f2ffa2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/rerank_models.md @@ -0,0 +1,49 @@ +# 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} +They are executed with `--is-embedding` and some may require `--trust-remote-code` +``` + +## Example 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." + ] +} + +response = requests.post(url, json=payload) +response_json = response.json() + +for item in response_json: + print(f"Score: {item['score']:.2f} - Document: '{item['document']}'") +``` + +## 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. | diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/reward_models.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/reward_models.md new file mode 100644 index 00000000..ef447463 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/reward_models.md @@ -0,0 +1,28 @@ +# 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. | diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/support_new_models.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/support_new_models.md new file mode 100644 index 00000000..b71e06c4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/support_new_models.md @@ -0,0 +1,320 @@ +# 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,mllm} + ``` +- 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/srt/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](https://github.com/sgl-project/sglang/blob/main/docs/supported_models/generative_models.md) or [multimodal_language_models.md](https://github.com/sgl-project/sglang/blob/main/docs/supported_models/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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/transformers_fallback.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/transformers_fallback.md new file mode 100644 index 00000000..3c7dd961 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/supported_models/transformers_fallback.md @@ -0,0 +1,58 @@ +# 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! diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/wrap_run_llm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/wrap_run_llm.py new file mode 100644 index 00000000..2d21442c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/docs/wrap_run_llm.py @@ -0,0 +1,47 @@ +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/assets/.gitignore b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/assets/.gitignore new file mode 100644 index 00000000..fc787e33 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/assets/.gitignore @@ -0,0 +1 @@ +!example_image.png diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/assets/example_image.png b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/assets/example_image.png new file mode 100644 index 00000000..851d0856 Binary files /dev/null and b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/assets/example_image.png differ diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekr1.jinja b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekr1.jinja new file mode 100644 index 00000000..bb86cbd9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekr1.jinja @@ -0,0 +1,92 @@ +{% 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'] %} + {%- 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['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}} + {%- else %} + {{content + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\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['function']['arguments'] + '\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 %} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekv3.jinja b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekv3.jinja new file mode 100644 index 00000000..fdde62ee --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekv3.jinja @@ -0,0 +1,92 @@ +{% 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 %} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekv31.jinja b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekv31.jinja new file mode 100644 index 00000000..a97f011f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekv31.jinja @@ -0,0 +1,92 @@ +{% 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 %} + {% 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 %} + {{'<｜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 %} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekv32.jinja b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekv32.jinja new file mode 100644 index 00000000..b6d239dc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_deepseekv32.jinja @@ -0,0 +1,101 @@ +{% 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 %} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_llama3.1_json.jinja b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_llama3.1_json.jinja new file mode 100644 index 00000000..f609c388 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_llama3.1_json.jinja @@ -0,0 +1,121 @@ +{# 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 %} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_llama4_pythonic.jinja b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_llama4_pythonic.jinja new file mode 100644 index 00000000..74b315c3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/tool_chat_template_llama4_pythonic.jinja @@ -0,0 +1,112 @@ +{# 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 %} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/vision_template_sarashina_vl.jinja b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/vision_template_sarashina_vl.jinja new file mode 100644 index 00000000..caff3441 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/chat_template/vision_template_sarashina_vl.jinja @@ -0,0 +1,9 @@ +{# + 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 %} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/checkpoint_engine/update.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/checkpoint_engine/update.py new file mode 100644 index 00000000..86b588cc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/checkpoint_engine/update.py @@ -0,0 +1,241 @@ +""" +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/anthropic_example_chat.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/anthropic_example_chat.py new file mode 100644 index 00000000..03d699be --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/anthropic_example_chat.py @@ -0,0 +1,73 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/anthropic_example_complete.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/anthropic_example_complete.py new file mode 100644 index 00000000..bce2a61e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/anthropic_example_complete.py @@ -0,0 +1,68 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/azure_openai_example_chat.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/azure_openai_example_chat.py new file mode 100644 index 00000000..d53f935f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/azure_openai_example_chat.py @@ -0,0 +1,83 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/gemini_example_chat.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/gemini_example_chat.py new file mode 100644 index 00000000..0ae62310 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/gemini_example_chat.py @@ -0,0 +1,73 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/gemini_example_complete.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/gemini_example_complete.py new file mode 100644 index 00000000..5188bf41 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/gemini_example_complete.py @@ -0,0 +1,68 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/gemini_example_multimodal_chat.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/gemini_example_multimodal_chat.py new file mode 100644 index 00000000..afe0c723 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/gemini_example_multimodal_chat.py @@ -0,0 +1,30 @@ +""" +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"]) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/images/cat.jpeg b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/images/cat.jpeg new file mode 100644 index 00000000..a6a8e48c Binary files /dev/null and b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/images/cat.jpeg differ diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/images/dog.jpeg b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/images/dog.jpeg new file mode 100644 index 00000000..dc4b40e2 Binary files /dev/null and b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/images/dog.jpeg differ diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/local_example_chat.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/local_example_chat.py new file mode 100644 index 00000000..e1e4b62c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/local_example_chat.py @@ -0,0 +1,75 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/local_example_complete.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/local_example_complete.py new file mode 100644 index 00000000..00a451cf --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/local_example_complete.py @@ -0,0 +1,70 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/local_example_llava_next.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/local_example_llava_next.py new file mode 100644 index 00000000..c941a549 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/local_example_llava_next.py @@ -0,0 +1,78 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_chat.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_chat.py new file mode 100644 index 00000000..9511e21c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_chat.py @@ -0,0 +1,74 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_complete.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_complete.py new file mode 100644 index 00000000..d64bcaf1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_complete.py @@ -0,0 +1,68 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_n.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_n.py new file mode 100644 index 00000000..25372b9f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_n.py @@ -0,0 +1,71 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_o1.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_o1.py new file mode 100644 index 00000000..2e5c1400 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openai_example_o1.py @@ -0,0 +1,57 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openrouter_example_chat.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openrouter_example_chat.py new file mode 100644 index 00000000..a0b6f15b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/openrouter_example_chat.py @@ -0,0 +1,81 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/together_example_chat.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/together_example_chat.py new file mode 100644 index 00000000..2d205906 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/together_example_chat.py @@ -0,0 +1,81 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/together_example_complete.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/together_example_complete.py new file mode 100644 index 00000000..d9119ed6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/quick_start/together_example_complete.py @@ -0,0 +1,76 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/chinese_regex.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/chinese_regex.py new file mode 100644 index 00000000..78e9c7e1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/chinese_regex.py @@ -0,0 +1,53 @@ +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/choices_logprob.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/choices_logprob.py new file mode 100644 index 00000000..6cd733fe --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/choices_logprob.py @@ -0,0 +1,44 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/cot_decoding.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/cot_decoding.py new file mode 100644 index 00000000..7a7a04bc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/cot_decoding.py @@ -0,0 +1,115 @@ +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, +) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/json_decode.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/json_decode.py new file mode 100644 index 00000000..5dc3522d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/json_decode.py @@ -0,0 +1,83 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/json_logprobs.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/json_logprobs.py new file mode 100644 index 00000000..15206a61 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/json_logprobs.py @@ -0,0 +1,103 @@ +# 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") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.py new file mode 100644 index 00000000..ec5b334b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.py @@ -0,0 +1,260 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.sh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.sh new file mode 100755 index 00000000..ffb1af96 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.sh @@ -0,0 +1,131 @@ +#!/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." diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/openai_chat_speculative.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/openai_chat_speculative.py new file mode 100644 index 00000000..f3fd74ed --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/openai_chat_speculative.py @@ -0,0 +1,155 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/openai_speculative.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/openai_speculative.py new file mode 100644 index 00000000..4389cb05 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/openai_speculative.py @@ -0,0 +1,54 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/parallel_sample.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/parallel_sample.py new file mode 100644 index 00000000..0f3cf170 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/parallel_sample.py @@ -0,0 +1,40 @@ +""" +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}") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/rag_using_parea/trace_and_evaluate_rag_using_parea.ipynb b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/rag_using_parea/trace_and_evaluate_rag_using_parea.ipynb new file mode 100644 index 00000000..f309142a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/rag_using_parea/trace_and_evaluate_rag_using_parea.ipynb @@ -0,0 +1,408 @@ +{ + "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", + "\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", + "\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", + "\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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/readme_examples.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/readme_examples.py new file mode 100644 index 00000000..7269ef14 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/readme_examples.py @@ -0,0 +1,109 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/sgl_gen_min_tokens.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/sgl_gen_min_tokens.py new file mode 100644 index 00000000..a5088199 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/sgl_gen_min_tokens.py @@ -0,0 +1,35 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/streaming.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/streaming.py new file mode 100644 index 00000000..506ee35c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/streaming.py @@ -0,0 +1,49 @@ +""" +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()) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/triton/Dockerfile b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/triton/Dockerfile new file mode 100644 index 00000000..e4741a1d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/triton/Dockerfile @@ -0,0 +1,10 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/triton/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/triton/README.md new file mode 100644 index 00000000..b2e55961 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/triton/README.md @@ -0,0 +1,35 @@ +# 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"] +}' + +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/triton/models/character_generation/config.pbtxt b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/triton/models/character_generation/config.pbtxt new file mode 100644 index 00000000..7546f993 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/frontend_language/usage/triton/models/character_generation/config.pbtxt @@ -0,0 +1,23 @@ +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 ] + } +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/README.md new file mode 100644 index 00000000..3eef0b09 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/README.md @@ -0,0 +1,76 @@ +# 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/docker-compose.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/docker-compose.yaml new file mode 100644 index 00000000..ce6457fa --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/docker-compose.yaml @@ -0,0 +1,28 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/grafana/dashboards/config/dashboard.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/grafana/dashboards/config/dashboard.yaml new file mode 100644 index 00000000..6c17a6c6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/grafana/dashboards/config/dashboard.yaml @@ -0,0 +1,11 @@ +apiVersion: 1 +providers: + - name: 'SGLang' + orgId: 1 + folder: 'SGLang Monitoring' + type: file + disableDeletion: false + updateIntervalSeconds: 10 + allowUiUpdates: false + options: + path: /var/lib/grafana/dashboards diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/grafana/dashboards/json/sglang-dashboard.json b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/grafana/dashboards/json/sglang-dashboard.json new file mode 100644 index 00000000..ebd663e4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/grafana/dashboards/json/sglang-dashboard.json @@ -0,0 +1,984 @@ +{ + "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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/grafana/datasources/datasource.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/grafana/datasources/datasource.yaml new file mode 100644 index 00000000..1ab0e4a5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/grafana/datasources/datasource.yaml @@ -0,0 +1,8 @@ +apiVersion: 1 +datasources: + - name: Prometheus + type: prometheus + access: proxy + url: http://localhost:9090 + isDefault: true + editable: false diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/opentelemetry.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/opentelemetry.yaml new file mode 100644 index 00000000..8593d918 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/opentelemetry.yaml @@ -0,0 +1,38 @@ +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] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/prometheus.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/prometheus.yaml new file mode 100644 index 00000000..ba16ac3b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/prometheus.yaml @@ -0,0 +1,10 @@ +# prometheus.yaml +global: + scrape_interval: 5s + evaluation_interval: 30s + +scrape_configs: + - job_name: sglang + static_configs: + - targets: + - '127.0.0.1:30000' diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/tracing_compose.yaml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/tracing_compose.yaml new file mode 100644 index 00000000..7ed1ecdd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/monitoring/tracing_compose.yaml @@ -0,0 +1,21 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/profiler/nsys_profile_tools/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/profiler/nsys_profile_tools/README.md new file mode 100644 index 00000000..687200e0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/profiler/nsys_profile_tools/README.md @@ -0,0 +1,176 @@ +# 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/profiler/nsys_profile_tools/gputrc2graph.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/profiler/nsys_profile_tools/gputrc2graph.py new file mode 100755 index 00000000..f17bd185 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/profiler/nsys_profile_tools/gputrc2graph.py @@ -0,0 +1,344 @@ +""" + 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 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: + engine_model.update(json.load(f)) + 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 + self.pd.options.mode.copy_on_write = True + + # 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 = " ".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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/profiler/nsys_profile_tools/sglang_engine_model.json b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/profiler/nsys_profile_tools/sglang_engine_model.json new file mode 100644 index 00000000..253cc762 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/profiler/nsys_profile_tools/sglang_engine_model.json @@ -0,0 +1,61 @@ +{ + "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" + } + } +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/README.md new file mode 100644 index 00000000..8b623fc3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/README.md @@ -0,0 +1,45 @@ +# 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/custom_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/custom_server.py new file mode 100644 index 00000000..b190a463 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/custom_server.py @@ -0,0 +1,53 @@ +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/embedding.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/embedding.py new file mode 100644 index 00000000..b927a188 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/embedding.py @@ -0,0 +1,27 @@ +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/fastapi_engine_inference.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/fastapi_engine_inference.py new file mode 100644 index 00000000..f5da9d71 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/fastapi_engine_inference.py @@ -0,0 +1,195 @@ +""" +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/launch_engine.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/launch_engine.py new file mode 100644 index 00000000..b97e0bcd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/launch_engine.py @@ -0,0 +1,17 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference.py new file mode 100644 index 00000000..92e68dcd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference.py @@ -0,0 +1,43 @@ +""" +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_async.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_async.py new file mode 100644 index 00000000..578962d7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_async.py @@ -0,0 +1,65 @@ +""" +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)) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_eagle.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_eagle.py new file mode 100644 index 00000000..a7a89ef5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_eagle.py @@ -0,0 +1,38 @@ +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_qwen_1m.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_qwen_1m.py new file mode 100644 index 00000000..664efa6d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_qwen_1m.py @@ -0,0 +1,74 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_vlm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_vlm.py new file mode 100644 index 00000000..939e6910 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/offline_batch_inference_vlm.py @@ -0,0 +1,52 @@ +""" +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/readme.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/readme.md new file mode 100644 index 00000000..fc6e89a6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/readme.md @@ -0,0 +1,54 @@ +# 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/save_remote_state.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/save_remote_state.py new file mode 100644 index 00000000..a428195c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/save_remote_state.py @@ -0,0 +1,59 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/save_sharded_state.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/save_sharded_state.py new file mode 100644 index 00000000..80ad5321 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/engine/save_sharded_state.py @@ -0,0 +1,74 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/hidden_states/hidden_states_engine.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/hidden_states/hidden_states_engine.py new file mode 100644 index 00000000..60ab302c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/hidden_states/hidden_states_engine.py @@ -0,0 +1,66 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/hidden_states/hidden_states_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/hidden_states/hidden_states_server.py new file mode 100644 index 00000000..b04f7437 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/hidden_states/hidden_states_server.py @@ -0,0 +1,81 @@ +""" +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}") + + 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/lora.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/lora.py new file mode 100644 index 00000000..181dc231 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/lora.py @@ -0,0 +1,67 @@ +""" +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 ..." + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/llama3_llava_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/llama3_llava_server.py new file mode 100644 index 00000000..a8409af7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/llama3_llava_server.py @@ -0,0 +1,111 @@ +""" +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 + + +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"{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 = f"{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="http://127.0.0.1") + parser.add_argument("--port", type=int, default=30000) + args = parser.parse_args() + asyncio.run(test_concurrent(args)) + test_streaming(args) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/llava_onevision_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/llava_onevision_server.py new file mode 100644 index 00000000..2cf16e3b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/llava_onevision_server.py @@ -0,0 +1,264 @@ +""" +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 decord import VideoReader, cpu +from PIL import Image + +# pip install httpx==0.23.3 +# pip install decord +# 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 + vr = VideoReader(video_path, ctx=cpu(0)) + total_frame_num = len(vr) + uniform_sampled_frames = np.linspace( + 0, total_frame_num - 1, max_frames_num, dtype=int + ) + frame_idx = uniform_sampled_frames.tolist() + frames = vr.get_batch(frame_idx).asnumpy() + + 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/pixtral_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/pixtral_server.py new file mode 100644 index 00000000..d907de14 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/pixtral_server.py @@ -0,0 +1,127 @@ +""" +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 + +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"{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"{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="http://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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/qwen_llava_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/qwen_llava_server.py new file mode 100644 index 00000000..d8b3226e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal/qwen_llava_server.py @@ -0,0 +1,111 @@ +""" +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 + + +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"{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 = f"{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="http://127.0.0.1") + parser.add_argument("--port", type=int, default=30000) + args = parser.parse_args() + asyncio.run(test_concurrent(args)) + test_streaming(args) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal_embedding.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal_embedding.py new file mode 100644 index 00000000..4e8d748b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/multimodal_embedding.py @@ -0,0 +1,18 @@ +# 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", [])]) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/openai_chat_with_response_prefill.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/openai_chat_with_response_prefill.py new file mode 100644 index 00000000..6d803a1d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/openai_chat_with_response_prefill.py @@ -0,0 +1,53 @@ +""" +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/reward_model.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/reward_model.py new file mode 100644 index 00000000..1a1177e6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/reward_model.py @@ -0,0 +1,32 @@ +# 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]) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_engine.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_engine.py new file mode 100644 index 00000000..11453f93 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_engine.py @@ -0,0 +1,43 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_server.py new file mode 100644 index 00000000..7e498f51 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_server.py @@ -0,0 +1,68 @@ +""" +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}") + + # 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_engine.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_engine.py new file mode 100644 index 00000000..bf6f1971 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_engine.py @@ -0,0 +1,74 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_server.py new file mode 100644 index 00000000..392e1bf0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_server.py @@ -0,0 +1,78 @@ +""" +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}") + + 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/vertex_predict.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/vertex_predict.py new file mode 100644 index 00000000..58a41b1c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/runtime/vertex_predict.py @@ -0,0 +1,66 @@ +""" +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/sagemaker/deploy_and_serve_endpoint.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/sagemaker/deploy_and_serve_endpoint.py new file mode 100644 index 00000000..e518183c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/sagemaker/deploy_and_serve_endpoint.py @@ -0,0 +1,69 @@ +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}'") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/usage/modelopt_quantize_and_export.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/usage/modelopt_quantize_and_export.py new file mode 100755 index 00000000..4394d917 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/examples/usage/modelopt_quantize_and_export.py @@ -0,0 +1,303 @@ +#!/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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject.toml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject.toml new file mode 100755 index 00000000..b5919cd0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject.toml @@ -0,0 +1,167 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "sglang" +version = "0.5.6" +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", + "av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' and platform_machine == 'armv7l')", + "blobfile==3.0.0", + "build", + "compressed-tensors", + "cuda-python", + "decord2", + "datasets", + "einops", + "fastapi", + "flashinfer_python==0.5.3", # keep it aligned with jit-cache version in Dockerfile + "flashinfer_cubin==0.5.3", + "gguf", + "hf_transfer", + "huggingface_hub", + "interegular", + "llguidance>=0.7.11,<0.8.0", + "modelscope", + "msgspec", + "ninja", + "numpy", + "nvidia-cutlass-dsl==4.2.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", + "requests", + "scipy", + "sentencepiece", + "setproctitle", + "sgl-kernel==0.3.18.post2", + "soundfile==0.13.1", + "tiktoken", + "timm==1.0.16", + "torch_memory_saver==0.0.9", + "torch==2.9.1", + "torchaudio==2.9.1", + "torchcodec==0.7.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. If not provided, transformer will use torchvision instead by default. + "torchvision", + "torchao==0.9.0", + "tqdm", + "transformers==4.57.1", + "uvicorn", + "uvloop", + "xgrammar==0.1.27", + + "grpcio==1.75.1", # keep it align with compile_proto.py + "grpcio-tools==1.75.1", # keep it align with compile_proto.py + "grpcio-reflection==1.75.1", # required by srt/entrypoints/grpc_server.py + "grpcio-health-checking==1.75.1", # required for Kubernetes gRPC health probes +] + +[project.optional-dependencies] +checkpoint-engine = ["checkpoint-engine==0.1.2"] +diffusion = [ + "PyYAML==6.0.1", + "cloudpickle", + "diffusers==0.35.2", + "imageio==2.36.0", + "imageio-ffmpeg==0.5.1", + "moviepy>=2.0.0", + "opencv-python==4.10.0.84", + "remote-pdb", + "st_attn ==0.0.7", + "vsa==0.0.4", + "yunchang==0.6.3.post1", +] + +[tool.uv.extra-build-dependencies] +st-attn = ["torch", "setuptools"] +vsa = ["torch", "setuptools"] + +test = [ + "accelerate", + "expecttest", + "jsonlines", + "matplotlib", + "pandas", + "peft", + "pytest", + "sentence_transformers", + "tabulate", +] +dev = ["sglang[test]"] +tracing = [ + "opentelemetry-api", + "opentelemetry-exporter-otlp", + "opentelemetry-exporter-otlp-proto-grpc", + "opentelemetry-sdk", +] + +[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/layers/moe/fused_moe_triton/configs/*/*.json", + "srt/layers/quantization/configs/*.json", + "srt/mem_cache/storage/hf3fs/hf3fs_utils.cpp", + "srt/mem_cache/cpp_radix_tree/*.cpp", + "srt/mem_cache/cpp_radix_tree/*.h", + "srt/speculative/cpp_ngram/*.cpp", + "srt/speculative/cpp_ngram/*.h", + "jit_kernel/include/sgl_kernel/*.h", + "jit_kernel/include/sgl_kernel/*.cuh", + "jit_kernel/csrc/*.cuh" +] + +[tool.setuptools.packages.find] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] + +[tool.wheel] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] + +[tool.codespell] +ignore-words-list = "ans, als, hel, boostrap, childs, te, vas, hsa, ment" +skip = "*.json,*.jsonl,*.patch,*.txt" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject_cpu.toml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject_cpu.toml new file mode 100644 index 00000000..a537bc36 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject_cpu.toml @@ -0,0 +1,134 @@ +# https://docs.sglang.io/platforms/cpu_server.html +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "sglang" +version = "0.5.6" +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", + "decord", + "einops", + "fastapi", + "gguf", + "hf_transfer", + "huggingface_hub", + "intel-openmp", + "interegular", + "llguidance>=0.7.11,<0.8.0", + "modelscope", + "msgspec", + "ninja", + "numpy", + "openai-harmony==0.0.4", + "openai==1.99.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", + "grpcio==1.75.1", # keep it align with compile_proto.py + "grpcio-tools==1.75.1", # keep it align with compile_proto.py + "grpcio-reflection==1.75.1", # required by srt/entrypoints/grpc_server.py +] + +[project.optional-dependencies] +tracing = [ + "opentelemetry-sdk", + "opentelemetry-api", + "opentelemetry-exporter-otlp", + "opentelemetry-exporter-otlp-proto-grpc", +] +test = [ + "accelerate", + "expecttest", + "jsonlines", + "matplotlib", + "pandas", + "peft", + "pytest", + "sentence_transformers", + "tabulate", +] +all = [] +dev = ["sglang[test]"] + +[project.urls] +"Homepage" = "https://github.com/sgl-project/sglang" +"Bug Tracker" = "https://github.com/sgl-project/sglang/issues" + +[tool.setuptools.package-data] +"sglang" = [ + "srt/layers/moe/fused_moe_triton/configs/*/*.json", + "srt/layers/quantization/configs/*.json", + "srt/mem_cache/storage/hf3fs/hf3fs_utils.cpp", + "srt/mem_cache/cpp_radix_tree/*.cpp", + "srt/mem_cache/cpp_radix_tree/*.h", + "srt/speculative/cpp_ngram/*.cpp", + "srt/speculative/cpp_ngram/*.h", + "jit_kernel/include/sgl_kernel/*.h", + "jit_kernel/include/sgl_kernel/*.cuh", + "jit_kernel/csrc/*.cuh" +] + +[tool.setuptools.packages.find] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] + +[tool.wheel] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] + +[tool.codespell] +ignore-words-list = "ans, als, hel, boostrap, childs, te, vas, hsa, ment" +skip = "*.json,*.jsonl,*.patch,*.txt" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject_other.toml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject_other.toml new file mode 100755 index 00000000..3e854caa --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject_other.toml @@ -0,0 +1,156 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "sglang" +version = "0.5.6" +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==1.99.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", + "grpcio==1.75.1", # keep it align with compile_proto.py + "grpcio-tools==1.75.1", # keep it align with compile_proto.py + "grpcio-reflection==1.75.1", # required by srt/entrypoints/grpc_server.py +] + +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", +] + +# https://docs.sglang.io/platforms/ascend_npu.html +srt_npu = ["sglang[runtime_common]"] + +# 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]"] + +test = [ + "accelerate", + "expecttest", + "gguf", + "jsonlines", + "matplotlib", + "pandas", + "peft", + "pytest", + "sentence_transformers", + "tabulate", +] +all_hip = ["sglang[srt_hip]"] +all_npu = ["sglang[srt_npu]"] +all_hpu = ["sglang[srt_hpu]"] + +dev_hip = ["sglang[all_hip]", "sglang[test]"] +dev_npu = ["sglang[all_npu]", "sglang[test]"] +dev_hpu = ["sglang[all_hpu]", "sglang[test]"] + +[project.urls] +"Homepage" = "https://github.com/sgl-project/sglang" +"Bug Tracker" = "https://github.com/sgl-project/sglang/issues" + +[tool.setuptools.package-data] +"sglang" = [ + "srt/layers/moe/fused_moe_triton/configs/*/*.json", + "srt/layers/quantization/configs/*.json", + "srt/mem_cache/storage/hf3fs/hf3fs_utils.cpp", + "srt/mem_cache/cpp_radix_tree/*.cpp", + "srt/mem_cache/cpp_radix_tree/*.h", + "srt/speculative/cpp_ngram/*.cpp", + "srt/speculative/cpp_ngram/*.h", + "jit_kernel/include/sgl_kernel/*.h", + "jit_kernel/include/sgl_kernel/*.cuh", + "jit_kernel/csrc/*.cuh" +] + +[tool.setuptools.packages.find] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] + +[tool.wheel] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] + +[tool.codespell] +ignore-words-list = "ans, als, hel, boostrap, childs, te, vas, hsa, ment" +skip = "*.json,*.jsonl,*.patch,*.txt" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject_xpu.toml b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject_xpu.toml new file mode 100644 index 00000000..a058e6f4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/pyproject_xpu.toml @@ -0,0 +1,138 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "sglang" +version = "0.5.6" +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.9.0", + "torchcodec==0.8.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. If not provided, transformer will use torchvision instead by default. + "av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' and platform_machine == 'armv7l')", + "torchaudio==2.9.0", + "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", + "decord", + "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==1.99.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.24", , xgrammar depends on CUDA PyTorch and Triton only + "grpcio==1.75.1", # keep it align with compile_proto.py + "grpcio-tools==1.75.1", # keep it align with compile_proto.py + "grpcio-reflection==1.75.1", # required by srt/entrypoints/grpc_server.py +] + +[project.optional-dependencies] +tracing = [ + "opentelemetry-sdk", + "opentelemetry-api", + "opentelemetry-exporter-otlp", + "opentelemetry-exporter-otlp-proto-grpc", +] +test = [ + "accelerate", + "expecttest", + "jsonlines", + "matplotlib", + "pandas", + "peft", + "pytest", + "sentence_transformers", + "tabulate", +] +all = [] +dev = ["sglang[test]"] + +[project.urls] +"Homepage" = "https://github.com/sgl-project/sglang" +"Bug Tracker" = "https://github.com/sgl-project/sglang/issues" + +[tool.setuptools.package-data] +"sglang" = [ + "srt/layers/moe/fused_moe_triton/configs/*/*.json", + "srt/layers/quantization/configs/*.json", + "srt/mem_cache/storage/hf3fs/hf3fs_utils.cpp", + "srt/mem_cache/cpp_radix_tree/*.cpp", + "srt/mem_cache/cpp_radix_tree/*.h", + "srt/speculative/cpp_ngram/*.cpp", + "srt/speculative/cpp_ngram/*.h", + "jit_kernel/include/sgl_kernel/*.h", + "jit_kernel/include/sgl_kernel/*.cuh", + "jit_kernel/csrc/*.cuh" +] + +[tool.setuptools.packages.find] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] + +[tool.wheel] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] + +[tool.codespell] +ignore-words-list = "ans, als, hel, boostrap, childs, te, vas, hsa, ment" +skip = "*.json,*.jsonl,*.patch,*.txt" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/README.md new file mode 100644 index 00000000..4d9cf8c2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/README.md @@ -0,0 +1,17 @@ +# Code Structure + +- `eval`: The evaluation utilities. +- `lang`: The frontend language. +- `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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/__init__.py new file mode 100644 index 00000000..509b145a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/__init__.py @@ -0,0 +1,83 @@ +# SGLang public APIs + +# 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__", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_offline_throughput.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_offline_throughput.py new file mode 100644 index 00000000..294d3f68 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_offline_throughput.py @@ -0,0 +1,476 @@ +""" +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.bench_serving import ( + DatasetRow, + get_dataset, + get_tokenizer, + sample_random_requests, + 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 throughput_test( + server_args: ServerArgs, + bench_args: BenchArgs, +): + if bench_args.backend == "engine": + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_one_batch.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_one_batch.py new file mode 100644 index 00000000..25b16d31 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_one_batch.py @@ -0,0 +1,795 @@ +""" +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 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.managers.schedule_batch import Req, ScheduleBatch +from sglang.srt.managers.scheduler_dp_attn_mixin import prepare_mlp_sync_batch_raw +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, + is_cuda_alike, + is_xpu, + 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 + +profile_activities = [torch.profiler.ProfilerActivity.CPU] + [ + profiler_activity + for available, profiler_activity in [ + (is_cuda_alike(), torch.profiler.ProfilerActivity.CUDA), + (is_xpu(), torch.profiler.ProfilerActivity.XPU), + ] + if available +] + + +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 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" + + @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"], + help="Profiler activities: CPU, GPU, CUDA_PROFILER. If CPU/GPU, 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"', + ) + + @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} + ) + + +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): + 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.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + req.logprob_start_len = len(req.origin_input_ids) - 1 + 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 = 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 + ] + req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + req.logprob_start_len = len(req.origin_input_ids) - 1 + 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.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + req.logprob_start_len = len(req.origin_input_ids) - 1 + reqs.append(req) + + return reqs + + +@torch.no_grad +def extend(reqs, model_runner): + # Create dummy tree_cache for benchmarks (no prefix caching, just allocation) + dummy_tree_cache = SimpleNamespace( + 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) + 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) + 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, +): + 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 = [] + profile_step_of_interest = output_len // 2 + enable_profile_decode = profile and profile_stage in ["all", "decode"] + for i in range(output_len - 1): + synchronize(device) + profiler = None + if enable_profile_decode and i == profile_step_of_interest: + 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 + + if enable_profile_decode and i == profile_step_of_interest: + 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", + ) + + 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) + + # 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, + ) + + 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, + ) + 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_one_batch_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_one_batch_server.py new file mode 100644 index 00000000..159a3a93 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_one_batch_server.py @@ -0,0 +1,605 @@ +""" +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 +import dataclasses +import itertools +import json +import multiprocessing +import os +import random +import time +from typing import List, Optional, Tuple + +import numpy as np +import requests +from pydantic import BaseModel +from transformers import AutoProcessor, PreTrainedTokenizer + +from sglang.bench_serving import ( + get_processor, + get_tokenizer, + sample_mmmu_requests, + sample_random_requests, +) +from sglang.profiler import run_profile +from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import is_blackwell, kill_process_tree +from sglang.test.nightly_bench_utils import save_results_as_pydantic_models +from sglang.test.test_utils import is_in_ci, write_github_step_summary + + +@dataclasses.dataclass +class BenchArgs: + run_name: str = "default" + batch_size: Tuple[int] = (1,) + input_len: Tuple[int] = (1024,) + output_len: Tuple[int] = (16,) + temperature: float = 0.0 + return_logprob: bool = False + client_stream_interval: int = 1 + input_len_step_percentage: float = 0.0 + base_url: str = "" + skip_warmup: bool = False + show_report: bool = False + profile: bool = False + profile_steps: int = 5 + profile_by_stage: bool = False + profile_prefix: Optional[str] = None + profile_output_dir: Optional[str] = None + dataset_path: str = "" + dataset_name: str = "random" + parallel_batch: bool = False + result_filename: str = "result.jsonl" + pydantic_result_filename: Optional[str] = None + append_to_github_summary: bool = True + seed: int = 42 + + @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("--temperature", type=float, default=BenchArgs.temperature) + parser.add_argument("--return-logprob", action="store_true") + parser.add_argument( + "--client-stream-interval", + type=int, + default=BenchArgs.client_stream_interval, + ) + parser.add_argument( + "--input-len-step-percentage", + type=float, + default=BenchArgs.input_len_step_percentage, + ) + parser.add_argument("--base-url", type=str, default=BenchArgs.base_url) + parser.add_argument("--skip-warmup", action="store_true") + parser.add_argument("--show-report", action="store_true") + parser.add_argument("--profile", action="store_true") + parser.add_argument( + "--profile-steps", type=int, default=BenchArgs.profile_steps + ) + parser.add_argument("--profile-by-stage", action="store_true") + parser.add_argument( + "--profile-prefix", + type=str, + default=BenchArgs.profile_prefix, + ) + parser.add_argument( + "--profile-output-dir", + type=str, + default=BenchArgs.profile_output_dir, + ) + parser.add_argument( + "--dataset-path", + type=str, + default=BenchArgs.dataset_path, + help="Path to the dataset.", + ) + parser.add_argument( + "--dataset-name", + type=str, + default=BenchArgs.dataset_name, + choices=["mmmu", "random"], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument("--parallel-batch", action="store_true") + parser.add_argument( + "--result-filename", + type=str, + default=BenchArgs.result_filename, + help="Store the results line by line in the JSON Line format to this file.", + ) + parser.add_argument( + "--pydantic-result-filename", + type=str, + default=BenchArgs.pydantic_result_filename, + help="Store the results as pydantic models in the JSON format to this file.", + ) + parser.add_argument( + "--no-append-to-github-summary", + action="store_false", + dest="append_to_github_summary", + help="Disable appending the output of this run to github ci summary", + ) + parser.add_argument("--seed", type=int, default=BenchArgs.seed) + + @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}) + + +class BenchOneCaseResult(BaseModel): + run_name: str + batch_size: int + input_len: int + output_len: int + latency: float + input_throughput: float + output_throughput: float + overall_throughput: float + last_ttft: float + last_gen_throughput: float + acc_length: float + profile_link: Optional[str] = None + + def dump_to_jsonl(self, result_filename: str): + with open(result_filename, "a") as fout: + res = { + "run_name": self.run_name, + "batch_size": self.batch_size, + "input_len": self.input_len, + "output_len": self.output_len, + "latency": round(self.latency, 4), + "input_throughput": round(self.input_throughput, 2), + "output_throughput": round(self.output_throughput, 2), + "overall_throughput": round(self.overall_throughput, 2), + "last_ttft": round(self.last_ttft, 4), + "last_gen_throughput": round(self.last_gen_throughput, 2), + "acc_length": round(self.acc_length, 2), + } + fout.write(json.dumps(res) + "\n") + + +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(server_args: ServerArgs): + proc = multiprocessing.Process(target=launch_server_internal, args=(server_args,)) + proc.start() + base_url = f"http://{server_args.host}:{server_args.port}" + timeout = 600 + + start_time = time.time() + while time.time() - start_time < timeout: + try: + headers = { + "Content-Type": "application/json; charset=utf-8", + } + response = requests.get(f"{base_url}/v1/models", headers=headers) + if response.status_code == 200: + return proc, base_url + except requests.RequestException: + pass + time.sleep(10) + raise TimeoutError("Server failed to start within the timeout period.") + + +def run_one_case( + url: str, + batch_size: int, + input_len: int, + output_len: int, + temperature: float, + return_logprob: bool, + stream_interval: int, + input_len_step_percentage: float, + run_name: str, + result_filename: str, + tokenizer: PreTrainedTokenizer | AutoProcessor, + profile: bool = False, + profile_steps: int = BenchArgs.profile_steps, + profile_by_stage: bool = False, + profile_prefix: Optional[str] = BenchArgs.profile_prefix, + profile_output_dir: Optional[str] = BenchArgs.profile_output_dir, + dataset_name: str = BenchArgs.dataset_name, + dataset_path: str = BenchArgs.dataset_path, + parallel_batch: bool = False, +): + requests.post(url + "/flush_cache") + + # Load input token ids + # TODO: reuse bench_serving.get_dataset ? + if dataset_name == "mmmu": + input_requests = sample_mmmu_requests( + num_requests=batch_size, + processor=tokenizer, + fixed_output_len=output_len, + random_sample=False, + ) + elif dataset_name == "random": + input_requests = sample_random_requests( + input_len=input_len, + output_len=output_len, + num_prompts=batch_size, + range_ratio=1.0, + tokenizer=tokenizer, + dataset_path=dataset_path, + random_sample=True, + return_text=False, + ) + + # Load sampling parameters + use_structured_outputs = False + if use_structured_outputs: + texts = [] + for _ in range(batch_size): + texts.append( + "Human: What is the capital city of france? can you give as many trivial information as possible about that city? answer in json.\n" + * 50 + + "Assistant:" + ) + json_schema = "$$ANY$$" + else: + json_schema = None + + payload = { + "sampling_params": { + "temperature": temperature, + "max_new_tokens": output_len, + "ignore_eos": True, + "json_schema": json_schema, + "stream_interval": stream_interval, + }, + "return_logprob": return_logprob, + "stream": True, + **({"parallel_batch": parallel_batch} if parallel_batch else {}), + } + if dataset_name == "mmmu": + # vlm + input_ids = [] + # for vlms, tokenizer is an instance of AutoProcessor + tokenizer = tokenizer.tokenizer + for input_req in input_requests: + input_ids += [tokenizer.encode(input_req.prompt)] + payload["image_data"] = [req.image_data for req in input_requests] + + else: + input_ids = [req.prompt for req in input_requests] + + payload["input_ids"] = input_ids + + # Turn on profiler + profile_link = None + if profile: + profile_link: str = run_profile( + url=url, + num_steps=profile_steps, + activities=["CPU", "GPU"], + output_dir=profile_output_dir, + profile_by_stage=profile_by_stage, + profile_prefix=profile_prefix, + ) + + # Run the request + tic = time.perf_counter() + response = requests.post( + url + "/generate", + json=payload, + stream=True, + ) + + # Get the TTFT of the last request in the batch + last_ttft = 0.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")) + if "error" in data: + raise RuntimeError(f"Request has failed. {data}.") + + assert ( + data["meta_info"]["finish_reason"] is None + or data["meta_info"]["finish_reason"]["type"] == "length" + ) + if data["meta_info"]["completion_tokens"] == 1: + last_ttft = time.perf_counter() - tic + + # Compute metrics + latency = time.perf_counter() - tic + input_throughput = batch_size * input_len / last_ttft + output_throughput = batch_size * output_len / (latency - last_ttft) + overall_throughput = batch_size * (input_len + output_len) / latency + + server_info = requests.get(url + "/get_server_info").json() + internal_state = server_info.get("internal_states", [{}]) + last_gen_throughput = internal_state[0].get("last_gen_throughput", None) or -1 + acc_length = internal_state[0].get("avg_spec_accept_length", None) or -1 + + # Print results + print(f"batch size: {batch_size}") + print(f"input_len: {input_len}") + print(f"output_len: {output_len}") + print(f"latency: {latency:.2f} s") + print(f"input throughput: {input_throughput:.2f} tok/s") + if output_len != 1: + print(f"output throughput: {output_throughput:.2f} tok/s") + print(f"last_ttft: {last_ttft:.2f} s") + print(f"last generation throughput: {last_gen_throughput:.2f} tok/s") + if acc_length > 0: + print(f"acc_length: {acc_length:.2f} ") + + # Dump results + result = BenchOneCaseResult( + run_name=run_name, + batch_size=batch_size, + input_len=input_len, + output_len=output_len, + latency=latency, + input_throughput=input_throughput, + output_throughput=output_throughput, + overall_throughput=overall_throughput, + last_ttft=last_ttft, + last_gen_throughput=last_gen_throughput, + acc_length=acc_length, + profile_link=profile_link, + ) + + # Save and return the results + if result_filename: + result.dump_to_jsonl(result_filename) + + return result + + +def should_skip_due_to_token_capacity( + batch_size, input_len, output_len, skip_token_capacity_threshold +): + if batch_size * (input_len + output_len) > skip_token_capacity_threshold: + print( + "=" * 8 + + f"Skip benchmark {batch_size=} * ({input_len=} + {output_len=}) = {batch_size * (input_len + output_len)} > {skip_token_capacity_threshold=} due to kv cache limit." + + "=" * 8 + ) + return True + return False + + +def get_report_summary( + results: List[BenchOneCaseResult], bench_args: BenchArgs, server_args: ServerArgs +): + summary = ( + f"\nInput lens: {bench_args.input_len}. Output lens: {bench_args.output_len}.\n" + ) + summary += "| batch size | input len | latency (s) | input throughput (tok/s) | output throughput (tok/s) | acc length | ITL (ms) | input cost ($/1M) | output cost ($/1M) |" + + if bench_args.profile: + summary += " profile |" + + summary += "\n" + summary += "| ---------- | --------- | ----------- | ------------------------- | ------------------------- | ---------- | -------- | ----------------- | ------------------ |" + + if bench_args.profile: + summary += "-------------|" + summary += "\n" + + if is_blackwell(): + hourly_cost_per_gpu = 4 # $4/hour for one B200 + else: + hourly_cost_per_gpu = 2 # $2/hour for one H100 + input_util = 0.7 + + # sort result by input_len + results.sort(key=lambda x: x.input_len) + for res in results: + hourly_cost = hourly_cost_per_gpu * server_args.tp_size + accept_length = round(res.acc_length, 2) if res.acc_length > 0 else "n/a" + line = ( + f"| {res.batch_size} | " + f"{res.input_len} | " + f"{res.latency:.2f} | " + f"{res.input_throughput:.2f} | " + f"{res.output_throughput:.2f} | " + f"{accept_length} | " + f"{1 / (res.output_throughput/res.batch_size) * 1000:.2f} | " + f"{1e6 / (res.input_throughput * input_util) / 3600 * hourly_cost:.2f} | " + f"{1e6 / res.output_throughput / 3600 * hourly_cost:.2f} |" + ) + if bench_args.profile: + if res.profile_link: + line += f" [Profile]({res.profile_link}) |" + else: + line += f" n/a |" + line += "\n" + summary += line + + return summary + + +def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs): + if bench_args.base_url: + proc, base_url = None, bench_args.base_url + else: + proc, base_url = launch_server_process(server_args) + + # Get tokenizer + server_info = requests.get(base_url + "/get_server_info").json() + if "tokenizer_path" in server_info: + tokenizer_path = server_info["tokenizer_path"] + elif "prefill" in server_info: + tokenizer_path = server_info["prefill"][0]["tokenizer_path"] + if bench_args.dataset_name == "mmmu": + # mmmu implies this is a MLLM + tokenizer = get_processor(tokenizer_path) + else: + tokenizer = get_tokenizer(tokenizer_path) + + # Get token capacity + internal_state = server_info.get("internal_states", [{}]) + skip_token_capacity_threshold = ( + internal_state[0].get("memory_usage", {}).get("token_capacity", 1000000000) + ) + + # Warmup + if not bench_args.skip_warmup: + print("=" * 8 + " Warmup Begin " + "=" * 8) + print(f"Warmup with batch_size={bench_args.batch_size}") + for bs in bench_args.batch_size: + run_one_case( + base_url, + batch_size=bs, + input_len=1024, + output_len=16, + temperature=bench_args.temperature, + return_logprob=bench_args.return_logprob, + stream_interval=bench_args.client_stream_interval, + input_len_step_percentage=bench_args.input_len_step_percentage, + run_name="", + result_filename="", + tokenizer=tokenizer, + dataset_name=bench_args.dataset_name, + dataset_path=bench_args.dataset_path, + parallel_batch=bench_args.parallel_batch, + ) + print("=" * 8 + " Warmup End " + "=" * 8 + "\n") + + results = [] + profile_results = [] + try: + # Benchmark all cases + for bs, il, ol in itertools.product( + bench_args.batch_size, bench_args.input_len, bench_args.output_len + ): + if should_skip_due_to_token_capacity( + bs, il, ol, skip_token_capacity_threshold + ): + continue + results.append( + run_one_case( + base_url, + bs, + il, + ol, + temperature=bench_args.temperature, + return_logprob=bench_args.return_logprob, + stream_interval=bench_args.client_stream_interval, + input_len_step_percentage=bench_args.input_len_step_percentage, + run_name=bench_args.run_name, + result_filename=bench_args.result_filename, + tokenizer=tokenizer, + dataset_name=bench_args.dataset_name, + dataset_path=bench_args.dataset_path, + parallel_batch=bench_args.parallel_batch, + ) + ) + + # Profile all cases + if bench_args.profile: + try: + for bs, il, ol in itertools.product( + bench_args.batch_size, bench_args.input_len, bench_args.output_len + ): + if should_skip_due_to_token_capacity( + bs, il, ol, skip_token_capacity_threshold + ): + continue + profile_prefix = ( + bench_args.profile_prefix or "" + ) + f"bs-{bs}-il-{il}" + profile_results.append( + run_one_case( + base_url, + bs, + il, + ol, + temperature=bench_args.temperature, + return_logprob=bench_args.return_logprob, + stream_interval=bench_args.client_stream_interval, + input_len_step_percentage=bench_args.input_len_step_percentage, + run_name=bench_args.run_name, + result_filename=bench_args.result_filename, + tokenizer=tokenizer, + dataset_name=bench_args.dataset_name, + dataset_path=bench_args.dataset_path, + parallel_batch=bench_args.parallel_batch, + profile=bench_args.profile, + profile_steps=bench_args.profile_steps, + profile_by_stage=bench_args.profile_by_stage, + profile_prefix=profile_prefix, + profile_output_dir=bench_args.profile_output_dir, + ) + ) + + # Replace the profile link + for res, profile_res in zip(results, profile_results): + res.profile_link = profile_res.profile_link + except Exception as e: + print(f"Error profiling, there will be no profile trace dump: {e}") + finally: + if proc: + kill_process_tree(proc.pid) + + print(f"\nResults are saved to {bench_args.result_filename}") + + if not bench_args.show_report: + return + + # Print summary + summary = get_report_summary(results, bench_args, server_args) + print(summary) + + if is_in_ci() and bench_args.append_to_github_summary: + write_github_step_summary(summary) + else: + print(summary) + + # 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, + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + BenchArgs.add_cli_args(parser) + args = parser.parse_args() + + random.seed(args.seed) + np.random.seed(args.seed) + + server_args = ServerArgs.from_cli_args(args) + bench_args = BenchArgs.from_cli_args(args) + + run_benchmark(server_args, bench_args) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_serving.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_serving.py new file mode 100644 index 00000000..00267489 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/bench_serving.py @@ -0,0 +1,2867 @@ +# 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 importlib.util +import io +import json +import os +import pickle +import random +import resource +import shutil +import sys +import time +import traceback +import warnings +from argparse import ArgumentParser +from dataclasses import dataclass, field +from datetime import datetime +from functools import lru_cache +from json import JSONDecodeError +from pathlib import Path +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union + +import aiohttp +import numpy as np +import pybase64 +import requests +from datasets import load_dataset +from PIL import Image +from tqdm.asyncio import tqdm +from transformers import ( + AutoProcessor, + AutoTokenizer, + PreTrainedTokenizer, + PreTrainedTokenizerBase, + PreTrainedTokenizerFast, +) + +ASSISTANT_SUFFIX = "Assistant:" + +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: 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 + + +@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 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 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 {} + + +# 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 "" + 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: + 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, + } + + # 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_auth_headers() + + 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 "" + 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'." + + if 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: + payload = { + "model": request_func_input.model, + "messages": messages, + "temperature": 0.0, + "max_completion_tokens": request_func_input.output_len, + "stream": not args.disable_stream, + "ignore_eos": not args.disable_ignore_eos, + **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_auth_headers() + + 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 "" + 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_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_auth_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 "" + 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, + "logprob_start_len": -1, + **request_func_input.extra_request_body, + } + + # 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_auth_headers() + + 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 "" + 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_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: + body = { + "activities": getattr(args, "profile_activities", []), + "num_steps": getattr(args, "profile_num_steps", None), + "profile_by_stage": getattr(args, "profile_by_stage", None), + "profile_stages": getattr(args, "profile_stages", None), + } + print(f"async_request_profile {api_url=} {body=}") + async with session.post(url=api_url, json=body) 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 _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}" + ) + + +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, +) -> 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_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 get_dataset(args, tokenizer, model_id=None): + tokenize_prompt = getattr(args, "tokenize_prompt", False) + if args.dataset_name == "sharegpt": + assert not tokenize_prompt + input_requests = sample_sharegpt_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + 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, + ) + elif args.dataset_name.startswith("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, + random_sample=args.dataset_name == "random", + return_text=not tokenize_prompt, + ) + elif args.dataset_name == "image": + processor = get_processor(model_id) + input_requests = sample_image_requests( + 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, + processor=processor, + image_content=args.image_content, + image_format=args.image_format, + image_resolution=args.image_resolution, + backend=args.backend, + ) + elif args.dataset_name == "generated-shared-prefix": + assert not tokenize_prompt + 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, + tokenizer=tokenizer, + args=args, + ) + elif args.dataset_name == "mmmu": + processor = get_processor(model_id) + input_requests = sample_mmmu_requests( + num_requests=args.num_prompts, + processor=processor, + backend=args.backend, + fixed_output_len=args.random_output_len, + random_sample=True, + ) + elif args.dataset_name == "mooncake": + # For mooncake, we don't generate the prompts here. + # We just load the raw trace data. The async generator will handle the rest. + if not args.dataset_path: + local_path = os.path.join("/tmp", args.mooncake_workload + "_trace.jsonl") + else: + local_path = args.dataset_path + + if not os.path.exists(local_path): + download_and_cache_file( + MOONCAKE_DATASET_URL[args.mooncake_workload], local_path + ) + + with open(local_path, "r") as f: + all_requests_data = [json.loads(line) for line in f if line.strip()] + + # Limit the number of requests based on --num-prompts + input_requests = all_requests_data[: args.num_prompts] + else: + raise ValueError(f"Unknown dataset: {args.dataset_name}") + return input_requests + + +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, + "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 + p99_e2e_latency_ms: float + concurrency: float + max_output_tokens_per_s: float = 0.0 + max_concurrent_requests: int = 0 + + +SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/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", +} + + +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 + + +@dataclass +class DatasetRow: + prompt: str + 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 + + 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 + + +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}) + + +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 + + +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_file(SHAREGPT_URL) + + # 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 + + +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 = 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 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_file(SHAREGPT_URL) + + # 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=int(input_lens[i]), + output_len=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): + input_content = [ + (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=int(input_lens[i]), + output_len=int(output_lens[i]), + ) + ) + + print(f"#Input tokens: {np.sum(input_lens)}") + print(f"#Output tokens: {np.sum(output_lens)}") + return input_requests + + +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, +) -> List[DatasetRow]: + """Generate requests with images. + + - Each request includes ``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) + + # Check for potentially problematic combinations and warn user + if width * height >= 1920 * 1080 and image_count * num_requests >= 100: + warnings.warn( + f"High resolution ({width}x{height}) with {image_count * num_requests} total images " + f"may take a long time. Consider reducing resolution or image count.", + UserWarning, + stacklevel=2, + ) + + # Sample text lengths + input_lens = np.random.randint( + max(int(input_len * range_ratio), 1), input_len + 1, size=num_requests + ) + output_lens = np.random.randint( + int(output_len * range_ratio), output_len + 1, size=num_requests + ) + + def _gen_random_image_data_uri( + width: int = width, height: int = height + ) -> (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): + # 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(image_count)] + ) + total_image_bytes += sum(list(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(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"\nCreated {len(dataset)} {image_content} {image_format} images with average {total_image_bytes // num_requests} bytes per request" + ) + return dataset + + +@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) + + +def get_gen_prefix_cache_path(args, tokenizer): + """Create cache directory under ~/.cache/sglang/benchmark""" + cache_dir = Path.home() / ".cache" / "sglang" / "benchmark" + + # Create a unique cache filename based on the generation parameters + cache_key = ( + f"gen_shared_prefix_{args.seed}_{args.gsp_num_groups}_{args.gsp_prompts_per_group}_" + f"{args.gsp_system_prompt_len}_{args.gsp_question_len}_{args.gsp_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, + tokenizer: PreTrainedTokenizerBase, + args: argparse.Namespace, +) -> List[DatasetRow]: + """Generate benchmark requests with shared system prompts using random tokens and caching.""" + cache_path = get_gen_prefix_cache_path(args, 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] + 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.append( + DatasetRow( + prompt=full_prompt, + prompt_len=prompt_len, + output_len=output_len, + ) + ) + total_input_tokens += prompt_len + total_output_tokens += output_len + + # 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)}") + 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 + + +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: 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) + 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, + 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 + + +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}") + + # 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) + + # 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 + + 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=extra_request_body, + timestamp=request.timestamp, + ) + + tasks.append( + asyncio.create_task( + limited_request_func(request_func_input=request_func_input, pbar=pbar) + ) + ) + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + + # Stop profiler + if profile: + 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=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)) + print( + "{:<40} {:<10}".format("Total input vision tokens:", metrics.total_input_vision) + ) + 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( + "Peak output token throughput (tok/s):", metrics.max_output_tokens_per_s + ) + ) + print( + "{:<40} {:<10}".format( + "Peak concurrent requests:", metrics.max_concurrent_requests + ) + ) + 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("{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, + "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, "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 + + 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 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 + ) + + # 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 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.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), + ) + ) + + +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}") + + +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( + "--dataset-name", + type=str, + default="sharegpt", + choices=[ + "sharegpt", + "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( + "--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( + "--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( + "--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"], + ) + 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( + "--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", + ) + 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." + ) + args = parser.parse_args() + run_benchmark(args) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/check_env.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/check_env.py new file mode 100644 index 00000000..18fa94af --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/check_env.py @@ -0,0 +1,433 @@ +"""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_npu + + +def is_cuda_v2(): + return torch.version.cuda is not None + + +# List of packages to check versions +PACKAGE_LIST = [ + "sglang", + "sgl_kernel", + "flashinfer_python", + "flashinfer_cubin", + "flashinfer_jit_cache", + "triton", + "transformers", + "torchao", + "numpy", + "aiohttp", + "fastapi", + "hf_transfer", + "huggingface_hub", + "interegular", + "modelscope", + "orjson", + "outlines", + "packaging", + "psutil", + "pydantic", + "python-multipart", + "pyzmq", + "torchao", + "uvicorn", + "uvloop", + "vllm", + "xgrammar", + "openai", + "tiktoken", + "anthropic", + "litellm", + "decord2", +] + + +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 {} + + +if __name__ == "__main__": + if is_cuda_v2(): + env = GPUEnv() + elif is_hip(): + env = HIPEnv() + elif is_npu(): + env = NPUEnv() + env.check_env() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/generate.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/generate.py new file mode 100644 index 00000000..894a1175 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/generate.py @@ -0,0 +1,33 @@ +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 = parser.parse_args(extra_argv) + generate_cmd(parsed_args) + else: + raise Exception( + f"Generate subcommand is not yet supported for model: {model_path}" + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/main.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/main.py new file mode 100644 index 00000000..e8d3b755 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/main.py @@ -0,0 +1,26 @@ +import argparse + +from sglang.cli.generate import generate +from sglang.cli.serve import serve + + +def main(): + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand", required=True) + + serve_parser = subparsers.add_parser( + "serve", + help="Launch the SGLang server.", + add_help=False, # Defer help to the specific parser + ) + serve_parser.set_defaults(func=serve) + + generate_parser = subparsers.add_parser( + "generate", + help="Run inference on a multimodal model.", + add_help=False, # Defer help to the specific parser + ) + generate_parser.set_defaults(func=generate) + + args, extra_argv = parser.parse_known_args() + args.func(args, extra_argv) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/serve.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/serve.py new file mode 100644 index 00000000..855d6335 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/serve.py @@ -0,0 +1,75 @@ +# 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 + +logger = logging.getLogger(__name__) + + +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" + ) + print( + "This command can launch either a standard language model server or a diffusion model server." + ) + print("The server type is determined by the model path.\n") + print("For specific arguments, please provide a model_path.") + 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 ---") + from sglang.multimodal_gen.runtime.entrypoints.cli.serve import ( + add_multimodal_gen_serve_args, + ) + + parser = argparse.ArgumentParser(description="SGLang Diffusion Model Serving") + add_multimodal_gen_serve_args(parser) + parser.print_help() + return + + model_path = get_model_path(extra_argv) + try: + is_diffusion_model = get_is_diffusion_model(model_path) + if is_diffusion_model: + logger.info("Diffusion model detected") + + 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(extra_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 + + # Add a dummy argument for the program name, expected by prepare_server_args + # as it typically processes sys.argv + server_args = prepare_server_args(extra_argv) + + run_server(server_args) + finally: + kill_process_tree(os.getpid(), include_parent=False) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/utils.py new file mode 100644 index 00000000..57068fc4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/cli/utils.py @@ -0,0 +1,152 @@ +import hashlib +import json +import logging +import os +import tempfile +from typing import Optional + +import filelock +from huggingface_hub import hf_hub_download + +logger = logging.getLogger(__name__) + +temp_dir = tempfile.gettempdir() + + +def _get_lock(model_name_or_path: str, cache_dir: Optional[str] = None): + lock_dir = cache_dir or temp_dir + 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() + lock_file_name = hash_name + model_name + ".lock" + lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666) + return lock + + +# Copied and adapted from hf_diffusers_utils.py +def _maybe_download_model( + model_name_or_path: str, local_dir: str | None = None, download: bool = True +) -> str: + """ + Resolve a model path. If it's a local directory, return it. + If it's a Hugging Face Hub ID, download only the config file + (`model_index.json` or `config.json`) and return its directory. + + Args: + model_name_or_path: Local path or Hugging Face Hub model ID + local_dir: Local directory to save the downloaded file (if any) + download: Whether to download from Hugging Face Hub when needed + + Returns: + Local directory path that contains the downloaded config file, or the original local directory. + """ + + if os.path.exists(model_name_or_path): + logger.info("Model already exists locally") + return model_name_or_path + + if not download: + return model_name_or_path + + with _get_lock(model_name_or_path): + # Try `model_index.json` first (diffusers models) + try: + logger.info( + "Downloading model_index.json from HF Hub for %s...", + model_name_or_path, + ) + file_path = hf_hub_download( + repo_id=model_name_or_path, + filename="model_index.json", + local_dir=local_dir, + ) + logger.info("Downloaded to %s", file_path) + return os.path.dirname(file_path) + except Exception as e_index: + logger.debug("model_index.json not found or failed: %s", e_index) + + # Fallback to `config.json` + try: + logger.info( + "Downloading config.json from HF Hub for %s...", model_name_or_path + ) + file_path = hf_hub_download( + repo_id=model_name_or_path, + filename="config.json", + local_dir=local_dir, + ) + logger.info("Downloaded to %s", file_path) + return os.path.dirname(file_path) + except Exception as e_config: + raise ValueError( + ( + "Could not find model locally at %s and failed to download " + "model_index.json/config.json from HF Hub: %s" + ) + % (model_name_or_path, e_config) + ) from e_config + + +# Copied and adapted from hf_diffusers_utils.py +def is_diffusers_model_path(model_path: str) -> True: + """ + Verify if the model directory contains a valid diffusers configuration. + + Args: + model_path: Path to the model directory + + Returns: + The loaded model configuration as a dictionary if the model is a diffusers model + None if the model is not a diffusers model + """ + + # Prefer model_index.json which indicates a diffusers pipeline + config_path = os.path.join(model_path, "model_index.json") + if not os.path.exists(config_path): + return False + + # Load the config + with open(config_path) as f: + config = json.load(f) + + # Verify diffusers version exists + if "_diffusers_version" not in config: + return False + return True + + +def get_is_diffusion_model(model_path: str): + model_path = _maybe_download_model(model_path) + is_diffusion_model = is_diffusers_model_path(model_path) + if is_diffusion_model: + logger.info("Diffusion model detected") + return is_diffusion_model + + +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" + "For specific arguments, please provide a model_path." + ) + else: + raise Exception( + "Error: --model-path is required. " + "Please provide the path to the model." + ) + return model_path diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/compile_deep_gemm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/compile_deep_gemm.py new file mode 100644 index 00000000..2386a46c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/compile_deep_gemm.py @@ -0,0 +1,191 @@ +""" +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 +os.environ["SGL_CHUNKED_PREFIX_CACHE_THRESHOLD"] = "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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/eval/llama3_eval.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/eval/llama3_eval.py new file mode 100644 index 00000000..253cdf27 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/eval/llama3_eval.py @@ -0,0 +1,315 @@ +# 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") == 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/eval/loogle_eval.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/eval/loogle_eval.py new file mode 100644 index 00000000..895362cd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/eval/loogle_eval.py @@ -0,0 +1,164 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/global_config.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/global_config.py new file mode 100644 index 00000000..fcd65b5e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/global_config.py @@ -0,0 +1,29 @@ +"""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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/.clang-format b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/.clang-format new file mode 100644 index 00000000..75fe1387 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/.clang-format @@ -0,0 +1,19 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/csrc/cuda_wait_value.cuh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/csrc/cuda_wait_value.cuh new file mode 100644 index 00000000..5874ce1e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/csrc/cuda_wait_value.cuh @@ -0,0 +1,38 @@ +#include +#include + +#include + +#include + +namespace { + +__global__ void wait_flag_kernel(const int32_t* flag, int32_t target) { + const volatile int32_t* vflag = (volatile const int32_t*)flag; + + while (*vflag != target) { +#if __CUDA_ARCH__ >= 700 + __nanosleep(100); +#else + // Note: This falls back to an inefficient busy-wait on pre-Volta architectures. +#endif + } +} + +auto stream_wait_value(const tvm::ffi::TensorView flag, std::int32_t value) -> void { + using namespace host; + + auto length = SymbolicSize{"length"}; + TensorMatcher({length}).with_dtype().with_device().verify(flag); + RuntimeCheck(length.unwrap() >= 1, "wait_flag expects a non-empty tensor."); + + auto* ptr = static_cast(flag.data_ptr()); + const auto stream = LaunchKernel::resolve_device(flag.device()); + + constexpr int blocks = 1; + constexpr int threads = 1; + wait_flag_kernel<<>>(ptr, value); + RuntimeDeviceCheck(cudaGetLastError()); +} + +} // namespace diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/csrc/hicache.cuh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/csrc/hicache.cuh new file mode 100644 index 00000000..e52ecbd3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/csrc/hicache.cuh @@ -0,0 +1,264 @@ +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace { + +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; + std::size_t length; + std::size_t kv_cache_src_stride; + std::size_t kv_cache_dst_stride; + std::size_t num_layers = 0; // only used in all_layer transfer +}; + +template < + std::integral T, + std::size_t kElementSize, + std::size_t kUnroll, + std::size_t kBlockQuota, + std::size_t kNumThreads, + std::size_t kMaxOccupancy> +__global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void hicache_transfer_per_layer( + const __grid_constant__ HicacheKernelParams params) { + // each warp acts as a worker + using namespace device; + static_assert(kNumThreads % kWarpThreads == 0); + static_assert(kWarpThreads % kUnroll == 0); + + constexpr auto kWarpThreads = device::kWarpThreads / kUnroll; + constexpr auto kWarpsPerBlock = kNumThreads / kWarpThreads; + constexpr auto kWorkers = kWarpsPerBlock * kBlockQuota; + + const auto& [ + k_cache_dst, v_cache_dst, indices_dst, // dst + k_cache_src, v_cache_src, indices_src, // src + length, kv_cache_src_stride, kv_cache_dst_stride, _ // metadata + ] = params; + const auto warp_id = blockIdx.x * kWarpsPerBlock + threadIdx.x / kWarpThreads; + + // force to transfer 128 bytes per iteration + // since the PCIe transaction size is 128 bytes aligned + constexpr auto kGranularity = 128 / kWarpThreads; + + for (auto i = warp_id; i < length; i += kWorkers) { + 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 = warp::load_vec(src_k); + const auto vec_v = warp::load_vec(src_v); + warp::store_vec(dst_k, vec_k); + warp::store_vec(dst_v, vec_v); + } +} + +template < + std::integral T, + std::size_t kElementSize, + std::size_t kUnroll, + std::size_t kBlockQuota, + std::size_t kNumThreads, + std::size_t kMaxOccupancy> +__global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void hicache_transfer_all_layer( + const __grid_constant__ HicacheKernelParams params) { + // each warp acts as a worker + using namespace device; + using src_ptr_t = std::add_pointer_t; + using dst_ptr_t = std::add_pointer_t; + + static_assert(kNumThreads % kWarpThreads == 0); + constexpr auto kWarpThreads = device::kWarpThreads / kUnroll; + constexpr auto kWarpsPerBlock = static_cast(kNumThreads) / kWarpThreads; + constexpr auto kWorkers = kWarpsPerBlock * kBlockQuota; + + const auto& [ + k_ptr_dst, v_ptr_dst, indices_dst, // dst + k_ptr_src, v_ptr_src, indices_src, // src + length, kv_cache_src_stride, kv_cache_dst_stride, num_layers // metadata + ] = params; + const auto warp_id = blockIdx.x * kWarpsPerBlock + threadIdx.x / kWarpThreads; + + // force to transfer 128 bytes per iteration + // since the PCIe transaction size is 128 bytes aligned + constexpr auto kGranularity = 128 / kWarpThreads; + + for (auto i = warp_id; i < length; i += kWorkers) { + const auto pos_src = static_cast(indices_src)[i]; + const auto pos_dst = static_cast(indices_dst)[i]; + for (std::size_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 = warp::load_vec(src_k); + const auto vec_v = warp::load_vec(src_v); + warp::store_vec(dst_k, vec_k); + warp::store_vec(dst_v, vec_v); + } + } +} + +template < + std::size_t kElementSize, + std::size_t kUnroll, + std::size_t kBlockQuota, + std::size_t kNumThreads, + std::size_t kMaxOccupancy> +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{"D"}; // cache dimension + auto N = SymbolicSize{"N"}; // src kv stride + auto M = SymbolicSize{"M"}; // dst kv stride + auto L = SymbolicSize{"L"}; // 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 = kNumThreads / (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, + .length = length, + .kv_cache_src_stride = kv_cache_src_stride, + .kv_cache_dst_stride = kv_cache_dst_stride, + }; + const auto kernel = use_int32 ? _kernel_one : _kernel_one; + LaunchKernel(num_blocks, kNumThreads, 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 std::size_t kv_src_stride, + const std::size_t kv_dst_stride) { + using namespace host; + + auto N = SymbolicSize{"N"}; // num layers + auto L = SymbolicSize{"L"}; // 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 = kNumThreads / (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, + .length = length, + .kv_cache_src_stride = kv_src_stride, + .kv_cache_dst_stride = kv_dst_stride, + .num_layers = static_cast(N.unwrap()), + }; + const auto kernel = use_int32 ? _kernel_all : _kernel_all; + LaunchKernel(num_blocks, kNumThreads, device)(kernel, params); + } +}; + +} // namespace diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/cuda_wait_value.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/cuda_wait_value.py new file mode 100644 index 00000000..27665dbb --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/cuda_wait_value.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import load_jit + +if TYPE_CHECKING: + import torch + from tvm_ffi.module import Module + + +@lru_cache(maxsize=1) +def _jit_stream_wait_value_module() -> Module: + return load_jit( + "cuda_wait_value", + cuda_files=["cuda_wait_value.cuh"], + cuda_wrappers=[("stream_wait_value", "stream_wait_value")], + ) + + +def stream_wait_value(flag: torch.Tensor, value: int) -> None: + module = _jit_stream_wait_value_module() + module.stream_wait_value(flag, value) + + +class Event: + def __init__(self) -> None: + self.flag = torch.zeros(1, dtype=torch.int32, device="cuda") + + def record(self, value: int = 1) -> None: + self.flag[0] = value + + def wait(self, value: int = 1) -> None: + stream_wait_value(self.flag, value) + + +def test_wait_before_record(event: Event | torch.cuda.Event): + stream_a = torch.cuda.Stream() + stream_b = torch.cuda.Stream() + + with torch.cuda.stream(stream_a): + event.wait() + + stream_a.synchronize() + + with torch.cuda.stream(stream_b): + event.record() + + +def main(): + import threading + import time + + block_thead = threading.Thread( + target=test_wait_before_record, args=(Event(),), daemon=True + ) + block_thead.start() + + non_block_thread = threading.Thread( + target=test_wait_before_record, args=(torch.cuda.Event(),) + ) + non_block_thread.start() + + print("Checking if custom Event blocks the stream...", flush=True) + for _ in range(5): + print(f"{block_thead.is_alive()=}, {non_block_thread.is_alive()=}", flush=True) + time.sleep(1) + + assert block_thead.is_alive(), "Custom Event did not block as expected" + assert not non_block_thread.is_alive(), "torch.cuda.Event should not block" + print("=" * 40) + print("Test completed successfully.") + + +if __name__ == "__main__": + main() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/hicache.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/hicache.py new file mode 100644 index 00000000..1d015fe0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/hicache.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import logging +from functools import lru_cache +from typing import TYPE_CHECKING + +from sglang.jit_kernel.utils import load_jit, make_cpp_args + +if TYPE_CHECKING: + import torch + from tvm_ffi.module import Module + +DEFAULT_BLOCK_QUOTA = 2 + + +@lru_cache(maxsize=None) +def _jit_hicache_module(*, element_size: int, unroll: int, block_quota: int) -> Module: + num_threads, occupancy = 1024, 1 + args = make_cpp_args( + element_size, + unroll, + block_quota, + num_threads, + occupancy, + ) + 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: + 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 = logging.getLogger(__name__) + 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, + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/tensor.h b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/tensor.h new file mode 100644 index 00000000..8208149e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/tensor.h @@ -0,0 +1,487 @@ +#pragma once +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace host { + +namespace stdr = std::ranges; +namespace stdv = std::views; + +namespace details { + +struct SizeRef; +struct DTypeRef; +struct DeviceRef; + +template +struct dtype_trait {}; + +template +struct dtype_trait { + inline static constexpr auto value = DLDataType{ + .code = std::is_signed_v ? DLDataTypeCode::kDLInt : DLDataTypeCode::kDLUInt, + .bits = static_cast(sizeof(T) * 8), + .lanes = 1}; +}; + +template +struct dtype_trait { + inline static constexpr auto value = + DLDataType{.code = DLDataTypeCode::kDLFloat, .bits = static_cast(sizeof(T) * 8), .lanes = 1}; +}; + +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); + +template +inline constexpr auto kDTypeList = std::array{dtype_trait::value...}; + +template +inline constexpr auto kDeviceList = std::array{ + DLDevice{.device_type = static_cast(Codes), .device_id = kAnyDeviceID}...}; + +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); + host::RuntimeCheck(entry < mapping.size()); + const auto name = mapping[entry]; + host::RuntimeCheck(!name.empty(), "Unknown device: ", int(device.device_type)); + os << name; + if (device.device_id != kAnyDeviceID) 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 : stdv::iota(std::size_t{0}, span.data.size())) { + if (i > 0) { + os << ", "; + } + os << span.data[i]; + } + os << "]"; + return os; +} + +} // namespace details + +struct SymbolicSize { + public: + SymbolicSize(std::string_view annotation = {}) : m_value(details::kNullSize), m_annotation(annotation) {} + + auto get_name() const -> std::string_view { + return m_annotation; + } + auto set_value(int64_t value) -> void { + host::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() const -> int64_t { + host::RuntimeCheck(this->has_value(), "Size value is not set"); + return m_value; + } + + SymbolicSize(const SymbolicSize&) = delete; + SymbolicSize& operator=(const SymbolicSize&) = delete; + + auto verify(int64_t dim) -> void { + if (this->has_value()) { + host::RuntimeCheck(m_value == dim, "Size mismatch: expected ", m_value, " but got ", dim); + } else { + this->set_value(dim); + } + } + + private: + 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; +} + +struct SymbolicDType { + public: + SymbolicDType() : m_value({details::kNullDType, 0, 0}) {} + + auto set_value(DLDataType value) -> void { + host::RuntimeCheck(!this->has_value(), "Dtype value already set"); + host::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() const -> DLDataType { + host::RuntimeCheck(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()) { + host::RuntimeCheck(m_value == dtype, "DType mismatch: expected ", m_value, " but got ", dtype); + } else { + this->set_value(dtype); + } + } + + 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; +}; + +struct SymbolicDevice { + public: + SymbolicDevice() : m_value({details::kNullDevice, details::kAnyDeviceID}) {} + + auto set_value(DLDevice value) -> void { + host::RuntimeCheck(!this->has_value(), "Device value already set"); + host::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() const -> DLDevice { + host::RuntimeCheck(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()) { + host::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 + } + } + + auto value_or_name(std::size_t dim) const -> std::string { + if (const auto value = (**this).get_value()) { + return std::to_string(*value); + } else { + const auto annotation = (**this).get_name(); + if (annotation.empty()) { + return "dim#" + std::to_string(dim); + } else { + return static_cast(annotation); + } + } + } +}; + +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 + +struct TensorMatcher { + private: + using SizeRef = details::SizeRef; + using DTypeRef = details::DTypeRef; + using DeviceRef = details::DeviceRef; + using Loc_t = std::source_location; + + 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 + host::RuntimeCheck(m_strides.size() == 0, "Strides already specified"); + host::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); + 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); + 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, Loc_t loc = Loc_t::current()) const&& -> const TensorMatcher&& { + try { + this->m_verify_impl(view); + } catch (PanicError& e) { + auto oss = std::ostringstream{}; + oss << "Tensor match failed for " << this->debug_str() << " at " << loc.file_name() << ":" << loc.line() + << "\n- Root cause: " << e.detail(); + throw PanicError(std::move(oss).str()); + } + return std::move(*this); + } + + auto debug_str() const -> std::string { + auto oss = std::ostringstream{}; + oss << "Tensor<"; + std::size_t dim = 0; + for (const auto& size_ref : m_shape) { + if (dim > 0) { + oss << ", "; + } + oss << size_ref.value_or_name(dim++); + } + oss << ">"; + if (m_strides.size() > 0) { + oss << " [strides=<"; + dim = 0; + for (const auto& stride_ref : m_strides) { + if (dim > 0) { + oss << ", "; + } + oss << stride_ref.value_or_name(dim++); + } + oss << ">]"; + } + return std::move(oss).str(); + } + + private: + auto m_verify_impl(tvm::ffi::TensorView view) const -> void { + const auto dim = static_cast(view.dim()); + host::RuntimeCheck(dim == m_shape.size(), "Tensor dimension mismatch: expected ", m_shape.size(), " but got ", dim); + for (const auto i : stdv::iota(std::size_t{0}, dim)) { + m_shape[i]->verify(view.size(i)); + } + if (this->m_has_strides()) { + for (const auto i : stdv::iota(std::size_t{0}, dim)) { + m_strides[i]->verify(view.stride(i)); + } + } else { + host::RuntimeCheck(view.is_contiguous(), "Tensor is not contiguous as expected"); + } + // since we may use the same matcher to verify again, we will force to check + m_dtype->verify(view.dtype()); + m_device->verify(view.device()); + } + + auto m_init_dtype() -> void { + host::RuntimeCheck(!m_has_dtype, "DType already specified"); + m_has_dtype = true; + } + auto m_init_device() -> void { + host::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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh new file mode 100644 index 00000000..cf03d8c0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -0,0 +1,101 @@ +#pragma once + +#include + +#include +#include + +#include +#include +#include +#include + +namespace device { + +inline constexpr auto kWarpThreads = 32u; + +namespace pointer { + +// we only allow void * pointer arithmetic for safety + +template +__always_inline __device__ auto offset(T* ptr, U... offset) -> void* { + static_assert(std::is_same_v, "Pointer arithmetic is only allowed for void* pointers"); + return static_cast(ptr) + (... + offset); +} + +template +__always_inline __device__ auto offset(const T* ptr, U... offset) -> const void* { + static_assert(std::is_same_v, "Pointer arithmetic is only allowed for void* pointers"); + return static_cast(ptr) + (... + offset); +} + +} // namespace pointer + +} // namespace device + +namespace host { + +inline auto +RuntimeDeviceCheck(::cudaError_t error, std::source_location location = std::source_location::current()) -> void { + if (error != ::cudaSuccess) { + [[unlikely]]; + ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error)); + } +} + +inline auto RuntimeCudaCheck(std::source_location location = std::source_location::current()) -> void { + return RuntimeDeviceCheck(::cudaGetLastError(), location); +} + +template +inline void set_smem_once(std::size_t smem_size) { + static const auto last_smem_size = [&] { + RuntimeDeviceCheck(::cudaFuncSetAttribute(F, ::cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + return smem_size; + }(); + RuntimeCheck( + smem_size <= last_smem_size, + "Dynamic shared memory size exceeds the previously set maximum size: ", + last_smem_size, + " bytes"); +} + +struct LaunchKernel { + public: + explicit LaunchKernel( + dim3 grid_dim, dim3 block_dim, DLDevice device, std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_config(s_make_config(grid_dim, block_dim, resolve_device(device), dynamic_shared_mem_bytes)) {} + + explicit LaunchKernel( + dim3 grid_dim, dim3 block_dim, cudaStream_t stream, std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_config(s_make_config(grid_dim, block_dim, stream, dynamic_shared_mem_bytes)) {} + + static auto resolve_device(DLDevice device) -> cudaStream_t { + return static_cast(::TVMFFIEnvGetStream(device.device_type, device.device_id)); + } + + LaunchKernel(const LaunchKernel&) = delete; + LaunchKernel& operator=(const LaunchKernel&) = delete; + + template + auto operator()(T&& kernel, Args&&... args) const -> void { + host::RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward(args)...)); + } + + private: + static auto + s_make_config(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; + /// TODO: We can add a queue to store the attributes if needed in the future. +}; + +} // namespace host diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/utils.h b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/utils.h new file mode 100644 index 00000000..fd9723df --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/utils.h @@ -0,0 +1,88 @@ +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace host { + +struct PanicError : public std::runtime_error { + public: + // copy and move constructors + explicit PanicError(std::string msg) : runtime_error(msg), m_message(std::move(msg)) {} + auto detail() const -> std::string_view { + const auto sv = std::string_view{m_message}; + const auto pos = sv.find(": "); + return pos == std::string_view::npos ? sv : sv.substr(pos + 2); + } + + private: + std::string m_message; +}; + +template +[[noreturn]] +inline auto panic(std::source_location 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()); +} + +template +struct RuntimeCheck { + using Loc_t = std::source_location; + template + explicit RuntimeCheck(Cond&& condition, Args&&... args, Loc_t location = Loc_t::current()) { + if (!condition) { + [[unlikely]]; + ::host::panic(location, std::forward(args)...); + } + } +}; + +template +explicit RuntimeCheck(Cond&&, Args&&...) -> RuntimeCheck; + +template +inline constexpr auto div_ceil(T a, U b) { + return (a + b - 1) / b; +} + +template +inline constexpr auto div_ceil(T a, U b) { + return (a + b - 1) / b; +} + +inline auto dtype_bytes(DLDataType dtype) -> std::size_t { + return static_cast(dtype.bits / 8); +} + +namespace pointer { + +// we only allow void * pointer arithmetic for safety + +template +inline auto offset(T* ptr, U... offset) -> void* { + static_assert(std::is_same_v, "Pointer arithmetic is only allowed for void* pointers"); + return static_cast(ptr) + (... + offset); +} + +template +inline auto offset(const T* ptr, U... offset) -> const void* { + static_assert(std::is_same_v, "Pointer arithmetic is only allowed for void* pointers"); + return static_cast(ptr) + (... + offset); +} + +} // namespace pointer + +} // namespace host diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh new file mode 100644 index 00000000..904531f3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh @@ -0,0 +1,145 @@ +#pragma once +#include + +#include +#include +#include + +namespace device::warp { + +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"); + } +} + +inline constexpr auto default_unit_size(std::size_t x) -> std::size_t { + if (x % (16 * kWarpThreads) == 0) return 16; + if (x % (8 * kWarpThreads) == 0) return 8; + if (x % (4 * kWarpThreads) == 0) return 4; + return 0; // trigger static assert in _get_mem_package +} + +template +using mem_package_t = decltype(get_mem_package()); + +template +struct storage_vec { + T data[N]; +}; + +__always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { + uint32_t tmp; + asm volatile("ld.global.cs.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); + return uint1{tmp}; +} + +__always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { + uint32_t tmp0, tmp1; + asm volatile("ld.global.cs.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); + return uint2{tmp0, tmp1}; +} + +__always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 { + uint32_t tmp0, tmp1, tmp2, tmp3; + asm volatile("ld.global.cs.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src)); + return uint4{tmp0, tmp1, tmp2, tmp3}; +} + +__always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) { + uint32_t tmp = value.x; + asm volatile("st.global.cs.b32 [%0],%1;" ::"l"(dst), "r"(tmp)); +} + +__always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) { + uint32_t tmp0 = value.x; + uint32_t tmp1 = value.y; + asm volatile("st.global.cs.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1)); +} + +__always_inline __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.cs.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3)); +} + +} // namespace details + +template < + std::size_t kBytes, + std::size_t kUnit = details::default_unit_size(kBytes), + std::size_t kThreads = ::device::kWarpThreads> +__always_inline __device__ void copy(void* __restrict__ dst, const void* __restrict__ src) { + using Package = details::mem_package_t; + constexpr auto kBytesPerLoop = sizeof(Package) * kThreads; + constexpr auto kLoopCount = kBytes / kBytesPerLoop; + static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes"); + + const auto dst_packed = static_cast(dst); + const auto src_packed = static_cast(src); + const auto lane_id = threadIdx.x % kThreads; + +#pragma unroll kLoopCount + for (std::size_t i = 0; i < kLoopCount; ++i) { + const auto j = i * kThreads + lane_id; + dst_packed[j] = src_packed[j]; + } +} + +template < + std::size_t kBytes, + std::size_t kUnit = details::default_unit_size(kBytes), + std::size_t kThreads = ::device::kWarpThreads> +__always_inline __device__ auto load_vec(const void* __restrict__ src) { + using Package = details::mem_package_t; + constexpr auto kBytesPerLoop = sizeof(Package) * kThreads; + constexpr auto kLoopCount = kBytes / kBytesPerLoop; + static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes"); + + const auto src_packed = static_cast(src); + const auto lane_id = threadIdx.x % kThreads; + details::storage_vec vec; + +#pragma unroll kLoopCount + for (std::size_t i = 0; i < kLoopCount; ++i) { + const auto j = i * kThreads + lane_id; + vec.data[i] = details::load_nc(src_packed + j); + } + + return vec; +} + +template < + std::size_t kBytes, + std::size_t kUnit = details::default_unit_size(kBytes), + std::size_t kThreads = ::device::kWarpThreads, + typename Tp> +__always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) { + using Package = details::mem_package_t; + constexpr auto kBytesPerLoop = sizeof(Package) * kThreads; + constexpr auto kLoopCount = kBytes / kBytesPerLoop; + static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes"); + static_assert(std::is_same_v>); + + const auto dst_packed = static_cast(dst); + const auto lane_id = threadIdx.x % kThreads; + +#pragma unroll kLoopCount + for (std::size_t i = 0; i < kLoopCount; ++i) { + const auto j = i * kThreads + lane_id; + details::store_nc(dst_packed + j, vec.data[i]); + } +} + +} // namespace device::warp diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/utils.py new file mode 100644 index 00000000..6462cf41 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/jit_kernel/utils.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import pathlib +from functools import lru_cache +from typing import TYPE_CHECKING, List, Tuple, TypeAlias, Union + +if TYPE_CHECKING: + from tvm_ffi import Module + + +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}));" + + +@lru_cache() +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 sgl-kernel/jit 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_LDFLAGS = [] +CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] + + +class CPPArgList(list[str]): + def __str__(self) -> str: + return ", ".join(self) + + +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) + 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: + 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] + + 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=DEFAULT_CUDA_CFLAGS + extra_cuda_cflags, + extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, + build_directory=build_directory, + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/api.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/api.py new file mode 100644 index 00000000..745c656e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/api.py @@ -0,0 +1,292 @@ +"""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)]) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/anthropic.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/anthropic.py new file mode 100644 index 00000000..4918a170 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/anthropic.py @@ -0,0 +1,73 @@ +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, {} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/base_backend.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/base_backend.py new file mode 100644 index 00000000..62dd5041 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/base_backend.py @@ -0,0 +1,82 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/litellm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/litellm.py new file mode 100644 index 00000000..5803b543 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/litellm.py @@ -0,0 +1,90 @@ +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, {} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/openai.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/openai.py new file mode 100644 index 00000000..a2d006bb --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/openai.py @@ -0,0 +1,475 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/runtime_endpoint.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/runtime_endpoint.py new file mode 100644 index 00000000..1573ca68 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/runtime_endpoint.py @@ -0,0 +1,527 @@ +import atexit +import json +import multiprocessing +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", + *args, + **kwargs, + ): + """See the arguments in server_args.py::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 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 + pipe_reader, pipe_writer = multiprocessing.Pipe(duplex=False) + + ctx = multiprocessing.get_context("spawn") + proc = ctx.Process( + target=launch_server, + args=(self.server_args, pipe_writer), + ) + proc.start() + pipe_writer.close() + self.pid = proc.pid + + # Before python program terminates, call shutdown implicitly. Therefore, users don't have to explicitly call .shutdown() + atexit.register(self.shutdown) + + # TODO: remove this pipe_writer mechanism and use `/health_generate` instead. + try: + init_state = pipe_reader.recv() + except EOFError: + init_state = "" + + if init_state != "ready": + self.shutdown() + raise RuntimeError( + "Initialization failed. Please see the error messages above." + ) + + 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/vertexai.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/vertexai.py new file mode 100644 index 00000000..3d51fb13 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/backend/vertexai.py @@ -0,0 +1,148 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/chat_template.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/chat_template.py new file mode 100644 index 00000000..212d07e0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/chat_template.py @@ -0,0 +1,668 @@ +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"), + }, + 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)) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/choices.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/choices.py new file mode 100644 index 00000000..e52c6b36 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/choices.py @@ -0,0 +1,164 @@ +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/interpreter.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/interpreter.py new file mode 100644 index 00000000..0b59e91b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/interpreter.py @@ -0,0 +1,1061 @@ +"""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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/ir.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/ir.py new file mode 100644 index 00000000..43da723b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/ir.py @@ -0,0 +1,643 @@ +"""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})" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/tracer.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/tracer.py new file mode 100644 index 00000000..0a2a744f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/lang/tracer.py @@ -0,0 +1,279 @@ +"""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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/launch_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/launch_server.py new file mode 100644 index 00000000..9e3e82a7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/launch_server.py @@ -0,0 +1,29 @@ +"""Launch the inference server.""" + +import asyncio +import os +import sys + +from sglang.srt.server_args import prepare_server_args +from sglang.srt.utils import kill_process_tree + + +def run_server(server_args): + """Run the server based on server_args.grpc_mode.""" + if server_args.grpc_mode: + from sglang.srt.entrypoints.grpc_server import serve_grpc + + asyncio.run(serve_grpc(server_args)) + else: + from sglang.srt.entrypoints.http_server import launch_server + + launch_server(server_args) + + +if __name__ == "__main__": + server_args = prepare_server_args(sys.argv[1:]) + + try: + run_server(server_args) + finally: + kill_process_tree(os.getpid(), include_parent=False) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/README.md new file mode 100644 index 00000000..68c9fb4b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/README.md @@ -0,0 +1,76 @@ +
+ +
+ +**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 + - 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 + - Diverse hardware support: H100, H200, A100, B200, 4090 + +## Getting Started + +```bash +uv pip install 'sglang[diffusion]' --prerelease=allow +``` + +For more installation methods (e.g. pypi, uv, docker), check [install.md](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen/docs/install.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 + ) + + # Provide a prompt for your video + prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest." + + # Generate the video + video = generator.generate( + prompt, + 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 +``` + +For more usage examples (e.g. OpenAI compatible API, server mode), check [cli.md](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen/docs/cli.md). + +## Contributing + +All contributions are welcome. The contribution guide is available [here](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen/docs/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 FastVide 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. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/__init__.py new file mode 100644 index 00000000..75182221 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/__init__.py @@ -0,0 +1,6 @@ +# 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"] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/benchmarks/compare_perf.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/benchmarks/compare_perf.py new file mode 100644 index 00000000..2dfb087c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/benchmarks/compare_perf.py @@ -0,0 +1,216 @@ +import argparse +import json +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 compare_benchmarks( + baseline_path: str, new_path: str, output_format: str = "markdown" +): + """ + Compares two benchmark JSON files and prints a report. + """ + try: + base_data = _load_benchmark_file(baseline_path) + new_data = _load_benchmark_file(new_path) + except Exception as e: + print(f"Error loading benchmark files: {e}") + return + + base_e2e = base_data.get("total_duration_ms", 0) + new_e2e = new_data.get("total_duration_ms", 0) + + diff_ms, diff_pct = calculate_diff(base_e2e, new_e2e) + + if diff_pct < -2.0: + status = "✅" + elif diff_pct > 2.0: + status = "❌" + else: + status = "" + + # --- Stage Breakdown --- + base_durations, base_order, base_counts = consolidate_steps( + base_data.get("steps", []) + ) + new_durations, new_order, new_counts = consolidate_steps(new_data.get("steps", [])) + + # Merge orders: Start with New order (execution order), append any missing from Base + combined_order = list(new_order) + for name in base_order: + if name not in combined_order: + combined_order.append(name) + + stage_rows = [] + 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) + + # Format count string if aggregated + count_str = "" + if stage == "Denoising Loop": + count_str = ( + f" ({n_count} steps)" + if n_count == b_count + else f" ({b_count}->{n_count} steps)" + ) + + # filter noise: show if diff is > 0.5ms OR if it's a major stage (like Denoising Loop) + # always show Denoising Loop or stages with significant duration/diff + stage_rows.append((stage + count_str, b_val, n_val, s_diff, s_pct)) + + if output_format == "markdown": + print("### Performance Comparison Report\n") + + # Summary Table + 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") + + # Detailed Breakdown + print("#### 2. Stage Breakdown") + print( + "| Stage Name | Baseline (ms) | New (ms) | Diff (ms) | Diff (%) | Status |" + ) + print("| :--- | :--- | :--- | :--- | :--- | :--- |") + for name, b, n, d, p in stage_rows: + name_str = name + status_emoji = get_perf_status_emoji(b, n) + print( + f"| {name_str} | {b:.2f} | {n:.2f} | {d:+.2f} | {p:+.1f}% | {status_emoji} |" + ) + print("\n") + + # Metadata + print("
") + print("Metadata\n") + print(f"- Baseline Commit: `{base_data.get('commit_hash', 'N/A')}`") + print(f"- New Commit: `{new_data.get('commit_hash', 'N/A')}`") + print(f"- Timestamp: {datetime.now().isoformat()}") + print("
") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Compare two sglang-diffusion performance JSON files." + ) + parser.add_argument("baseline", help="Path to the baseline JSON file") + parser.add_argument("new", help="Path to the new JSON file") + args = parser.parse_args() + + compare_benchmarks(args.baseline, args.new) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/__init__.py new file mode 100644 index 00000000..dfff5f2c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/__init__.py @@ -0,0 +1,3 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# Configs for pipelines, and pipeline modules (in models folder) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/backend/vmoba/wan_1.3B_77_448_832.json b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/backend/vmoba/wan_1.3B_77_448_832.json new file mode 100644 index 00000000..1e55b5f2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/backend/vmoba/wan_1.3B_77_448_832.json @@ -0,0 +1,16 @@ +{ + "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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/backend/vmoba/wan_1.3B_77_480_832.json b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/backend/vmoba/wan_1.3B_77_480_832.json new file mode 100644 index 00000000..ddf66f48 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/backend/vmoba/wan_1.3B_77_480_832.json @@ -0,0 +1,16 @@ +{ + "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 +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/__init__.py new file mode 100644 index 00000000..62c0aadf --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/__init__.py @@ -0,0 +1,8 @@ +# 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"] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/base.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/base.py new file mode 100644 index 00000000..6de428ad --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/base.py @@ -0,0 +1,105 @@ +# 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 + valid_fields = {f.name for f in fields(arch_config)} + + for key, value in source_model_dict.items(): + setattr(arch_config, key, value) + # else: + # raise AttributeError( + # f"{type(arch_config).__name__} has no field '{key}'" + # ) + + 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__() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/__init__.py new file mode 100644 index 00000000..67e6d97b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/__init__.py @@ -0,0 +1,7 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig +from sglang.multimodal_gen.configs.models.dits.stepvideo import StepVideoConfig +from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig + +__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig"] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/base.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/base.py new file mode 100644 index 00000000..77fecfb4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/base.py @@ -0,0 +1,76 @@ +# 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.TORCH_SDPA, + AttentionBackendEnum.VIDEO_SPARSE_ATTN, + AttentionBackendEnum.VMOBA_ATTN, + AttentionBackendEnum.SAGE_ATTN_THREE, + } + ) + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/flux.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/flux.py new file mode 100644 index 00000000..f50c9b1e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/flux.py @@ -0,0 +1,42 @@ +# 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) + + param_names_mapping: dict = field( + default_factory=lambda: { + r"transformer\.(\w*)\.(.*)$": r"\1.\2", + } + ) + + 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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py new file mode 100644 index 00000000..1cae921f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py @@ -0,0 +1,184 @@ +# 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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py new file mode 100644 index 00000000..f9003264 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py @@ -0,0 +1,41 @@ +# 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) + + param_names_mapping: dict = field( + default_factory=lambda: { + 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 QwenImageDitConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=QwenImageArchConfig) + + prefix: str = "qwenimage" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/stepvideo.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/stepvideo.py new file mode 100644 index 00000000..1d7fe21a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/stepvideo.py @@ -0,0 +1,64 @@ +# 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_transformer_blocks(n, m): + return "transformer_blocks" in n and n.split(".")[-1].isdigit() + + +@dataclass +class StepVideoArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field( + default_factory=lambda: [is_transformer_blocks] + ) + + param_names_mapping: dict = field( + default_factory=lambda: { + # transformer block + r"^transformer_blocks\.(\d+)\.norm1\.(weight|bias)$": r"transformer_blocks.\1.norm1.norm.\2", + r"^transformer_blocks\.(\d+)\.norm2\.(weight|bias)$": r"transformer_blocks.\1.norm2.norm.\2", + r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.weight$": r"transformer_blocks.\1.ff.fc_in.weight", + r"^transformer_blocks\.(\d+)\.ff\.net\.2\.weight$": r"transformer_blocks.\1.ff.fc_out.weight", + # adanorm block + r"^adaln_single\.emb\.timestep_embedder\.linear_1\.(weight|bias)$": r"adaln_single.emb.mlp.fc_in.\1", + r"^adaln_single\.emb\.timestep_embedder\.linear_2\.(weight|bias)$": r"adaln_single.emb.mlp.fc_out.\1", + # caption projection + r"^caption_projection\.linear_1\.(weight|bias)$": r"caption_projection.fc_in.\1", + r"^caption_projection\.linear_2\.(weight|bias)$": r"caption_projection.fc_out.\1", + } + ) + + num_attention_heads: int = 48 + attention_head_dim: int = 128 + in_channels: int = 64 + out_channels: int | None = 64 + num_layers: int = 48 + dropout: float = 0.0 + patch_size: int = 1 + norm_type: str = "ada_norm_single" + norm_elementwise_affine: bool = False + norm_eps: float = 1e-6 + caption_channels: int | list[int] | tuple[int, ...] | None = field( + default_factory=lambda: [6144, 1024] + ) + attention_type: str | None = "torch" + use_additional_conditions: bool | None = False + exclude_lora_layers: list[str] = field(default_factory=lambda: []) + + def __post_init__(self): + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.out_channels = ( + self.in_channels if self.out_channels is None else self.out_channels + ) + self.num_channels_latents = self.out_channels + + +@dataclass +class StepVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=StepVideoArchConfig) + + prefix: str = "StepVideo" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py new file mode 100644 index 00000000..3a9939c4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py @@ -0,0 +1,102 @@ +# 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+)\.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 + + 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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/zimage.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/zimage.py new file mode 100644 index 00000000..7d3576fc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/dits/zimage.py @@ -0,0 +1,40 @@ +# 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 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) + + 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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/__init__.py new file mode 100644 index 00000000..70851bfa --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/__init__.py @@ -0,0 +1,25 @@ +# 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.llama import LlamaConfig +from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config + +__all__ = [ + "EncoderConfig", + "TextEncoderConfig", + "ImageEncoderConfig", + "BaseEncoderOutput", + "CLIPTextConfig", + "CLIPVisionConfig", + "LlamaConfig", + "T5Config", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/base.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/base.py new file mode 100644 index 00000000..514b1561 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/base.py @@ -0,0 +1,86 @@ +# 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, + } + ) + 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) + + +@dataclass +class ImageEncoderConfig(EncoderConfig): + arch_config: ArchConfig = field(default_factory=ImageEncoderArchConfig) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/clip.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/clip.py new file mode 100644 index 00000000..6b36fc88 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/clip.py @@ -0,0 +1,95 @@ +# 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, +) + + +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 + 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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/llama.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/llama.py new file mode 100644 index 00000000..41d98cab --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/llama.py @@ -0,0 +1,69 @@ +# 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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/qwen_image.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/qwen_image.py new file mode 100644 index 00000000..02f57ca6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/qwen_image.py @@ -0,0 +1,68 @@ +# 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 = 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 Qwen2_5VLConfig(TextEncoderConfig): + arch_config: TextEncoderArchConfig = field(default_factory=QwenImageArchConfig) + # prefix: str = "qwen_image" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/t5.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/t5.py new file mode 100644 index 00000000..3fd9b2f1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/encoders/t5.py @@ -0,0 +1,86 @@ +# 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 "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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/__init__.py new file mode 100644 index 00000000..e9b47861 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/__init__.py @@ -0,0 +1,11 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +from sglang.multimodal_gen.configs.models.vaes.hunyuanvae import HunyuanVAEConfig +from sglang.multimodal_gen.configs.models.vaes.stepvideovae import StepVideoVAEConfig +from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig + +__all__ = [ + "HunyuanVAEConfig", + "WanVAEConfig", + "StepVideoVAEConfig", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/base.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/base.py new file mode 100644 index 00000000..344a37e6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/base.py @@ -0,0 +1,156 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/flux.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/flux.py new file mode 100644 index 00000000..53bd4461 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/flux.py @@ -0,0 +1,58 @@ +# 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): + self.arch_config.vae_scale_factor = 2 ** ( + len(self.arch_config.block_out_channels) - 1 + ) + self.arch_config.spatial_compression_ratio = self.arch_config.vae_scale_factor + + +@dataclass +class Flux2VAEConfig(FluxVAEConfig): + arch_config: Flux2VAEArchConfig = field(default_factory=Flux2VAEArchConfig) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/hunyuanvae.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/hunyuanvae.py new file mode 100644 index 00000000..601b72d5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/hunyuanvae.py @@ -0,0 +1,41 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py new file mode 100644 index 00000000..d08ab5da --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py @@ -0,0 +1,53 @@ +# 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 + 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 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/stepvideovae.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/stepvideovae.py new file mode 100644 index 00000000..6794e979 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/stepvideovae.py @@ -0,0 +1,31 @@ +# 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 StepVideoVAEArchConfig(VAEArchConfig): + in_channels: int = 3 + out_channels: int = 3 + z_channels: int = 64 + num_res_blocks: int = 2 + version: int = 2 + frame_len: int = 17 + world_size: int = 1 + + spatial_compression_ratio: int = 16 + temporal_compression_ratio: int = 8 + + scaling_factor: float = 1.0 + + +@dataclass +class StepVideoVAEConfig(VAEConfig): + arch_config: VAEArchConfig = field(default_factory=StepVideoVAEArchConfig) + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False + use_temporal_scaling_frames: bool = False diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/wanvae.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/wanvae.py new file mode 100644 index 00000000..a1bd77eb --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/models/vaes/wanvae.py @@ -0,0 +1,88 @@ +# 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 + + def __post_init__(self): + self.blend_num_frames = ( + self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + ) * 2 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py new file mode 100644 index 00000000..250c0d7a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -0,0 +1,35 @@ +# 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.flux import FluxPipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( + FastHunyuanConfig, + HunyuanConfig, +) +from sglang.multimodal_gen.configs.pipeline_configs.stepvideo import StepVideoT2VConfig +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__ = [ + "HunyuanConfig", + "FastHunyuanConfig", + "FluxPipelineConfig", + "PipelineConfig", + "SlidingTileAttnConfig", + "WanT2V480PConfig", + "WanI2V480PConfig", + "WanT2V720PConfig", + "WanI2V720PConfig", + "StepVideoT2VConfig", + "SelfForcingWanT2V480PConfig", + "ZImagePipelineConfig", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/base.py new file mode 100644 index 00000000..5ffc9278 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -0,0 +1,669 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +import json +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.utils import update_config_from_args +from sglang.multimodal_gen.runtime.distributed import ( + get_sp_parallel_rank, + get_sp_world_size, + sequence_model_parallel_all_gather, +) +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, WorkloadType +# this may focus on the model's original ability +class ModelTaskType(Enum): + 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 + + def is_image_gen(self): + return self == ModelTaskType.T2I or self == ModelTaskType.I2I + + +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 + + +# config for a single pipeline +@dataclass +class PipelineConfig: + """The base configuration class for a generation pipeline.""" + + task_type: ModelTaskType = ModelTaskType.I2I + + model_path: str = "" + pipeline_config_path: str | None = None + + # 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_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,)) + ) + + # StepVideo specific parameters + pos_magic: str | None = None + neg_magic: str | None = None + timesteps_scale: bool | None = None + + # 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_image_processor_kwargs(self, batch): + 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 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): + 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 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 latents.dim() != 5: + return latents, False + time_dim = latents.shape[2] + if time_dim > 0 and time_dim % 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 + return latents, False + + 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): + return latents + + 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", + ) + + # 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-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", + ) + parser.add_argument( + f"--{prefix_with_dot}pos_magic", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}pos_magic", + default=PipelineConfig.pos_magic, + help="Positive magic prompt for sampling, used in stepvideo", + ) + parser.add_argument( + f"--{prefix_with_dot}neg_magic", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}neg_magic", + default=PipelineConfig.neg_magic, + help="Negative magic prompt for sampling, used in stepvideo", + ) + parser.add_argument( + f"--{prefix_with_dot}timesteps_scale", + type=bool, + dest=f"{prefix_with_dot.replace('-', '_')}timesteps_scale", + default=PipelineConfig.timesteps_scale, + help="Bool for applying scheduler scale in set_timesteps, used in stepvideo", + ) + + # 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") + + 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 + ) + + @classmethod + def from_kwargs( + cls, kwargs: dict[str, Any], config_cli_prefix: str = "" + ) -> "PipelineConfig": + """ + Load PipelineConfig from kwargs Dictionary. + 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") + + # 1. Get the pipeline config class from the registry + model_info = get_model_info(model_path) + + pipeline_config = model_info.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): + sp_world_size, rank_in_sp_group = get_sp_world_size(), get_sp_parallel_rank() + seq_len = latents.shape[1] + + # 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[0], pad_len, latents.shape[2]), + dtype=latents.dtype, + device=latents.device, + ) + latents = torch.cat([latents, pad], dim=1) + # Record padding length for later unpad + batch.sp_seq_pad = int(getattr(batch, "sp_seq_pad", 0)) + pad_len + + 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)) + + # If SP padding was applied, remove extra tokens before reshaping + target_tokens = (height // 2) * (width // 2) + if latents.shape[1] > target_tokens: + latents = latents[:, :target_tokens, :] + + 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 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(",")] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py new file mode 100644 index 00000000..d4bed105 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py @@ -0,0 +1,576 @@ +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.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) + + # 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) + l = torch.arange(1) # [0] - layer dimension + + # Create position IDs: (H*W, 4) + latent_ids = torch.cartesian_prod(t, h, w, l) + + # 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) + l = torch.arange(L) + + coords = torch.cartesian_prod(t, h, w, l) + 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 + + +@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.I2I + + 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]]: + target_area: int = 1024 * 1024 + if width is not None and height is not None: + if width * height > target_area: + scale = math.sqrt(target_area / (width * height)) + width = int(width * scale) + height = int(height * scale) + return width, 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 + # get image_latent_ids right after scale & shift + image_latent_ids = _prepare_image_ids([latent_condition]) + image_latent_ids = image_latent_ids.repeat(batch_size, 1, 1) + image_latent_ids = image_latent_ids.to(get_local_torch_device()) + batch.condition_image_latent_ids = image_latent_ids + + # 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 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 batch.image_latent is not None: + image_latent_ids = batch.condition_image_latent_ids + img_ids = torch.cat([img_ids, image_latent_ids], dim=1).to(device=device) + + 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) + + 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 preprocess_decoding(self, latents): + latents = _unpatchify_latents(latents) + return latents + + def get_decode_scale_and_shift(self, device, dtype, vae): + vae_arch_config = self.vae_config.arch_config + latents_bn_mean = ( + vae.bn.running_mean.view(1, -1, 1, 1).to(device=device).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 + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py new file mode 100644 index 00000000..d45dfadb --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py @@ -0,0 +1,114 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py new file mode 100644 index 00000000..035c71e0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py @@ -0,0 +1,330 @@ +# 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 +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, + 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 + + +# 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) + + # 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): + if batch.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) + + img_cos, img_sin = ( + img_freqs.real.to(dtype=dtype), + img_freqs.imag.to(dtype=dtype), + ) + txt_cos, txt_sin = ( + txt_freqs.real.to(dtype=dtype), + txt_freqs.imag.to(dtype=dtype), + ) + + return (img_cos, img_sin), (txt_cos, txt_sin) + + 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]] + + (img_cos, img_sin), (txt_cos, txt_sin) = self.get_freqs_cis( + img_shapes, txt_seq_lens, rotary_emb, device, dtype + ) + + img_cos = shard_rotary_emb_for_sp(img_cos) + img_sin = shard_rotary_emb_for_sp(img_sin) + return { + "txt_seq_lens": txt_seq_lens, + "freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)), + } + + 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]] + (img_cos, img_sin), (txt_cos, txt_sin) = 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) + ) + + 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)), + } + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/stepvideo.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/stepvideo.py new file mode 100644 index 00000000..aff18e5c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/stepvideo.py @@ -0,0 +1,36 @@ +# 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 import DiTConfig, VAEConfig +from sglang.multimodal_gen.configs.models.dits import StepVideoConfig +from sglang.multimodal_gen.configs.models.vaes import StepVideoVAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig + + +@dataclass +class StepVideoT2VConfig(PipelineConfig): + """Base configuration for StepVideo pipeline architecture.""" + + # WanConfig-specific parameters with defaults + # DiT + dit_config: DiTConfig = field(default_factory=StepVideoConfig) + # VAE + vae_config: VAEConfig = field(default_factory=StepVideoVAEConfig) + vae_tiling: bool = False + vae_sp: bool = False + + # Denoising stage + flow_shift: int = 13 + timesteps_scale: bool = False + pos_magic: str = ( + "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。" + ) + neg_magic: str = ( + "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。" + ) + + # Precision for each component + precision: str = "bf16" + vae_precision: str = "bf16" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py new file mode 100644 index 00000000..9e7f83ec --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py @@ -0,0 +1,212 @@ +# 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 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.""" + + # 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.""" + + # WanConfig-specific parameters with defaults + + # Denoising stage + flow_shift: float | None = 5.0 + + +@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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py new file mode 100644 index 00000000..b9e86b6a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py @@ -0,0 +1,142 @@ +# 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.zimage import ZImageDitConfig +from sglang.multimodal_gen.configs.models.encoders import ( + BaseEncoderOutput, + TextEncoderConfig, +) +from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ImagePipelineConfig, + ModelTaskType, +) + + +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: (TextEncoderConfig(),) + ) + + preprocess_text_funcs: tuple[Callable, ...] = field( + default_factory=lambda: (zimage_preprocess_text,) + ) + + postprocess_text_funcs: tuple[Callable, ...] = field( + default_factory=lambda: (zimage_postprocess_text,) + ) + + 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 + + def post_denoising_loop(self, latents, batch): + bs, channels, num_frames, height, width = latents.shape + 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) + + PATCH_SIZE = 2 + F_PATCH_SIZE = 1 + SEQ_MULTI_OF = 32 + + cap_ori_len = prompt_embeds.size(0) + cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF + 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) + + C = self.dit_config.num_channels_latents + F = 1 + H = height // self.vae_config.arch_config.spatial_compression_ratio + W = width // self.vae_config.arch_config.spatial_compression_ratio + + pH, pW = PATCH_SIZE, PATCH_SIZE + pF = 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) % SEQ_MULTI_OF + + 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, + ), + } diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/__init__.py new file mode 100644 index 00000000..1d2574de --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/__init__.py @@ -0,0 +1,5 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams + +__all__ = ["SamplingParams"] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/flux.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/flux.py new file mode 100644 index 00000000..657142f2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/flux.py @@ -0,0 +1,25 @@ +# 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 FluxSamplingParams(SamplingParams): + # Video parameters + # height: int = 1024 + # width: int = 1024 + num_frames: int = 1 + # Denoising stage + guidance_scale: float = 1.0 + negative_prompt: str = None + num_inference_steps: int = 50 + + def __post_init__(self): + default_sample_size = 128 + vae_scale_factor = 8 + # FIXME + # self.height = default_sample_size * vae_scale_factor + # self.width = default_sample_size * vae_scale_factor diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/hunyuan.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/hunyuan.py new file mode 100644 index 00000000..d6e33a81 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/hunyuan.py @@ -0,0 +1,37 @@ +# 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 + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/qwenimage.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/qwenimage.py new file mode 100644 index 00000000..c3270395 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/qwenimage.py @@ -0,0 +1,18 @@ +# 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): + # Video parameters + # height: int = 1024 + # width: int = 1024 + negative_prompt: str = " " + num_frames: int = 1 + # Denoising stage + guidance_scale: float = 4.0 + num_inference_steps: int = 50 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/sampling_params.py new file mode 100644 index 00000000..ef691b1a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -0,0 +1,629 @@ +# 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.path +import re +import time +import unicodedata +import uuid +from dataclasses import dataclass +from enum import Enum, auto +from typing import Any + +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 align_to + +logger = init_logger(__name__) + + +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() + + def get_default_extension(self) -> str: + if self == DataType.IMAGE: + return "jpg" + else: + return "mp4" + + +@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 | 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 = "outputs/" + output_file_name: str | None = None + + # Batch info + num_outputs_per_prompt: int = 1 + seed: int = 1024 + + # Original dimensions (before VAE scaling) + num_frames: int = 125 + num_frames_round_down: bool = ( + False # Whether to round down num_frames if it's not divisible by num_gpus + ) + height: int | None = None + width: int | None = None + # NOTE: this is temporary, we need a way to know if width or height is not provided, or do the image resize earlier + height_not_provided: bool = False + width_not_provided: bool = False + fps: int = 24 + + # Denoising parameters + num_inference_steps: int = 50 + guidance_scale: float = 1.0 + guidance_rescale: float = 0.0 + boundary_ratio: float | None = None + + # TeaCache parameters + enable_teacache: bool = False + + # Profiling + profile: bool = False + num_profiled_timesteps: int = 2 + + # 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 + + 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"] + ): + 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 + self.data_type = DataType.VIDEO if self.num_frames > 1 else DataType.IMAGE + + if self.width is None: + self.width_not_provided = True + self.width = 1280 + if self.height is None: + self.height_not_provided = True + self.height = 720 + + def check_sampling_param(self): + if self.prompt_path and not self.prompt_path.endswith(".txt"): + raise ValueError("prompt_path must be a txt file") + + def _adjust( + self, + server_args: ServerArgs, + ): + """ + final adjustment, called after merged with user params + """ + pipeline_config = server_args.pipeline_config + if not isinstance(self.prompt, str): + raise TypeError(f"`prompt` must be a string, but got {type(self.prompt)}") + + # 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}" + ) + + if pipeline_config.task_type.is_image_gen(): + # settle num_frames + logger.debug(f"Setting num_frames to 1 because this is a image-gen model") + self.num_frames = 1 + self.data_type = DataType.IMAGE + else: + # 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. + self.num_frames = server_args.pipeline_config.adjust_num_frames( + self.num_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: # stepvideo only + orig_latent_num_frames = self.num_frames // 17 * 3 + + 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: # stepvideo only + # Find the least common multiple of 3 and num_gpus + divisor = math.lcm(3, num_gpus) + # Round up to the nearest multiple of this LCM + new_latent_num_frames = ( + (new_latent_num_frames + divisor - 1) // divisor + ) * divisor + # Convert back to actual frames using the StepVideo formula + new_num_frames = new_latent_num_frames // 3 * 17 + + 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 + + self._set_output_file_name() + self.log(server_args=server_args) + + def update(self, source_dict: dict[str, Any]) -> None: + for key, value in source_dict.items(): + if hasattr(self, key): + setattr(self, key, value) + else: + logger.exception("%s has no attribute %s", type(self).__name__, key) + + self.__post_init__() + + @classmethod + def from_pretrained(cls, model_path: str, **kwargs) -> "SamplingParams": + from sglang.multimodal_gen.registry import get_model_info + + model_info = get_model_info(model_path) + sampling_params: SamplingParams = model_info.sampling_param_cls(**kwargs) + return sampling_params + + @staticmethod + def from_user_sampling_params_args(model_path: str, server_args, *args, **kwargs): + sampling_params = SamplingParams.from_pretrained(model_path) + + user_sampling_params = SamplingParams(*args, **kwargs) + # TODO: refactor + sampling_params._merge_with_user_params(user_sampling_params) + sampling_params._adjust(server_args) + + 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""" + parser.add_argument("--data-type", type=str, nargs="+", default=DataType.VIDEO) + parser.add_argument( + "--num-frames-round-down", + action="store_true", + default=SamplingParams.num_frames_round_down, + ) + parser.add_argument( + "--enable-teacache", + action="store_true", + default=SamplingParams.enable_teacache, + ) + parser.add_argument( + "--profile", + action="store_true", + default=SamplingParams.profile, + help="Enable torch profiler for denoising stage", + ) + parser.add_argument( + "--debug", + action="store_true", + default=SamplingParams.debug, + help="", + ) + parser.add_argument( + "--num-profiled-timesteps", + type=int, + default=SamplingParams.num_profiled_timesteps, + help="Number of timesteps to profile after warmup", + ) + parser.add_argument( + "--prompt", + type=str, + default=SamplingParams.prompt, + help="Text prompt for generation", + ) + parser.add_argument( + "--negative-prompt", + type=str, + default=SamplingParams.negative_prompt, + help="Negative text prompt for generation", + ) + parser.add_argument( + "--prompt-path", + type=str, + default=SamplingParams.prompt_path, + help="Path to a text file containing the prompt", + ) + parser.add_argument( + "--output-path", + type=str, + default=SamplingParams.output_path, + help="Path to save the generated image/video", + ) + parser.add_argument( + "--output-file-name", + type=str, + default=SamplingParams.output_file_name, + help="Name of the output file", + ) + parser.add_argument( + "--num-outputs-per-prompt", + type=int, + default=SamplingParams.num_outputs_per_prompt, + help="Number of outputs to generate per prompt", + ) + parser.add_argument( + "--seed", + type=int, + default=SamplingParams.seed, + help="Random seed for generation", + ) + parser.add_argument( + "--num-frames", + type=int, + default=SamplingParams.num_frames, + help="Number of frames to generate", + ) + parser.add_argument( + "--height", + type=int, + default=SamplingParams.height, + help="Height of generated output", + ) + parser.add_argument( + "--width", + type=int, + default=SamplingParams.width, + help="Width of generated output", + ) + # resolution shortcuts + parser.add_argument( + "--4k", + action="store_true", + dest="resolution_4k", + help="Set resolution to 4K (3840x2160)", + ) + parser.add_argument( + "--2k", + action="store_true", + dest="resolution_2k", + help="Set resolution to 2K (2560x1440)", + ) + parser.add_argument( + "--1080p", + action="store_true", + dest="resolution_1080p", + help="Set resolution to 1080p (1920x1080)", + ) + parser.add_argument( + "--720p", + action="store_true", + dest="resolution_720p", + help="Set resolution to 720p (1280x720)", + ) + + parser.add_argument( + "--fps", + type=int, + default=SamplingParams.fps, + help="Frames per second for saved output", + ) + parser.add_argument( + "--num-inference-steps", + type=int, + default=SamplingParams.num_inference_steps, + help="Number of denoising steps", + ) + parser.add_argument( + "--guidance-scale", + type=float, + default=SamplingParams.guidance_scale, + help="Classifier-free guidance scale", + ) + parser.add_argument( + "--guidance-rescale", + type=float, + default=SamplingParams.guidance_rescale, + help="Guidance rescale factor", + ) + parser.add_argument( + "--boundary-ratio", + type=float, + default=SamplingParams.boundary_ratio, + help="Boundary timestep ratio", + ) + parser.add_argument( + "--save-output", + action="store_true", + default=SamplingParams.save_output, + help="Whether to save the output to disk", + ) + parser.add_argument( + "--no-save-output", + action="store_false", + dest="save_output", + help="Don't save the output to disk", + ) + parser.add_argument( + "--return-frames", + action="store_true", + default=SamplingParams.return_frames, + help="Whether to return the raw frames", + ) + parser.add_argument( + "--image-path", + type=str, + default=SamplingParams.image_path, + help="Path to input image for image-to-video generation", + ) + parser.add_argument( + "--moba-config-path", + type=str, + default=None, + help="Path to a JSON file containing V-MoBA specific configurations.", + ) + parser.add_argument( + "--return-trajectory-latents", + action="store_true", + default=SamplingParams.return_trajectory_latents, + help="Whether to return the trajectory", + ) + parser.add_argument( + "--return-trajectory-decoded", + action="store_true", + default=SamplingParams.return_trajectory_decoded, + help="Whether to return the decoded trajectory", + ) + 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 + + attrs = [attr.name for attr in dataclasses.fields(cls)] + args.height_not_provided = False + args.width_not_provided = False + return {attr: getattr(args, attr) for attr in attrs} + + def output_file_path(self): + return os.path.join(self.output_path, self.output_file_name) + + def _merge_with_user_params(self, user_params): + """ + Merges parameters from a user-provided SamplingParams object. + + This method updates the current object with values from `user_params`, + but skips any fields that are explicitly defined in the current object's + subclass. This is to preserve model-specific optimal parameters. + It also skips fields that the user has not changed from the default + in `user_params`. + """ + if user_params is None: + return + + # user is not allowed to modify any param defined in the SamplingParams subclass + subclass_defined_fields = set(type(self).__annotations__.keys()) + + # Compare against current instance to avoid constructing a default instance + default_params = SamplingParams() + + for field in dataclasses.fields(user_params): + field_name = field.name + user_value = getattr(user_params, field_name) + default_value = getattr(default_params, field_name) + + # A field is considered user-modified if its value is different from + # the default, with an exception for `output_file_name` which is + # auto-generated with a random component. + is_user_modified = ( + user_value != default_value + if field_name != "output_file_name" + else user_params.output_file_path is not None + ) + if is_user_modified and field_name not in subclass_defined_fields: + if hasattr(self, field_name): + setattr(self, field_name, user_value) + + self.height_not_provided = user_params.height_not_provided + self.width_not_provided = user_params.width_not_provided + 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 + + def output_file_path(self): + return os.path.join(self.output_path, self.output_file_name) + + def log(self, server_args: ServerArgs): + # 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 + + # Log sampling parameters + debug_str = f"""Sampling params: + width: {target_width} + height: {target_height} + num_frames: {self.num_frames} + prompt: {self.prompt} + neg_prompt: {self.negative_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 CacheParams: + cache_type: str = "none" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/stepvideo.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/stepvideo.py new file mode 100644 index 00000000..4fff150b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/stepvideo.py @@ -0,0 +1,22 @@ +# 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 StepVideoT2VSamplingParams(SamplingParams): + # Video parameters + height: int = 720 + width: int = 1280 + num_frames: int = 81 + + # Denoising stage + guidance_scale: float = 9.0 + num_inference_steps: int = 50 + + # neg magic and pos magic + # pos_magic: str = "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。" + # neg_magic: str = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/teacache.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/teacache.py new file mode 100644 index 00000000..ada71d0b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/teacache.py @@ -0,0 +1,43 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/wan.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/wan.py new file mode 100644 index 00000000..9a23fb0b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/wan.py @@ -0,0 +1,217 @@ +# 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 + + 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 + + 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 + + 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 + + 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压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" + ) + + +@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 + + +@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 + # NOTE(will): default boundary timestep is tracked by PipelineConfig, but + # can be overridden during sampling + + +@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 + # NOTE(will): default boundary timestep is tracked by PipelineConfig, but + # can be overridden during sampling + + +# ============================================= +# ============= Causal Self-Forcing ============= +# ============================================= +@dataclass +class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParams): + pass diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/zimage.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/zimage.py new file mode 100644 index 00000000..f4530e97 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/sample/zimage.py @@ -0,0 +1,33 @@ +# 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 ZImageSamplingParams(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 + + teacache_params: TeaCacheParams = field( + default_factory=lambda: TeaCacheParams( + teacache_thresh=0.15, + coefficients=[ + 7.33226126e02, + -4.01131952e02, + 6.75869174e01, + -3.14987800e00, + 9.61237896e-02, + ], + ) + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/utils.py new file mode 100644 index 00000000..d2cc69ad --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/configs/utils.py @@ -0,0 +1,61 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/README.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/README.md new file mode 100644 index 00000000..7b41bd51 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/README.md @@ -0,0 +1,31 @@ +# 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 +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/setup.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/setup.py new file mode 100644 index 00000000..3a1bdb67 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/setup.py @@ -0,0 +1,26 @@ +# 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", + ], +) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/tests/test_vmoba_attn.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/tests/test_vmoba_attn.py new file mode 100644 index 00000000..f4304bda --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/tests/test_vmoba_attn.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 + +import random + +import pytest +import torch +from 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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/vmoba/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/vmoba/__init__.py new file mode 100644 index 00000000..8119387c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/vmoba/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +from .vmoba import moba_attn_varlen, process_moba_input, process_moba_output diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/vmoba/vmoba.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/vmoba/vmoba.py new file mode 100644 index 00000000..8a29360a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/csrc/attn/vmoba_attn/vmoba/vmoba.py @@ -0,0 +1,1086 @@ +# 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", + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/cli.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/cli.md new file mode 100644 index 00000000..4b0f29c7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/cli.md @@ -0,0 +1,184 @@ +# 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`. +- Python 3.11+ if you plan to use the OpenAI Python SDK. + + +## Supported Arguments + +### Server Arguments + +- `--model-path {MODEL_PATH}`: Path to the model or model ID +- `--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-size {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 + + +### 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 + + +#### 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). + + +## 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[@]}" +``` + +Once the generation task has finished, the server will shut down automatically. + +> [!NOTE] +> The HTTP server-related arguments are ignored in this subcommand. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/contributing.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/contributing.md new file mode 100644 index 00000000..fb8b4456 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/contributing.md @@ -0,0 +1,56 @@ +# Contributing to SGLang Diffusion + +This guide outlines the requirements for contributing to the SGLang Diffusion module (`sglang.multimodal_gen`). + +## 1. 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"). + +## 2. 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 + ### Performance Comparison Report + ... + ``` +4. **Paste**: paste the table into the PR description + +## 3. CI-Based Change Protection + +Consider adding tests to the `pr-test` or `nightly-test` suites to safeguard your changes, especially for PRs that: + +1. support a new model +2. support or fix important features +3. significantly improve performance + +See [test](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen/test) for examples diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/install.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/install.md new file mode 100644 index 00000000..894a414b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/install.md @@ -0,0 +1,48 @@ +# Install SGLang-diffusion + +You can install sglang-diffusion using one of the methods below. + +This page primarily applies to common NVIDIA GPU platforms. + +## 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](), built from the [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:dev \ + sglang generate --model-path black-forest-labs/FLUX.1-dev \ + --prompt "A logo With Bold Large text: SGL Diffusion" \ + --save-output +``` diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/openai_api.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/openai_api.md new file mode 100644 index 00000000..568104f1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/openai_api.md @@ -0,0 +1,281 @@ +# 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. + +## 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`). + +--- + +## 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** +> The `response_format=url` option is not supported for `POST /v1/images/generations` and will return a `400` error. + +#### Edit an image + +**Endpoint:** `POST /v1/images/edits` + +This endpoint accepts a multipart form upload with an input image 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=@input.png" \ + -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=@input.png" \ + -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/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 a LoRA adapter and merges its weights into the model. + +**Endpoint:** `POST /v1/set_lora` + +**Parameters:** +- `lora_nickname` (string, required): A unique identifier for this LoRA +- `lora_path` (string, optional): Path to the `.safetensors` file or Hugging Face repo ID. Required for the first load; optional if re-activating a cached nickname + +**Curl 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" + }' +``` + + +#### 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` + +**Curl Example:** + +```bash +curl -X POST http://localhost:30010/v1/merge_lora_weights \ + -H "Content-Type: application/json" +``` + + +#### 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" +``` + +### 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... diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/support_matrix.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/support_matrix.md new file mode 100644 index 00000000..d8eeb037 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/support_matrix.md @@ -0,0 +1,48 @@ +# 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) | +|:-----------------------------|:--------------------------------------------------|:--------------------|:--------:|:-----------------:|:---------:|:----------------------------:| +| 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 | ✅ | ✅ | ✅ | ⭕ | + +**Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue. + +### 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 | +| Qwen Image | `Qwen/Qwen-Image` | Any resolution | +| Qwen Image Edit | `Qwen/Qwen-Image-Edit` | Any resolution | + +## Special requirements + +### Sliding Tile Attention + +- Currently, only Hopper GPUs (H100s) are supported. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/support_new_models.md b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/support_new_models.md new file mode 100644 index 00000000..e51bd68d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/docs/support_new_models.md @@ -0,0 +1,107 @@ +# 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 modular pipeline architecture. This +design allows developers to easily construct complex, customized pipelines for various diffusion models by combining and +reusing different components. + +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 common function within the diffusion process. Examples include prompt encoding, the denoising loop, or VAE decoding. These stages are designed to be self-contained and reusable across different pipelines. + +## Key Components for Implementation + +To add support for a new diffusion model, you will primarily need to define or configure the following components: + +1. **`PipelineConfig`**: This is a dataclass that holds all the static configurations for your model pipeline. It includes paths to model components (like UNet, VAE, text encoders), precision settings (e.g., `fp16`, `bf16`), and other model-specific architectural parameters. Each model typically has its own subclass of `PipelineConfig`. + +2. **`SamplingParams`**: This dataclass defines the parameters that control the generation process at runtime. These are the user-provided inputs for a generation request, such as the `prompt`, `negative_prompt`, `guidance_scale`, `num_inference_steps`, `seed`, output dimensions (`height`, `width`), etc. + +3. **`ComposedPipeline` (not a config)**: This is the central class where you define the structure of your model's generation pipeline. You will create a new class that inherits from `ComposedPipelineBase` and, within it, instantiate and chain together the necessary `PipelineStage`s in the correct order. See `ComposedPipelineBase` and `PipelineStage` base definitions: + - [`ComposedPipelineBase`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/pipelines/composed_pipeline_base.py) + - [`PipelineStage`]( https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/pipelines/stages/base.py) + - [Central registry (models/config mapping)](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/registry.py) + +4. **Modules (components referenced by the pipeline)**: Each pipeline references a set of modules that are loaded from the model repository (e.g., Diffusers `model_index.json`) and assembled via the registry/loader. Common modules include: + - `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 (may be distinct from or combined with `processor`). + - `dit/transformer`: The core denoising network (DiT/UNet architecture) operating in latent space. + - `scheduler`: Controls the timestep schedule and denoising dynamics throughout inference. + - `vae`: Variational Autoencoder for encoding/decoding between pixel space and latent space. + +## Available Pipeline Stages + +You can build your custom `ComposedPipeline` by combining the following available stages as your will. Each stage is responsible for a specific part of the generation process. + +| Stage Class | Description | +| -------------------------------- | ------------------------------------------------------------------------------------------------------- | +| `InputValidationStage` | Validates the user-provided `SamplingParams` to ensure they are correct before starting the pipeline. | +| `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` | Specifically encodes an input image into the latent space using a Variational Autoencoder (VAE). | +| `ConditioningStage` | Prepares the conditioning tensors (e.g., from text or image embeddings) for the denoising loop. | +| `TimestepPreparationStage` | Prepares the scheduler's timesteps for the diffusion process. | +| `LatentPreparationStage` | Creates the initial noisy latent tensor that will be denoised. | +| `DenoisingStage` | Executes the main denoising loop, iteratively applying the model (e.g., UNet) to refine the latents. | +| `DecodingStage` | Decodes the final latent tensor from the denoising loop back into pixel space (e.g., an image) using the VAE. | +| `DmdDenoisingStage` | A specialized denoising stage for certain model architectures. | +| `CausalDMDDenoisingStage` | A specialized causal denoising stage for specific video models. | + +## Example: Implementing `Qwen-Image-Edit` + +To illustrate the process, let's look at how `Qwen-Image-Edit` is implemented. The typical implementation order is: + +1. **Analyze Required Modules**: + - Study the target model's components by examining its `model_index.json` or Diffusers implementation to identify required modules: + - `processor`: Image preprocessing and feature extraction + - `scheduler`: Diffusion timestep scheduling + - `text_encoder`: Text-to-embedding conversion + - `tokenizer`: Text tokenization for the encoder + - `transformer`: Core DiT denoising network + - `vae`: Variational autoencoder for latent encoding/decoding + +2. **Create Configs**: + - **PipelineConfig**: [`QwenImageEditPipelineConfig`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/configs/pipelines/qwen_image.py) defines model-specific parameters, precision settings, preprocessing functions, and latent shape calculations. + - **SamplingParams**: [`QwenImageSamplingParams`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/configs/sample/qwenimage.py) sets runtime defaults like `num_frames=1`, `guidance_scale=4.0`, `num_inference_steps=50`. + +3. **Implement Model Components**: + - Adapt or implement specific model components in the appropriate directories: + - **DiT/Transformer**: Implement in [`runtime/models/dits/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/dits/) - e.g., [`qwen_image.py`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py) for Qwen's DiT architecture + - **Encoders**: Implement in [`runtime/models/encoders/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/encoders/) - e.g., text encoders like [`qwen2_5vl.py`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py) + - **VAEs**: Implement in [`runtime/models/vaes/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/vaes/) - e.g., [`autoencoder_kl_qwenimage.py`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py) + - **Schedulers**: Implement in [`runtime/models/schedulers/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/schedulers/) if needed + - These components handle the core model logic, attention mechanisms, and data transformations specific to the target diffusion model. + +4. **Define Pipeline Class**: + - The [`QwenImageEditPipeline`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/architectures/basic/qwen_image/qwen_image.py) class inherits from `ComposedPipelineBase` and orchestrates stages sequentially. + - Declare required modules via `_required_config_modules` and implement the pipeline stages: + + ```python + class QwenImageEditPipeline(ComposedPipelineBase): + pipeline_name = "QwenImageEditPipeline" # Matches Diffusers model_index.json + _required_config_modules = ["processor", "scheduler", "text_encoder", "tokenizer", "transformer", "vae"] + + def create_pipeline_stages(self, server_args: ServerArgs): + """Set up pipeline stages sequentially.""" + self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage()) + self.add_stage(stage_name="prompt_encoding_stage_primary", stage=ImageEncodingStage(...)) + self.add_stage(stage_name="image_encoding_stage_primary", stage=ImageVAEEncodingStage(...)) + self.add_stage(stage_name="timestep_preparation_stage", stage=TimestepPreparationStage(...)) + self.add_stage(stage_name="latent_preparation_stage", stage=LatentPreparationStage(...)) + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + self.add_stage(stage_name="denoising_stage", stage=DenoisingStage(...)) + self.add_stage(stage_name="decoding_stage", stage=DecodingStage(...)) + ``` + The pipeline is constructed by adding stages in order. `Qwen-Image-Edit` uses `ImageEncodingStage` (for prompt and image processing) and `ImageVAEEncodingStage` (for latent extraction) before standard denoising and decoding. + +5. **Register Configs**: + - Register the configs in the central registry ([`registry.py`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/registry.py)) via `_register_configs` to enable automatic loading and instantiation for the model. Modules are automatically loaded and injected based on the config and repository structure. + +By following this pattern of defining configurations and composing pipelines, you can integrate new diffusion models +into SGLang with ease. diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/envs.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/envs.py new file mode 100644 index 00000000..56418e72 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/envs.py @@ -0,0 +1,328 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +import importlib.util + +# 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 collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import diffusers +import torch +from packaging import version + +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 + + +def _is_hip(): + has_rocm = torch.version.hip is not None + return has_rocm + + +def _is_cuda(): + has_cuda = torch.version.cuda is not None + return has_cuda + + +def _is_musa(): + try: + if hasattr(torch, "musa") and torch.musa.is_available(): + return True + except ModuleNotFoundError: + return False + + +def _is_mps(): + return torch.backends.mps.is_available() + + +class PackagesEnvChecker: + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super(PackagesEnvChecker, cls).__new__(cls) + cls._instance.initialize() + return cls._instance + + def initialize(self): + self.packages_info = { + "has_aiter": self.check_aiter(), + "diffusers_version": self.check_diffusers_version(), + } + + def check_aiter(self): + """ + Checks whether ROCm AITER library is installed + """ + try: + + logger.info("Using AITER as the attention library") + return True + except: + if _is_hip(): + logger.warning( + f'Using AMD GPUs, but library "aiter" is not installed, ' + "defaulting to other attention mechanisms" + ) + return False + + def check_flash_attn(self): + if not torch.cuda.is_available(): + return False + if _is_musa(): + logger.info( + "Flash Attention library is not supported on MUSA for the moment." + ) + return False + try: + return True + except ImportError: + logger.warning( + f'Flash Attention library "flash_attn" not found, ' + f"using pytorch attention implementation" + ) + return False + + def check_long_ctx_attn(self): + if not torch.cuda.is_available(): + return False + try: + return importlib.util.find_spec("yunchang") is not None + except ImportError: + logger.warning( + f'Ring Flash Attention library "yunchang" not found, ' + f"using pytorch attention implementation" + ) + return False + + def check_diffusers_version(self): + if version.parse( + version.parse(diffusers.__version__).base_version + ) < version.parse("0.30.0"): + raise RuntimeError( + f"Diffusers version: {version.parse(version.parse(diffusers.__version__).base_version)} is not supported," + f"please upgrade to version > 0.30.0" + ) + return version.parse(version.parse(diffusers.__version__).base_version) + + def get_packages_info(self): + return self.packages_info + + +PACKAGES_CHECKER = PackagesEnvChecker() + + +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: + if value is None: + return None + return int(value) + + +# 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": lambda: os.getenv( + "SGLANG_DIFFUSION_TARGET_DEVICE", "cuda" + ), + # Maximum number of compilation jobs to run in parallel. + # By default this is the number of CPUs + "MAX_JOBS": lambda: os.getenv("MAX_JOBS", None), + # 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": lambda: os.getenv("NVCC_THREADS", None), + # If set, sgl_diffusion will use precompiled binaries (*.so) + "SGLANG_DIFFUSION_USE_PRECOMPILED": lambda: bool( + os.environ.get("SGLANG_DIFFUSION_USE_PRECOMPILED") + ) + or bool(os.environ.get("SGLANG_DIFFUSION_PRECOMPILED_WHEEL_LOCATION")), + # CMake build type + # If not set, defaults to "Debug" or "RelWithDebInfo" + # Available options: "Debug", "Release", "RelWithDebInfo" + "CMAKE_BUILD_TYPE": lambda: os.getenv("CMAKE_BUILD_TYPE"), + # If set, sgl_diffusion will print verbose logs during installation + "VERBOSE": lambda: bool(int(os.getenv("VERBOSE", "0"))), + # Root directory for FASTVIDEO 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": lambda: os.path.expanduser( + os.getenv( + "SGLANG_DIFFUSION_CONFIG_ROOT", + os.path.join(get_default_config_root(), "sgl_diffusion"), + ) + ), + # ================== Runtime Env Vars ================== + # Root directory for FASTVIDEO cache files + # Defaults to `~/.cache/sgl_diffusion` unless `XDG_CACHE_HOME` is set + "SGLANG_DIFFUSION_CACHE_ROOT": lambda: os.path.expanduser( + os.getenv( + "SGLANG_DIFFUSION_CACHE_ROOT", + 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": lambda: int( + os.environ.get("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": lambda: os.environ.get( + "SGLANG_DIFFUSION_NCCL_SO_PATH", None + ), + # 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": lambda: os.environ.get("LD_LIBRARY_PATH", None), + # Internal flag to enable Dynamo fullgraph capture + "SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE": lambda: bool( + os.environ.get("SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1") != "0" + ), + # local rank of the process in the distributed setting, used to determine + # the GPU device id + "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), + # used to control the visible devices in the distributed setting + "CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None), + # timeout for each iteration in the engine + "SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S": lambda: int( + os.environ.get("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": lambda: int( + os.getenv("SGLANG_DIFFUSION_CONFIGURE_LOGGING", "1") + ), + "SGLANG_DIFFUSION_LOGGING_CONFIG_PATH": lambda: os.getenv( + "SGLANG_DIFFUSION_LOGGING_CONFIG_PATH" + ), + # this is used for configuring the default logging level + "SGLANG_DIFFUSION_LOGGING_LEVEL": lambda: os.getenv( + "SGLANG_DIFFUSION_LOGGING_LEVEL", "INFO" + ), + # if set, SGLANG_DIFFUSION_LOGGING_PREFIX will be prepended to all log messages + "SGLANG_DIFFUSION_LOGGING_PREFIX": lambda: os.getenv( + "SGLANG_DIFFUSION_LOGGING_PREFIX", "" + ), + # Trace function calls + # If set to 1, sgl_diffusion will trace function calls + # Useful for debugging + "SGLANG_DIFFUSION_TRACE_FUNCTION": lambda: int( + os.getenv("SGLANG_DIFFUSION_TRACE_FUNCTION", "0") + ), + # Path to the attention configuration file. Only used for sliding tile + # attention for now. + "SGLANG_DIFFUSION_ATTENTION_CONFIG": lambda: ( + None + if os.getenv("SGLANG_DIFFUSION_ATTENTION_CONFIG", None) is None + else os.path.expanduser(os.getenv("SGLANG_DIFFUSION_ATTENTION_CONFIG", ".")) + ), + # Use dedicated multiprocess context for workers. + # Both spawn and fork work + "SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD": lambda: os.getenv( + "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": lambda: ( + None + if os.getenv("SGLANG_DIFFUSION_TORCH_PROFILER_DIR", None) is None + else os.path.expanduser(os.getenv("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": lambda: get_bool_env_var( + "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": lambda: get_bool_env_var( + "SGLANG_DIFFUSION_STAGE_LOGGING" + ), +} + + +# 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()) + + +def get_torch_distributed_backend() -> str: + if torch.cuda.is_available(): + return "nccl" + elif _is_musa(): + return "mccl" + elif _is_mps(): + return "gloo" + else: + raise NotImplementedError( + "No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available" + ) + + +def get_device(local_rank: int) -> torch.device: + if torch.cuda.is_available(): + return torch.device("cuda", local_rank) + elif _is_musa(): + return torch.device("musa", local_rank) + elif _is_mps(): + return torch.device("mps") + else: + return torch.device("cpu") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/registry.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/registry.py new file mode 100644 index 00000000..1e2fafeb --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/registry.py @@ -0,0 +1,423 @@ +# 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 +import re +from functools import lru_cache +from typing import Any, Callable, Dict, List, Optional, Tuple, Type + +from sglang.multimodal_gen.configs.pipeline_configs import ( + FastHunyuanConfig, + FluxPipelineConfig, + HunyuanConfig, + StepVideoT2VConfig, + WanI2V480PConfig, + WanI2V720PConfig, + WanT2V480PConfig, + WanT2V720PConfig, + ZImagePipelineConfig, +) +from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.flux import Flux2PipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( + QwenImageEditPipelineConfig, + QwenImagePipelineConfig, +) +from sglang.multimodal_gen.configs.pipeline_configs.wan import ( + FastWan2_1_T2V_480P_Config, + FastWan2_2_TI2V_5B_Config, + Wan2_2_I2V_A14B_Config, + Wan2_2_T2V_A14B_Config, + Wan2_2_TI2V_5B_Config, +) +from sglang.multimodal_gen.configs.sample.flux import FluxSamplingParams +from sglang.multimodal_gen.configs.sample.hunyuan import ( + FastHunyuanSamplingParam, + HunyuanSamplingParams, +) +from sglang.multimodal_gen.configs.sample.qwenimage import QwenImageSamplingParams +from sglang.multimodal_gen.configs.sample.stepvideo import StepVideoT2VSamplingParams +from sglang.multimodal_gen.configs.sample.wan import ( + FastWanT2V480PConfig, + 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 +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]] = {} + + +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 hasattr(cls, "pipeline_name"): + if cls.pipeline_name in _PIPELINE_REGISTRY: + logger.warning( + f"Duplicate pipeline name '{cls.pipeline_name}' found. Overwriting." + ) + _PIPELINE_REGISTRY[cls.pipeline_name] = cls + logger.debug( + f"Registering pipelines complete, {len(_PIPELINE_REGISTRY)} pipelines registered" + ) + + +# --- 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_config_info(model_path: str) -> Optional[ConfigInfo]: + """ + Gets the ConfigInfo for a given model path using mappings and detectors. + """ + # 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. + cleaned_model_path = re.sub(r"--", "/", model_path.lower()) + all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True) + for model_hf_path in all_model_hf_paths: + if model_hf_path.lower() in cleaned_model_path: + logger.debug( + f"Resolved model name '{model_hf_path}' from partial path match." + ) + model_id = _MODEL_HF_PATH_TO_NAME[model_hf_path] + 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] + + +@lru_cache(maxsize=1) +def get_model_info(model_path: str) -> 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. + """ + # 1. Discover all available pipeline classes and cache them + _discover_and_register_pipelines() + + # 2. Get pipeline class from model's 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}") + 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}'") + return None + + pipeline_cls = _PIPELINE_REGISTRY.get(pipeline_class_name) + if not pipeline_cls: + 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())}" + ) + return None + + # 3. Get configuration classes (sampling, pipeline config) + config_info = _get_config_info(model_path) + if not config_info: + 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())}" + ) + return None + + # 4. Combine the complete model info + model_info = ModelInfo( + pipeline_cls=pipeline_cls, + sampling_param_cls=config_info.sampling_param_cls, + pipeline_config_cls=config_info.pipeline_config_cls, + ) + logger.info(f"Found model info: {model_info}") + + return model_info + + +# Registration of model configs +def _register_configs(): + # Hunyuan + register_configs( + sampling_param_cls=HunyuanSamplingParams, + pipeline_config_cls=HunyuanConfig, + hf_model_paths=[ + "hunyuanvideo-community/HunyuanVideo", + ], + model_detectors=[lambda hf_id: "hunyuan" in hf_id.lower()], + ) + register_configs( + sampling_param_cls=FastHunyuanSamplingParam, + pipeline_config_cls=FastHunyuanConfig, + hf_model_paths=[ + "FastVideo/FastHunyuan-diffusers", + ], + ) + + # StepVideo + register_configs( + sampling_param_cls=StepVideoT2VSamplingParams, + pipeline_config_cls=StepVideoT2VConfig, + hf_model_paths=[ + "FastVideo/stepvideo-t2v-diffusers", + ], + model_detectors=[lambda hf_id: "stepvideo" in hf_id.lower()], + ) + + # 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_14B_SamplingParams, + pipeline_config_cls=WanT2V720PConfig, + hf_model_paths=[ + "Wan-AI/Wan2.1-T2V-14B-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=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", + ], + ) + + # 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=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()], + ) + register_configs( + sampling_param_cls=ZImageSamplingParams, + pipeline_config_cls=ZImagePipelineConfig, + hf_model_paths=[ + "Tongyi-MAI/Z-Image-Turbo", + ], + model_detectors=[lambda hf_id: "z-image" in hf_id.lower()], + ) + + # Qwen-Image + register_configs( + sampling_param_cls=QwenImageSamplingParams, + pipeline_config_cls=QwenImagePipelineConfig, + hf_model_paths=["Qwen/Qwen-Image"], + ) + register_configs( + sampling_param_cls=QwenImageSamplingParams, + pipeline_config_cls=QwenImageEditPipelineConfig, + hf_model_paths=["Qwen/Qwen-Image-Edit"], + ) + + +_register_configs() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/__init__.py new file mode 100644 index 00000000..9edfd5c6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/__init__.py @@ -0,0 +1,55 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 + +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 * + +__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", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/communication_op.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/communication_op.py new file mode 100644 index 00000000..61672ca4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/communication_op.py @@ -0,0 +1,55 @@ +# 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 + +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) -> torch.Tensor: + """All-reduce the input tensor across model parallel group.""" + return get_tp_group().all_reduce(input_) + + +def tensor_model_parallel_all_gather( + input_: torch.Tensor, dim: int = -1 +) -> torch.Tensor: + """All-gather the input tensor across model parallel group.""" + return get_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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/__init__.py new file mode 100644 index 00000000..af2eb7d1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/__init__.py @@ -0,0 +1 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py new file mode 100644 index 00000000..01bdf1c2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py @@ -0,0 +1,297 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py new file mode 100644 index 00000000..434cf384 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py @@ -0,0 +1,161 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py new file mode 100644 index 00000000..c128c69f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py @@ -0,0 +1,79 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py new file mode 100644 index 00000000..2d1ef558 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py @@ -0,0 +1,258 @@ +# 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), + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py new file mode 100644 index 00000000..598e7be9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py @@ -0,0 +1,450 @@ +# 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 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", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py new file mode 100644 index 00000000..dd42b875 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py @@ -0,0 +1,1226 @@ +# 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 import envs +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.utils.logging_utils import init_logger + +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.""" + from sglang.multimodal_gen.runtime.platforms import current_platform + + return ( + torch.device(f"cuda:{envs.LOCAL_RANK}") + if current_platform.is_cuda_alike() + else torch.device("mps") + ) + + +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. + 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() + + from sglang.multimodal_gen.runtime.platforms import current_platform + + 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): + # Platform-aware graph capture + from sglang.multimodal_gen.runtime.platforms import current_platform + + 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. + 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. + 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 = envs.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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py new file mode 100644 index 00000000..82dbb588 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -0,0 +1,1144 @@ +# 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 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: Optional[GroupCoordinator] = None +_TP: Optional[GroupCoordinator] = None +_SP: Optional[SequenceParallelGroupCoordinator] = None +_PP: Optional[PipelineGroupCoordinator] = None +_CFG: Optional[GroupCoordinator] = None +_DP: Optional[GroupCoordinator] = None +_DIT: Optional[GroupCoordinator] = None +_VAE: Optional[GroupCoordinator] = None + +logger = init_logger(__name__) + +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) + + +_WORLD: GroupCoordinator | None = None +_NODE: GroupCoordinator | None = None + + +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", + ) + + +# xDiT +def init_parallel_group_coordinator( + group_ranks: List[List[int]], + local_rank: int, + backend: str, + parallel_mode: str, + **kwargs, +) -> GroupCoordinator: + """ + Returns 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 init_parallel_group_coordinator( +# group_ranks: list[list[int]], +# local_rank: int, +# backend: str, +# use_message_queue_broadcaster: bool = False, +# group_name: str | None = None, +# ) -> GroupCoordinator: +# return GroupCoordinator( +# group_ranks=group_ranks, +# local_rank=local_rank, +# torch_distributed_backend=backend, +# use_device_communicator=True, +# use_message_queue_broadcaster=use_message_queue_broadcaster, +# group_name=group_name, +# ) + + +_TP: GroupCoordinator | None = None + + +def get_tp_group() -> GroupCoordinator: + assert _TP is not None, "tensor model parallel group is not initialized" + return _TP + + +_ENABLE_CUSTOM_ALL_REDUCE = True + + +def set_custom_all_reduce(enable: bool): + global _ENABLE_CUSTOM_ALL_REDUCE + _ENABLE_CUSTOM_ALL_REDUCE = enable + + +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, +): + # 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", + world_size, + rank, + local_rank, + distributed_init_method, + backend, + ) + 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, don't pass device_id as it doesn't support device indices + extra_args = {} if current_platform.is_mps() else dict(device_id=device_id) + 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" + + +_SP: GroupCoordinator | None = None + + +def get_sp_group() -> SequenceParallelGroupCoordinator: + assert _SP is not None, "pipeline model parallel group is not initialized" + return _SP + + +_DP: GroupCoordinator | None = None + + +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: + backend = envs.get_torch_distributed_backend() + # 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" + + from yunchang import set_seq_parallel_pg + from yunchang.globals import PROCESS_GROUP + + set_seq_parallel_pg( + sp_ulysses_degree=ulysses_degree, + sp_ring_degree=ring_degree, + rank=get_world_group().rank_in_group, + world_size=dit_parallel_size, + ) + + _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 initialize_model_parallel( +# tensor_model_parallel_size: int = 1, +# sequence_model_parallel_size: int = 1, +# data_parallel_size: int = 1, +# backend: str | None = None, +# ) -> None: +# """ +# Initialize model parallel groups. +# +# Arguments: +# tensor_model_parallel_size: number of GPUs used for tensor model +# parallelism (used for language encoder). +# sequence_model_parallel_size: number of GPUs used for sequence model +# parallelism (used for DiT). +# """ +# # Get world size and rank. Ensure some consistencies. +# assert ( +# _WORLD is not None +# ), "world group is not initialized, please call init_distributed_environment first" +# world_size: int = get_world_size() +# backend = backend or torch.distributed.get_backend(get_world_group().device_group) +# assert ( +# world_size >= tensor_model_parallel_size +# ), f"world_size({world_size}) must be greater than or equal to tensor_model_parallel_size({tensor_model_parallel_size})" +# num_tensor_model_parallel_groups: int = world_size // tensor_model_parallel_size +# global _TP +# assert _TP is None, "tensor model parallel group is already initialized" +# group_ranks = [] +# for i in range(num_tensor_model_parallel_groups): +# ranks = list( +# range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size) +# ) +# group_ranks.append(ranks) +# +# # message queue broadcaster is only used in tensor model parallel group +# _TP = init_parallel_group_coordinator( +# group_ranks, +# get_world_group().local_rank, +# backend, +# use_message_queue_broadcaster=True, +# group_name="tp", +# ) +# +# # Build the sequence model-parallel groups. +# num_sequence_model_parallel_groups: int = world_size // sequence_model_parallel_size +# global _SP +# assert _SP is None, "sequence model parallel group is already initialized" +# group_ranks = [] +# +# # Since SP is incompatible with TP and PP, we can use a simpler group creation logic +# for i in range(num_sequence_model_parallel_groups): +# # Create groups of consecutive ranks +# ranks = list( +# range( +# i * sequence_model_parallel_size, (i + 1) * sequence_model_parallel_size +# ) +# ) +# group_ranks.append(ranks) +# +# _SP = init_parallel_group_coordinator( +# group_ranks, get_world_group().local_rank, backend, group_name="sp" +# ) +# +# # Build the data parallel groups. +# num_data_parallel_groups: int = sequence_model_parallel_size +# global _DP +# assert _DP is None, "data parallel group is already initialized" +# group_ranks = [] +# +# for i in range(num_data_parallel_groups): +# ranks = list(range(i, world_size, num_data_parallel_groups)) +# group_ranks.append(ranks) +# +# _DP = init_parallel_group_coordinator( +# group_ranks, get_world_group().local_rank, backend, group_name="dp" +# ) +# + + +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://", +): + 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", + world_size, + device, + 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, + ) + 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) + + +def model_parallel_is_initialized() -> bool: + """Check if tensor, sequence parallel groups are initialized.""" + return _TP is not None and _SP is not None and _DP is not None and _CFG 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. + + Args: + tp_group (GroupCoordinator): the tp group coordinator + """ + 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 initialize_tensor_parallel_group( + tensor_model_parallel_size: int = 1, + backend: str | None = None, + group_name_suffix: str = "", +) -> GroupCoordinator: + """Initialize a tensor parallel group for a specific model. + + This function creates a tensor parallel group that can be used with the + patch_tensor_parallel_group context manager. It allows different models + to use different tensor parallelism configurations. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for tensor parallelism that can be used with + the patch_tensor_parallel_group context manager. + + Example usage: + ```python + # Initialize tensor parallel group for model1 + tp_group_model1 = initialize_tensor_parallel_group( + tensor_model_parallel_size=4, + group_name_suffix="model1" + ) + + # Use tensor parallelism for model1 + with patch_tensor_parallel_group(tp_group_model1): + # Run model1 with tensor parallelism + output1 = model1(input1) + ``` + """ + # 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) + + # Ensure the world size is compatible with the parallelism configuration + assert ( + world_size % tensor_model_parallel_size == 0 + ), f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size})" + + # Build the tensor model-parallel groups. + num_tensor_model_parallel_groups: int = world_size // tensor_model_parallel_size + tp_group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list( + range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size) + ) + tp_group_ranks.append(ranks) + + # Create TP group coordinator with a unique name + group_name = f"tp_{group_name_suffix}" if group_name_suffix else "tp" + tp_group = init_parallel_group_coordinator( + tp_group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name=group_name, + ) + + return tp_group + + +def initialize_sequence_parallel_group( + sequence_model_parallel_size: int = 1, + backend: str | None = None, + group_name_suffix: str = "", +) -> GroupCoordinator: + """Initialize a sequence parallel group for a specific model. + + This function creates a sequence parallel group that can be used with the + patch_sequence_parallel_group context manager. It allows different models + to use different sequence parallelism configurations. + + Arguments: + sequence_model_parallel_size: number of GPUs used for sequence model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for sequence parallelism that can be used with + the patch_sequence_parallel_group context manager. + + Example usage: + ```python + # Initialize sequence parallel group for model2 + sp_group_model2 = initialize_sequence_parallel_group( + sequence_model_parallel_size=2, + group_name_suffix="model2" + ) + + # Use sequence parallelism for model2 + with patch_sequence_parallel_group(sp_group_model2): + # Run model2 with sequence parallelism + output2 = model2(input2) + ``` + """ + # 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) + + # Ensure the world size is compatible with the parallelism configuration + assert ( + world_size % sequence_model_parallel_size == 0 + ), f"World size ({world_size}) must be divisible by sequence_model_parallel_size ({sequence_model_parallel_size})" + + # Build the sequence model-parallel groups. + num_sequence_model_parallel_groups: int = world_size // sequence_model_parallel_size + sp_group_ranks = [] + + for i in range(num_sequence_model_parallel_groups): + # Create groups of consecutive ranks + ranks = list( + range( + i * sequence_model_parallel_size, (i + 1) * sequence_model_parallel_size + ) + ) + sp_group_ranks.append(ranks) + + # Create SP group coordinator with a unique name + group_name = f"sp_{group_name_suffix}" if group_name_suffix else "sp" + sp_group = init_parallel_group_coordinator( + sp_group_ranks, get_world_group().local_rank, backend, group_name=group_name + ) + + return sp_group + + +# * QUERY +def get_world_group() -> GroupCoordinator: + assert _WORLD is not None, "world group is not initialized" + return _WORLD + + +# TP +def get_tp_group() -> GroupCoordinator: + assert _TP is not None, "tensor model parallel group is not initialized" + return _TP + + +def get_tensor_model_parallel_world_size(): + """Return world size for the tensor model parallel group.""" + return get_tp_group().world_size + + +def get_tensor_model_parallel_rank(): + """Return my rank for the tensor model parallel group.""" + return get_tp_group().rank_in_group + + +def get_sequence_parallel_world_size(): + """Return world size for the sequence parallel group.""" + return get_sp_group().world_size + + +def get_sequence_parallel_rank(): + """Return my rank for the sequence parallel group.""" + return get_sp_group().rank_in_group + + +def get_ulysses_parallel_world_size(): + return get_sp_group().ulysses_world_size + + +def get_ulysses_parallel_rank(): + return get_sp_group().ulysses_rank + + +def get_ring_parallel_world_size(): + return get_sp_group().ring_world_size + + +def get_ring_parallel_rank(): + 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(): + """Return world size for the pipeline model parallel group.""" + return get_pp_group().world_size + + +def get_pipeline_parallel_rank(): + """Return my rank for the pipeline model parallel group.""" + return get_pp_group().rank_in_group + + +def is_pipeline_first_stage(): + """Return True if in the first pipeline model parallel stage, False otherwise.""" + return get_pipeline_parallel_rank() == 0 + + +def is_pipeline_last_stage(): + """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(): + """Return world size for the classifier_free_guidance parallel group.""" + return get_cfg_group().world_size + + +def get_classifier_free_guidance_rank(): + """Return my rank for the classifier_free_guidance parallel group.""" + return get_cfg_group().rank_in_group + + +# DP +def get_dp_group() -> GroupCoordinator: + assert _DP is not None, "pipeline model parallel group is not initialized" + return _DP + + +def get_data_parallel_world_size(): + """Return world size for the data parallel group.""" + return get_dp_group().world_size + + +def get_data_parallel_rank(): + """Return my rank for the data parallel group.""" + return get_dp_group().rank_in_group + + +def is_dp_last_group(): + """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(): + """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() + ) + + +# Add VAE getter functions +def get_vae_parallel_group() -> GroupCoordinator: + assert _VAE is not None, "VAE parallel group is not initialized" + return _VAE + + +def get_vae_parallel_world_size(): + """Return world size for the VAE parallel group.""" + return get_vae_parallel_group().world_size + + +def get_vae_parallel_rank(): + """Return my rank for the VAE parallel group.""" + return get_vae_parallel_group().rank_in_group + + +# * SET + + +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, + ) + + +def model_parallel_is_initialized(): + """Check if tensor and pipeline 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 + ) + + +def init_dit_group( + dit_parallel_size: int, + backend: str, +): + global _DIT + _DIT = torch.distributed.new_group( + ranks=list(range(dit_parallel_size)), backend=backend + ) + + +def get_dit_group(): + 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 + if _TP: + _TP.destroy() + _TP = None + + global _SP + if _SP: + _SP.destroy() + _SP = None + + global _DP + if _DP: + _DP.destroy() + _DP = None + + +# xDit +# def destroy_model_parallel(): +# """Set the groups to none and destroy them.""" +# global _DP +# if _DP: +# _DP.destroy() +# _DP = None +# +# global _CFG +# if _CFG: +# _CFG.destroy() +# _CFG = None +# +# global _SP +# if _SP: +# _SP.destroy() +# _SP = None +# +# global _TP +# if _TP: +# _TP.destroy() +# _TP = None +# +# global _PP +# if _PP: +# _PP.destroy() +# _PP = None +# +# global _VAE +# if _VAE: +# _VAE.destroy() +# _VAE = None + + +def destroy_distributed_environment(): + global _WORLD + if _WORLD: + _WORLD.destroy() + _WORLD = None + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/utils.py new file mode 100644 index 00000000..2d84f8b5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/distributed/utils.py @@ -0,0 +1,195 @@ +# 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, + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/__init__.py new file mode 100644 index 00000000..af2eb7d1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/__init__.py @@ -0,0 +1 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/__init__.py new file mode 100644 index 00000000..af2eb7d1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/__init__.py @@ -0,0 +1 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/cli_types.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/cli_types.py new file mode 100644 index 00000000..2e5107ec --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/cli_types.py @@ -0,0 +1,28 @@ +# 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) -> 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py new file mode 100644 index 00000000..9b4f8b59 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py @@ -0,0 +1,151 @@ +# 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 os +from typing import cast + +import sglang.multimodal_gen.envs as envs +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.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.perf_logger import ( + PerformanceLogger, + RequestTimings, +) +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): + """dump performance if necessary""" + if not (args.perf_dump_path and results): + return + + if isinstance(results, list): + result = results[0] if results else {} + else: + result = results + + timings_dict = result.get("timings") + if not (args.perf_dump_path and timings_dict): + return + + timings = RequestTimings(request_id=timings_dict.get("request_id")) + timings.stages = timings_dict.get("stages", {}) + timings.total_duration_ms = timings_dict.get("total_duration_ms", 0) + + PerformanceLogger.dump_benchmark_report( + file_path=args.perf_dump_path, + timings=timings, + meta={ + "prompt": prompt, + "model": server_args.model_path, + }, + tag="cli_generate", + ) + + +def generate_cmd(args: argparse.Namespace): + """The entry point for the generate command.""" + # FIXME(mick): do not hard code + args.request_id = generate_request_id() + + # Auto-enable stage logging if dump path is provided + if args.perf_dump_path: + os.environ["SGLANG_DIFFUSION_STAGE_LOGGING"] = "True" + envs.SGLANG_DIFFUSION_STAGE_LOGGING = True + + server_args = ServerArgs.from_cli_args(args) + sampling_params_kwargs = SamplingParams.get_cli_args(args) + generator = DiffGenerator.from_pretrained( + model_path=server_args.model_path, server_args=server_args + ) + + results = generator.generate(sampling_params_kwargs=sampling_params_kwargs) + + prompt = sampling_params_kwargs.get("prompt", None) + 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) -> None: + generate_cmd(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="sgl_diffusion 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/main.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/main.py new file mode 100644 index 00000000..c35dec33 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/main.py @@ -0,0 +1,44 @@ +# 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 = parser.parse_args() + if args.subparser in cmds: + cmds[args.subparser].validate(args) + + if hasattr(args, "dispatch_function"): + args.dispatch_function(args) + else: + parser.print_help() + + +if __name__ == "__main__": + main() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py new file mode 100644 index 00000000..5f939a28 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py @@ -0,0 +1,69 @@ +# 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.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) + server_args.post_init_serve() + launch_server(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="sgl_diffusion 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()] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py new file mode 100644 index 00000000..a4fc7527 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py @@ -0,0 +1,74 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +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", " ".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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py new file mode 100644 index 00000000..d296450f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -0,0 +1,459 @@ +# 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 multiprocessing as mp +import os +import time +from typing import Any + +import imageio +import numpy as np +import torch +import torchvision +from einops import rearrange + +from sglang.multimodal_gen.configs.sample.sampling_params import ( + DataType, + SamplingParams, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( + MergeLoraWeightsReq, + SetLoraReq, + UnmergeLoraWeightsReq, +) +from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request +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.server_args import PortArgs, ServerArgs +from sglang.multimodal_gen.runtime.sync_scheduler_client import sync_scheduler_client +from sglang.multimodal_gen.runtime.utils.logging_utils import ( + init_logger, + log_batch_completion, + log_generation_timer, + suppress_loggers, + suppress_other_loggers, +) + +suppress_loggers(["imageio", "imageio_ffmpeg", "PIL", "PIL_Image"]) + +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 + + +# TODO: rename +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 + self._current_lora_path: str | None = None + self._current_lora_nickname: str | None = None + self._is_lora_merged: bool = False + + @classmethod + def from_pretrained( + cls, + **kwargs, + ) -> "DiffGenerator": + """ + Create a DiffGenerator from a pretrained model. + + Args: + **kwargs: Additional arguments to customize model loading, set any ServerArgs or PipelineConfig attributes here. + + Returns: + The created DiffGenerator + + 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) + + @classmethod + def from_server_args(cls, server_args: ServerArgs) -> "DiffGenerator": + """ + Create a DiffGenerator with the specified arguments. + + Args: + server_args: The inference arguments + + Returns: + The created DiffGenerator + """ + instance = cls( + server_args=server_args, + ) + is_local_mode = server_args.is_local_mode + logger.info(f"Local mode: {is_local_mode}") + if is_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()}." + ) + + def post_process_sample( + self, + sample: torch.Tensor, + data_type: DataType, + fps: int, + save_output: bool = True, + save_file_path: str = None, + ): + """ + Process a single sample output and save output if necessary + """ + # Process outputs + if sample.dim() == 3: + # for images, dim t is missing + sample = sample.unsqueeze(1) + sample = rearrange(sample, "c t h w -> t c h w") + frames = [] + # TODO: this can be batched + for x in sample: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + + # Save outputs if requested + if save_output: + if save_file_path: + os.makedirs(os.path.dirname(save_file_path), exist_ok=True) + with suppress_other_loggers(): + if data_type == DataType.VIDEO: + imageio.mimsave( + save_file_path, + frames, + fps=fps, + format=data_type.get_default_extension(), + ) + else: + imageio.imwrite(save_file_path, frames[0]) + logger.info("Saved output to %s", save_file_path) + else: + logger.warning("No output path provided, output not saved") + + return frames + + def generate( + self, + sampling_params_kwargs: dict | None = None, + ) -> dict[str, Any] | list[np.ndarray] | list[dict[str, Any]] | None: + """ + Generate a image/video based on the given prompt. + + Args: + + Returns: + Either the output dictionary, list of frames, or list of results for batch processing + """ + # 1. prepare requests + prompt = sampling_params_kwargs.get("prompt", None) + prompts: list[str] = [] + # Handle batch processing from text file + if self.server_args.prompt_file_path is not None: + prompt_txt_path = self.server_args.prompt_file_path + if not os.path.exists(prompt_txt_path): + raise FileNotFoundError( + f"Prompt text file not found: {prompt_txt_path}" + ) + # Read prompts from file + with open(prompt_txt_path, encoding="utf-8") as f: + prompts.extend(line.strip() for line in f if line.strip()) + + if not prompts: + raise ValueError(f"No prompts found in file: {prompt_txt_path}") + + logger.info("Found %d prompts in %s", len(prompts), prompt_txt_path) + elif prompt is not None: + if isinstance(prompt, str): + prompts.append(prompt) + elif isinstance(prompt, list): + prompts.extend(prompt) + else: + raise ValueError("Either prompt or prompt_txt must be provided") + + sampling_params = SamplingParams.from_user_sampling_params_args( + self.server_args.model_path, + server_args=self.server_args, + **sampling_params_kwargs, + ) + + requests: list[Req] = [] + for output_idx, p in enumerate(prompts): + sampling_params.prompt = p + requests.append( + prepare_request( + server_args=self.server_args, + sampling_params=sampling_params, + ) + ) + + results = [] + 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: + logger.error( + "Received empty output from scheduler for prompt %d", + request_idx + 1, + ) + continue + for output_idx, sample in enumerate(output_batch.output): + num_outputs = len(output_batch.output) + frames = self.post_process_sample( + sample, + fps=req.fps, + save_output=req.save_output, + save_file_path=req.output_file_path( + num_outputs, output_idx + ), + data_type=req.data_type, + ) + + result_item: dict[str, Any] = { + "samples": sample, + "frames": frames, + "prompts": req.prompt, + "size": (req.height, req.width, req.num_frames), + "generation_time": timer.duration, + "timings": ( + output_batch.timings.to_dict() + if output_batch.timings + else {} + ), + "trajectory": output_batch.trajectory_latents, + "trajectory_timesteps": output_batch.trajectory_timesteps, + "trajectory_decoded": output_batch.trajectory_decoded, + "prompt_index": output_idx, + } + results.append(result_item) + except Exception: + continue + + total_gen_time = time.perf_counter() - total_start_time + log_batch_completion(logger, len(results), total_gen_time) + + if len(results) == 0: + return None + else: + if requests[0].return_frames: + results = [r["frames"] for r in results] + if len(results) == 1: + return results[0] + return results + + 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 isinstance(response, dict) and response.get("status") == "ok": + logger.info(success_msg) + else: + error_msg = ( + response.get("message", "Unknown error") + if isinstance(response, dict) + else "Unknown response format" + ) + raise RuntimeError(f"{failure_msg}: {error_msg}") + + def set_lora(self, lora_nickname: str, lora_path: str | None = None) -> None: + req = SetLoraReq(lora_nickname=lora_nickname, lora_path=lora_path) + self._send_lora_request( + req, + f"Successfully set LoRA adapter: {lora_nickname}", + "Failed to set LoRA adapter", + ) + + def unmerge_lora_weights(self) -> None: + req = UnmergeLoraWeightsReq() + self._send_lora_request( + req, + "Successfully unmerged LoRA weights", + "Failed to unmerge LoRA weights", + ) + self._is_lora_merged = False + + def merge_lora_weights(self) -> None: + req = MergeLoraWeightsReq() + self._send_lora_request( + req, "Successfully merged LoRA weights", "Failed to merge LoRA weights" + ) + self._is_lora_merged = True + + def _ensure_lora_state( + self, + lora_path: str | None, + lora_nickname: str | None = None, + merge_lora: bool = True, + ) -> None: + if lora_path is None: + if self._is_lora_merged: + self.unmerge_lora_weights() + self._current_lora_path = None + self._current_lora_nickname = None + self._is_lora_merged = False + return + + lora_nickname = lora_nickname or self.server_args.lora_nickname + + if self._current_lora_path != lora_path: + if self._is_lora_merged: + self.unmerge_lora_weights() + self._is_lora_merged = False + self.set_lora(lora_nickname, lora_path) + self._current_lora_path = lora_path + self._current_lora_nickname = lora_nickname + self._is_lora_merged = False + + if merge_lora and not self._is_lora_merged: + self.merge_lora_weights() + elif not merge_lora: + self._is_lora_merged = False + + 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( + 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: + logger.info("Waiting for local worker processes to terminate...") + 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py new file mode 100644 index 00000000..489a08e2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -0,0 +1,80 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +import asyncio +from contextlib import asynccontextmanager + +from fastapi import APIRouter, FastAPI + +from sglang.multimodal_gen.runtime.entrypoints.openai import image_api, video_api +from sglang.multimodal_gen.runtime.server_args import ServerArgs, prepare_server_args +from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger + + +@asynccontextmanager +async def lifespan(app: FastAPI): + from sglang.multimodal_gen.runtime.scheduler_client import ( + run_zeromq_broker, + scheduler_client, + ) + + # 1. Initialize the singleton client that connects to the backend Scheduler + server_args = app.state.server_args + 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 + print("FastAPI app is shutting down...") + broker_task.cancel() + scheduler_client.close() + + +# Health router +health_router = APIRouter() + + +@health_router.get("/health") +async def health(): + return {"status": "ok"} + + +@health_router.get("/health_generate") +async def health_generate(): + # TODO : health generate endpoint + return {"status": "ok"} + + +def create_app(server_args: ServerArgs): + """ + Create and configure the FastAPI application instance. + """ + app = FastAPI(lifespan=lifespan) + + app.include_router(health_router) + + from sglang.multimodal_gen.runtime.entrypoints.openai import common_api + + app.include_router(common_api.router) + app.include_router(image_api.router) + app.include_router(video_api.router) + + app.state.server_args = server_args + return app + + +if __name__ == "__main__": + import uvicorn + + server_args = prepare_server_args([]) + configure_logger(server_args) + app = create_app(server_args) + uvicorn.run( + app, + host=server_args.host, + port=server_args.port, + log_config=None, + reload=False, # Set to True during development for auto-reloading + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py new file mode 100644 index 00000000..2dcbeb77 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py @@ -0,0 +1,62 @@ +from typing import Any, Optional + +from fastapi import APIRouter, Body, HTTPException + +from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( + MergeLoraWeightsReq, + SetLoraReq, + UnmergeLoraWeightsReq, +) +from sglang.multimodal_gen.runtime.scheduler_client import scheduler_client +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +router = APIRouter(prefix="/v1") +logger = init_logger(__name__) + + +async def _handle_lora_request(req: Any, success_msg: str, failure_msg: str): + try: + response = await scheduler_client.forward(req) + if isinstance(response, dict) and response.get("status") == "ok": + return {"status": "ok", "message": success_msg} + else: + error_msg = ( + response.get("message", "Unknown error") + if isinstance(response, dict) + else "Unknown response format" + ) + 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: str = Body(..., embed=True), + lora_path: Optional[str] = Body(None, embed=True), +): + req = SetLoraReq(lora_nickname=lora_nickname, lora_path=lora_path) + return await _handle_lora_request( + req, + f"Successfully set LoRA adapter: {lora_nickname}", + "Failed to set LoRA adapter", + ) + + +@router.post("/merge_lora_weights") +async def merge_lora_weights(): + req = MergeLoraWeightsReq() + return await _handle_lora_request( + req, "Successfully merged LoRA weights", "Failed to merge LoRA weights" + ) + + +@router.post("/unmerge_lora_weights") +async def unmerge_lora_weights(): + req = UnmergeLoraWeightsReq() + return await _handle_lora_request( + req, "Successfully unmerged LoRA weights", "Failed to unmerge LoRA weights" + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py new file mode 100644 index 00000000..7ac86c2c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -0,0 +1,228 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +import base64 +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 ( + SamplingParams, + generate_request_id, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( + ImageGenerationsRequest, + ImageResponse, + ImageResponseData, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.stores import IMAGE_STORE +from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( + _parse_size, + _save_upload_to_path, + process_generation_batch, +) +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.scheduler_client import 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 _choose_ext(output_format: Optional[str], background: Optional[str]) -> str: + # Normalize and choose extension + fmt = (output_format or "").lower() + if fmt in {"png", "webp", "jpeg", "jpg"}: + return "jpg" if fmt == "jpeg" else fmt + # If transparency requested, prefer png + if (background or "auto").lower() == "transparent": + return "png" + # Default + return "jpg" + + +def _build_sampling_params_from_request( + request_id: str, + prompt: str, + n: int, + size: Optional[str], + output_format: Optional[str], + background: Optional[str], + image_path: Optional[str] = None, +) -> SamplingParams: + if size is None: + width, height = None, None + else: + width, height = _parse_size(size) + ext = _choose_ext(output_format, background) + server_args = get_global_server_args() + # Build user params + sampling_params = SamplingParams.from_user_sampling_params_args( + model_path=server_args.model_path, + request_id=request_id, + prompt=prompt, + image_path=image_path, + num_frames=1, # image + width=width, + height=height, + num_outputs_per_prompt=max(1, min(int(n or 1), 10)), + save_output=True, + server_args=server_args, + output_file_name=f"{request_id}.{ext}", + ) + return sampling_params + + +def _build_req_from_sampling(s: SamplingParams) -> Req: + return Req( + request_id=s.request_id, + data_type=s.data_type, + prompt=s.prompt, + image_path=s.image_path, + height=s.height, + width=s.width, + fps=1, + num_frames=s.num_frames, + seed=s.seed, + output_path=s.output_path, + output_file_name=s.output_file_name, + num_outputs_per_prompt=s.num_outputs_per_prompt, + save_output=s.save_output, + ) + + +@router.post("/generations", response_model=ImageResponse) +async def generations( + request: ImageGenerationsRequest, +): + request_id = generate_request_id() + sampling = _build_sampling_params_from_request( + request_id=request_id, + prompt=request.prompt, + n=request.n or 1, + size=request.size, + output_format=request.output_format, + background=request.background, + ) + batch = prepare_request( + server_args=get_global_server_args(), + sampling_params=sampling, + ) + # Run synchronously for images and save to disk + save_file_path = await process_generation_batch(scheduler_client, batch) + + await IMAGE_STORE.upsert( + request_id, + { + "id": request_id, + "created_at": int(time.time()), + "file_path": save_file_path, + }, + ) + + resp_format = (request.response_format or "b64_json").lower() + if resp_format == "b64_json": + with open(save_file_path, "rb") as f: + b64 = base64.b64encode(f.read()).decode("utf-8") + return ImageResponse( + data=[ + ImageResponseData( + b64_json=b64, + revised_prompt=request.prompt, + ) + ] + ) + else: + # Return error, not supported + raise HTTPException( + status_code=400, detail="response_format=url is not supported" + ) + + +@router.post("/edits", response_model=ImageResponse) +async def edits( + image: Optional[List[UploadFile]] = File(None), + image_array: Optional[List[UploadFile]] = File(None, alias="image[]"), + 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"), + user: Optional[str] = Form(None), +): + request_id = generate_request_id() + # Resolve images from either `image` or `image[]` (OpenAI SDK sends `image[]` when list is provided) + images = image or image_array + if not images or len(images) == 0: + raise HTTPException(status_code=422, detail="Field 'image' is required") + + # Save first input image; additional images or mask are not yet used by the pipeline + uploads_dir = os.path.join("outputs", "uploads") + os.makedirs(uploads_dir, exist_ok=True) + first_image = images[0] + input_path = os.path.join(uploads_dir, f"{request_id}_{first_image.filename}") + await _save_upload_to_path(first_image, input_path) + + sampling = _build_sampling_params_from_request( + request_id=request_id, + prompt=prompt, + n=n or 1, + size=size, + output_format=output_format, + background=background, + image_path=input_path, + ) + batch = _build_req_from_sampling(sampling) + + save_file_path = await process_generation_batch(scheduler_client, batch) + + await IMAGE_STORE.upsert( + request_id, + { + "id": request_id, + "created_at": int(time.time()), + "file_path": save_file_path, + }, + ) + + # Default to b64_json to align with gpt-image-1 behavior in OpenAI examples + if (response_format or "b64_json").lower() == "b64_json": + with open(save_file_path, "rb") as f: + b64 = base64.b64encode(f.read()).decode("utf-8") + return ImageResponse( + data=[ImageResponseData(b64_json=b64, revised_prompt=prompt)] + ) + else: + url = f"/v1/images/{request_id}/content" + return ImageResponse(data=[ImageResponseData(url=url, revised_prompt=prompt)]) + + +@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") + + file_path = item.get("file_path") + if not file_path or 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) + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py new file mode 100644 index 00000000..00800ab1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py @@ -0,0 +1,65 @@ +import time +from typing import Any, Dict, List, Optional + +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 + + +class ImageResponse(BaseModel): + created: int = Field(default_factory=lambda: int(time.time())) + data: List[ImageResponseData] + + +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 + + +# 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 = "720x1280" + seconds: str = "4" + quality: str = "standard" + remixed_from_video_id: Optional[str] = None + completed_at: Optional[int] = None + expires_at: Optional[int] = None + error: Optional[Dict[str, Any]] = None + + +class VideoGenerationsRequest(BaseModel): + prompt: str + input_reference: Optional[str] = None + model: Optional[str] = None + seconds: Optional[int] = 4 + size: Optional[str] = "720x1280" + fps: Optional[int] = None + num_frames: Optional[int] = None + + +class VideoListResponse(BaseModel): + data: List[VideoResponse] + object: str = "list" + + +class VideoRemixRequest(BaseModel): + prompt: str diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/stores.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/stores.py new file mode 100644 index 00000000..f924de81 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/stores.py @@ -0,0 +1,46 @@ +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 +VIDEO_STORE = AsyncDictStore() +IMAGE_STORE = AsyncDictStore() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py new file mode 100644 index 00000000..ec7a3cb1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -0,0 +1,125 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +import dataclasses +import os +import time +from typing import Optional + +import imageio +import numpy as np +import torch +import torchvision +from einops import rearrange +from fastapi import UploadFile + +from sglang.multimodal_gen.configs.sample.sampling_params import DataType +from sglang.multimodal_gen.runtime.utils.logging_utils import ( + init_logger, + log_batch_completion, + log_generation_timer, +) + +logger = init_logger(__name__) + + +@dataclasses.dataclass +class SetLoraReq: + lora_nickname: str + lora_path: Optional[str] = None + + +@dataclasses.dataclass +class MergeLoraWeightsReq: + pass + + +@dataclasses.dataclass +class UnmergeLoraWeightsReq: + pass + + +def post_process_sample( + sample: torch.Tensor, + data_type: DataType, + fps: int, + save_output: bool = True, + save_file_path: str = None, +): + """ + Process sample output and save video if necessary + """ + # Process outputs + if sample.dim() == 3: + # for images, dim t is missing + sample = sample.unsqueeze(1) + videos = rearrange(sample, "c t h w -> t c h w") + frames = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + + # 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: + imageio.mimsave( + save_file_path, + frames, + fps=fps, + format=data_type.get_default_extension(), + ) + else: + imageio.imwrite(save_file_path, frames[0]) + logger.info(f"Saved output to {save_file_path}") + else: + logger.info(f"No output path provided, output not saved") + + return frames + + +def _parse_size(size: str) -> tuple[int, int]: + 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: + # Fallback to default portrait 720x1280 + return 720, 1280 + + +# 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 process_generation_batch( + scheduler_client, + batch, +): + total_start_time = time.perf_counter() + with log_generation_timer(logger, batch.prompt): + result = await scheduler_client.forward([batch]) + + if result.output is None: + raise RuntimeError("Model generation returned no output.") + + save_file_path = str(os.path.join(batch.output_path, batch.output_file_name)) + post_process_sample( + result.output[0], + batch.data_type, + batch.fps, + batch.save_output, + save_file_path, + ) + + total_time = time.perf_counter() - total_start_time + log_batch_completion(logger, 1, total_time) + + return save_file_path diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py new file mode 100644 index 00000000..c346fc52 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -0,0 +1,265 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +import asyncio +import json +import os +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.stores import VIDEO_STORE +from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( + _parse_size, + _save_upload_to_path, + process_generation_batch, +) +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"]) + + +# NOTE(mick): the sampling params needs to be further adjusted +# FIXME: duplicated with the one in `image_api.py` +def _build_sampling_params_from_request( + request_id: str, request: VideoGenerationsRequest +) -> SamplingParams: + if request.size is None: + width, height = None, None + else: + width, height = _parse_size(request.size) + seconds = request.seconds if request.seconds is not None else 4 + # Prefer user-provided fps/num_frames from request; fallback to defaults + fps_default = 24 + fps = request.fps if request.fps is not None else fps_default + # If user provides num_frames, use it directly; otherwise derive from seconds * fps + derived_num_frames = fps * seconds + num_frames = ( + request.num_frames if request.num_frames is not None else derived_num_frames + ) + server_args = get_global_server_args() + sampling_params = SamplingParams.from_user_sampling_params_args( + model_path=server_args.model_path, + request_id=request_id, + prompt=request.prompt, + num_frames=num_frames, + fps=fps, + width=width, + height=height, + image_path=request.input_reference, + save_output=True, + server_args=server_args, + output_file_name=request_id, + ) + + return sampling_params + + +# 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": sampling.output_file_path(), + } + + +async def _dispatch_job_async(job_id: str, batch: Req) -> None: + from sglang.multimodal_gen.runtime.scheduler_client import scheduler_client + + try: + await process_generation_batch(scheduler_client, batch) + await VIDEO_STORE.update_fields( + job_id, + {"status": "completed", "progress": 100, "completed_at": int(time.time())}, + ) + except Exception as e: + logger.error(f"{e}") + await VIDEO_STORE.update_fields( + job_id, {"status": "failed", "error": {"message": str(e)}} + ) + + +# 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), + 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), + extra_body: Optional[str] = Form(None), +): + content_type = request.headers.get("content-type", "").lower() + request_id = generate_request_id() + + if "multipart/form-data" in content_type: + if not prompt: + raise HTTPException(status_code=400, detail="prompt is required") + if input_reference is None: + raise HTTPException( + status_code=400, detail="input_reference file is required" + ) + + uploads_dir = os.path.join("outputs", "uploads") + os.makedirs(uploads_dir, exist_ok=True) + input_path = os.path.join( + uploads_dir, f"{request_id}_{input_reference.filename}" + ) + await _save_upload_to_path(input_reference, input_path) + + # 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 or "720x1280", + fps=fps_val, + num_frames=num_frames_val, + ) + 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) + req = VideoGenerationsRequest(**payload) + except Exception as e: + raise HTTPException(status_code=400, detail=f"Invalid request body: {e}") + + logger.debug(f"Server received from create_video endpoint: req={req}") + + sampling_params = _build_sampling_params_from_request(request_id, req) + 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=get_global_server_args(), + sampling_params=sampling_params, + ) + # Enqueue the job asynchronously and return immediately + asyncio.create_task(_dispatch_job_async(request_id, batch)) + 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") + + 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) + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/utils.py new file mode 100644 index 00000000..9cd56f1e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -0,0 +1,42 @@ +# 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. +""" + + +from sglang.multimodal_gen.configs.sample.sampling_params import 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 init_logger +from sglang.multimodal_gen.utils import shallow_asdict + +logger = init_logger(__name__) + + +def prepare_request( + server_args: ServerArgs, + sampling_params: SamplingParams, +) -> Req: + """ + Settle SamplingParams according to ServerArgs + + """ + # Create a copy of inference args to avoid modifying the original + req = Req( + **shallow_asdict(sampling_params), + VSA_sparsity=server_args.VSA_sparsity, + ) + req.adjust_size(server_args) + + if req.width <= 0 or req.height <= 0: + raise ValueError( + f"Height, width must be positive integers, got " + f"height={req.height}, width={req.width}" + ) + + return req diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/launch_server.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/launch_server.py new file mode 100644 index 00000000..0f34166a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/launch_server.py @@ -0,0 +1,142 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +import multiprocessing as mp + +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, set_global_server_args +from sglang.multimodal_gen.runtime.utils.logging_utils import ( + configure_logger, + logger, + suppress_other_loggers, +) + + +def launch_server(server_args: ServerArgs, launch_http_server: bool = True): + """ + Args: + launch_http_server: False for offline local mode + """ + configure_logger(server_args) + suppress_other_loggers() + + # 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.") + + # set for endpoints to access global_server_args + set_global_server_args(server_args) + + app = create_app(server_args) + uvicorn.run( + app, + log_config=None, + log_level=server_args.log_level, + host=server_args.host, + port=server_args.port, + reload=False, + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/__init__.py new file mode 100644 index 00000000..af2eb7d1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/__init__.py @@ -0,0 +1 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/activation.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/activation.py new file mode 100644 index 00000000..4eff9ba1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/activation.py @@ -0,0 +1,129 @@ +# 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 + +# 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, *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.silu(x[..., :d]) * x[..., d:] + + +@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]() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py new file mode 100644 index 00000000..9635a674 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py @@ -0,0 +1,414 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py new file mode 100644 index 00000000..1b40782b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/__init__.py @@ -0,0 +1,28 @@ +# 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 + +__all__ = [ + "USPAttention", + "LocalAttention", + "UlyssesAttention", + "UlyssesAttention_VSA", + "AttentionBackend", + "AttentionMetadata", + "AttentionMetadataBuilder", + # "AttentionState", + "get_attn_backend", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/__init__.py new file mode 100644 index 00000000..af2eb7d1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/__init__.py @@ -0,0 +1 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py new file mode 100644 index 00000000..b96aad6a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py @@ -0,0 +1,101 @@ +# 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, +) + + +class AITerBackend(AttentionBackend): + """ + Backend for AITemplate attention implementation. + """ + + @staticmethod + def get_name() -> str: + return "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: + super().__init__( + num_heads=num_heads, + head_size=head_size, + softmax_scale=softmax_scale, + causal=causal, + num_kv_heads=num_kv_heads, + prefix=prefix, + **extra_impl_args, + ) + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py new file mode 100644 index 00000000..3463ef05 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py @@ -0,0 +1,180 @@ +# 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 + + +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_name() -> str: + 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 + + Args: + qkv: The query-key-value tensor + attn_metadata: Metadata for the attention operation + + Returns: + Processed QKV tensor + """ + 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 + + Args: + output: The output tensor from the attention operation + attn_metadata: Metadata for the attention operation + + Returns: + Postprocessed output tensor + """ + + return output + + @abstractmethod + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: T, + ) -> torch.Tensor: + raise NotImplementedError diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py new file mode 100644 index 00000000..2394de5f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py @@ -0,0 +1,139 @@ +# 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 + +import torch + +from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context + +try: + from sgl_kernel.flash_attn import flash_attn_varlen_func + + # flash_attn 3 no longer have a different API, see following commit: + # https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7 + flash_attn_func = flash_attn_varlen_func +except ImportError as e: + raise e + +from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +fa_ver = 3 + + +def set_fa_ver(ver: int): + 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): + pass + + def prepare(self): + 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_name() -> str: + return "FLASH_ATTN" + + @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.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] + 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=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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py new file mode 100644 index 00000000..df795e06 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py @@ -0,0 +1,78 @@ +# 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.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_name() -> str: + return "FA" + + @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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py new file mode 100644 index 00000000..3563ddd1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py @@ -0,0 +1,70 @@ +# 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.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_name() -> str: + return "SAGE_ATTN" + + @staticmethod + def get_impl_cls() -> type["SageAttentionImpl"]: + return SageAttentionImpl + + # @staticmethod + # def get_metadata_cls() -> Type["AttentionMetadata"]: + # return FlashAttentionMetadata + + +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, + ) -> 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, + ) + return output diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py new file mode 100644 index 00000000..fd5b6f2b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py @@ -0,0 +1,78 @@ +# 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.sageattn.api import ( + sageattn_blackwell, +) +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_name() -> str: + return "SAGE_ATTN_THREE" + + @staticmethod + def get_impl_cls() -> type["SageAttention3Impl"]: + return SageAttention3Impl + + @staticmethod + def get_metadata_cls() -> type["AttentionMetadata"]: + raise NotImplementedError + + @staticmethod + def get_builder_cls() -> type["AttentionMetadataBuilder"]: + raise NotImplementedError + + # @staticmethod + # def get_metadata_cls() -> Type["AttentionMetadata"]: + # return FlashAttentionMetadata + + +class SageAttention3Impl(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: + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + output = sageattn_blackwell(query, key, value, is_causal=self.causal) + output = output.transpose(1, 2) + return output diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py new file mode 100644 index 00000000..bfa3b430 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py @@ -0,0 +1,77 @@ +# 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.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_name() -> str: + return "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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py new file mode 100644 index 00000000..6db3785f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py @@ -0,0 +1,313 @@ +# 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 + +import sglang.multimodal_gen.envs as envs +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.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_name() -> str: + return "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? + config_file = envs.SGLANG_DIFFUSION_ATTENTION_CONFIG + if config_file is None: + raise ValueError("SGLANG_DIFFUSION_ATTENTION_CONFIG is not set") + + # TODO(kevin): get mask strategy for different STA modes + with open(config_file) 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py new file mode 100644 index 00000000..6fe34292 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py @@ -0,0 +1,331 @@ +# 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.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_name() -> str: + return "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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py new file mode 100644 index 00000000..5709601d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py @@ -0,0 +1,258 @@ +# 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.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class VMOBAAttentionBackend(AttentionBackend): + + accept_output_buffer: bool = True + + @staticmethod + def get_name() -> str: + return "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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/layer.py new file mode 100644 index 00000000..df4f377d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -0,0 +1,396 @@ +# 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_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 ( + backend_name_to_enum, + 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 = backend_name_to_enum(attn_backend.get_name()) + self.dtype = dtype + + @torch.compiler.disable + 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.""" + + @torch.compiler.disable + 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 = backend_name_to_enum(attn_backend.get_name()) + 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, + **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: 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 = backend_name_to_enum(attn_backend.get_name()) + self.dtype = dtype + self.causal = causal + self.dropout_p = dropout_rate + + 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, + ) -> torch.Tensor: + """ + Forward pass for USPAttention. + + q, k, v: [B, S_local, H, D] + + Note: Replicated tensors are not supported in this implementation. + """ + assert ( + replicated_q is None and replicated_k is None and replicated_v is None + ), "USPAttention does not support replicated_qkv." + forward_context: ForwardContext = get_forward_context() + ctx_attn_metadata = forward_context.attn_metadata + if 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 + + # Ulysses-style All-to-All for sequence/head sharding + if get_ulysses_parallel_world_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 get_ulysses_parallel_world_size() > 1: + # -> [B, S_local, H, D] + out = _usp_output_all_to_all(out, head_dim=2) + + return out diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/selector.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/selector.py new file mode 100644 index 00000000..b5d589f7 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/attention/selector.py @@ -0,0 +1,197 @@ +# 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 not None: + # Sort the backend names to ensure consistent cache key + be_tuple = tuple( + sorted(list(supported_attention_backends), key=lambda b: b.name) + ) + else: + be_tuple = None + 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] | None = None, +) -> type[AttentionBackend]: + # Check whether a particular choice of backend was + # previously forced. + # + # THIS SELECTION OVERRIDES THE SGLANG_DIFFUSION_ATTENTION_BACKEND + # ENVIRONMENT VARIABLE. + from sglang.multimodal_gen.runtime.platforms import current_platform + + supported_attention_backends = set(supported_attention_backends) + if not supported_attention_backends: + raise ValueError("supported_attention_backends is empty") + 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 selected_backend is None: + logger.debug(f"Attention backend not specified") + elif ( + not supported_attention_backends + or 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/custom_op.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/custom_op.py new file mode 100644 index 00000000..abc2f123 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/custom_op.py @@ -0,0 +1,110 @@ +# 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.utils.common import ( + is_cpu, + is_cuda, + is_hip, + is_npu, + is_xpu, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +_is_cuda = is_cuda() +_is_hip = is_hip() +_is_cpu = is_cpu() +_is_npu = is_npu() +_is_xpu = is_xpu() + + +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_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_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 dispatch_forward(self) -> Callable: + if _is_cuda: + return self.forward_cuda + elif _is_hip: + return self.forward_hip + elif _is_npu: + return self.forward_npu + elif _is_xpu: + return self.forward_xpu + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/layernorm.py new file mode 100644 index 00000000..ec8f680a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -0,0 +1,429 @@ +# 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.layers.custom_op import CustomOp +from sglang.multimodal_gen.runtime.layers.triton_ops import ( + fuse_scale_shift_kernel, + norm_infer, + rms_norm_fn, +) +from sglang.multimodal_gen.runtime.utils.common import ( + get_bool_env_var, + is_cpu, + is_cuda, + is_hip, + is_npu, + is_xpu, +) + +_is_cuda = is_cuda() +_is_hip = is_hip() +_is_npu = is_npu() +_is_cpu = is_cpu() +_is_xpu = is_xpu() + +from sgl_kernel import fused_add_rmsnorm, rmsnorm + + +# 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 + 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) + 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: + 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 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") + 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 extra_repr(self) -> str: + s = f"hidden_size={self.weight.data.size(0)}" + s += f", eps={self.variance_epsilon}" + return s + + +class ScaleResidual(nn.Module): + """ + Applies gated residual connection. + """ + + def __init__(self, prefix: str = ""): + super().__init__() + + def forward( + self, residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor + ) -> torch.Tensor: + """Apply gated residual connection.""" + # x.shape: [batch_size, seq_len, inner_dim] + 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 + return residual + ( + x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate + ).flatten(1, 2) + else: + # gate.shape: [batch_size, 1, inner_dim] + return residual + x * gate + + +# 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 + return F.layer_norm( + inputs.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ).to(origin_dtype) + + +class ScaleResidualLayerNormScaleShift(nn.Module): + """ + Fused operation that combines: + 1. Gated residual connection + 2. LayerNorm + 3. Scale and shift operations + + This reduces memory bandwidth by combining memory-bound operations. + """ + + def __init__( + self, + hidden_size: int, + norm_type: str = "rms", + eps: float = 1e-6, + elementwise_affine: bool = False, + dtype: torch.dtype = torch.float32, + compute_dtype: torch.dtype | None = None, + prefix: str = "", + ): + super().__init__() + if norm_type == "rms": + self.norm = RMSNorm( + hidden_size, has_weight=elementwise_affine, eps=eps, dtype=dtype + ) + elif norm_type == "layer": + if compute_dtype == torch.float32: + self.norm = FP32LayerNorm( + hidden_size, elementwise_affine=elementwise_affine, eps=eps + ) + else: + self.norm = LayerNorm( + hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + dtype=dtype, + ) + else: + raise NotImplementedError(f"Norm type {norm_type} not implemented") + + def forward( + self, + residual: torch.Tensor, + x: torch.Tensor, + gate: torch.Tensor | int, + shift: torch.Tensor, + scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Apply gated residual connection, followed by layernorm and + scale/shift in a single fused operation. + + Returns: + Tuple containing: + - normalized and modulated output of shape: [batch_size, seq_len, inner_dim] + - residual value (value after residual connection + but before normalization) + """ + # x.shape: [batch_size, seq_len, inner_dim] + # Apply residual connection with gating + 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: + # used by bidirectional self attention + # gate.shape: [batch_size, 1, inner_dim] + residual_output = residual + x * gate + else: + raise ValueError(f"Gate type {type(gate)} not supported") + # residual_output.shape: [batch_size, seq_len, inner_dim] + + # Apply normalization + normalized = self.norm(residual_output) + + # modulated = fused_scale_shift( + # normalized, + # scale, + # shift, + # ) + modulated = fuse_scale_shift_kernel( + normalized, + scale, + shift, + ) + return modulated, residual_output + + +class LayerNormScaleShift(nn.Module): + """ + Fused operation that combines LayerNorm with scale and shift operations. + This reduces memory bandwidth by combining memory-bound operations. + """ + + def __init__( + self, + hidden_size: int, + norm_type: str = "rms", + eps: float = 1e-6, + elementwise_affine: bool = False, + dtype: torch.dtype = torch.float32, + compute_dtype: torch.dtype | None = None, + prefix: str = "", + ): + super().__init__() + self.compute_dtype = compute_dtype + if norm_type == "rms": + self.norm = RMSNorm(hidden_size, has_weight=elementwise_affine, eps=eps) + elif norm_type == "layer": + if self.compute_dtype == torch.float32: + self.norm = FP32LayerNorm( + hidden_size, elementwise_affine=elementwise_affine, eps=eps + ) + else: + self.norm = nn.LayerNorm( + hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + dtype=dtype, + ) + else: + raise NotImplementedError(f"Norm type {norm_type} not implemented") + + def forward( + self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor + ) -> torch.Tensor: + """Apply ln followed by scale and shift in a single fused operation.""" + # x.shape: [batch_size, seq_len, inner_dim] + normalized = self.norm(x) + if self.compute_dtype == torch.float32: + normalized = normalized.float() + + if scale.dim() == 4: + # scale.shape: [batch_size, num_frames, 1, inner_dim] + num_frames = scale.shape[1] + frame_seqlen = normalized.shape[1] // num_frames + output = ( + normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) + * (1.0 + scale) + + shift + ).flatten(1, 2) + else: + # scale.shape: [batch_size, 1, inner_dim] + # shift.shape: [batch_size, 1, inner_dim] + output = normalized * (1.0 + scale) + shift + + if self.compute_dtype == torch.float32: + output = output.to(x.dtype) + + return output diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/linear.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/linear.py new file mode 100644 index 00000000..65c71372 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/linear.py @@ -0,0 +1,1057 @@ +# 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.nn.functional as F +from torch.nn.parameter import Parameter + +from sglang.multimodal_gen.runtime.distributed import ( + divide, + get_tp_rank, + get_tp_world_size, + split_tensor_along_last_dim, + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, +) +from sglang.multimodal_gen.runtime.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) + +# 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.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 torch.cuda.is_available() 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. + 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. + """ + + 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 = "", + ): + # Divide the weight matrix along the last dimension. + self.tp_size = get_tp_world_size() + 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 = get_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) + 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={get_tp_world_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 = "", + ): + self.output_sizes = output_sizes + tp_size = get_tp_world_size() + assert all(output_size % tp_size == 0 for output_size in output_sizes) + 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, + ) + + 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 = get_tp_rank() + tp_size = get_tp_world_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 = get_tp_world_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 = "", + ): + 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_size = get_tp_world_size() + 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, + ) + + 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 = get_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 = "", + ): + # Divide the weight matrix along the first dimension. + self.tp_rank = get_tp_rank() + self.tp_size = get_tp_world_size() + 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 = get_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 = get_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) + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/lora/linear.py new file mode 100644 index 00000000..63092612 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/lora/linear.py @@ -0,0 +1,433 @@ +# 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 + self.cpu_weight = base_layer.weight.to("cpu") + # 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) + self.disable_lora: bool = False + self.lora_rank = lora_rank + self.lora_alpha = lora_alpha + self.lora_path: str | None = None + + self.lora_A = None + self.lora_B = 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() + + 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 + out, output_bias = self.base_layer(x) + return out + delta, output_bias + else: + out, output_bias = self.base_layer(x) + return out.to(x), 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, + ) -> None: + self.lora_A = torch.nn.Parameter( + A + ) # share storage with weights in the pipeline + self.lora_B = torch.nn.Parameter(B) + self.disable_lora = False + self.merge_lora_weights() + self.lora_path = lora_path + + @torch.no_grad() + def merge_lora_weights(self) -> None: + if self.disable_lora: + return + + if self.merged: + self.unmerge_lora_weights() + assert ( + self.lora_A is not None and self.lora_B is not None + ), "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() + data += self.slice_lora_b_weights(self.lora_B).to( + data + ) @ self.slice_lora_a_weights(self.lora_A).to(data) + 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()) + data += self.slice_lora_b_weights( + self.lora_B.to(data) + ) @ self.slice_lora_a_weights(self.lora_A.to(data)) + 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 + self.base_layer.weight = nn.Parameter( + self.cpu_weight.to(device, non_blocking=True) + ) + else: + self.base_layer.weight.data = self.cpu_weight.data.to( + self.base_layer.weight, non_blocking=True + ) + + 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.to(self.base_layer.weight) + + 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.set_lora: + output_parallel = self.apply_lora(output_parallel, 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() + + 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 + # 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/mlp.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/mlp.py new file mode 100644 index 00000000..17918e2a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/mlp.py @@ -0,0 +1,46 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn as nn + +from sglang.multimodal_gen.runtime.layers.activation import get_act_fn +from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear + + +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 = "", + ): + super().__init__() + self.fc_in = ReplicatedLinear( + input_dim, + mlp_hidden_dim, # For activation func like SiLU that need 2x width + bias=bias, + params_dtype=dtype, + ) + + self.act = get_act_fn(act_type) + if output_dim is None: + output_dim = input_dim + self.fc_out = ReplicatedLinear( + mlp_hidden_dim, output_dim, bias=bias, params_dtype=dtype + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.fc_in(x) + x = self.act(x) + x, _ = self.fc_out(x) + return x diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py new file mode 100644 index 00000000..0d6c7979 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py @@ -0,0 +1,71 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +from typing import Literal, get_args + +from sglang.multimodal_gen.runtime.layers.quantization.base_config import ( + QuantizationConfig, +) + +QuantizationMethods = Literal[None] + +QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) + +# The customized quantization methods which will be added to this dict. +_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {} + + +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. + + Examples: + >>> from sglang.multimodal_gen.runtime.layers.quantization import register_quantization_config + >>> from sglang.multimodal_gen.runtime.layers.quantization import get_quantization_config + >>> from sglang.multimodal_gen.runtime.layers.quantization.base_config import QuantizationConfig + >>> + >>> @register_quantization_config("my_quant") + ... class MyQuantConfig(QuantizationConfig): + ... pass + >>> + >>> get_quantization_config("my_quant") + + """ # 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", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/quantization/base_config.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/quantization/base_config.py new file mode 100644 index 00000000..ffb275a8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/quantization/base_config.py @@ -0,0 +1,152 @@ +# 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.""" + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py new file mode 100644 index 00000000..c0a58903 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py @@ -0,0 +1,889 @@ +# 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.""" +import functools +from collections import OrderedDict +from typing import Any + +import torch + +from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group +from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp +from sglang.multimodal_gen.runtime.layers.triton_ops import apply_rotary_embedding +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +def _rotate_neox(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _rotate_gptj(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., ::2] + x2 = x[..., 1::2] + x = torch.stack((-x2, x1), dim=-1) + return x.flatten(-2) + + +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) + + +@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) -> Any: + 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 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() + + +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_meshgrid_nd( + start: int | tuple[int, ...], + *args: int | tuple[int, ...], + dim: int = 2, + device: torch.device | str | None = None, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Get n-D meshgrid with start, stop and num. + + Args: + start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop, + step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num + should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in + n-tuples. + *args: See above. + dim (int): Dimension of the meshgrid. Defaults to 2. + + Returns: + grid (np.ndarray): [dim, ...] + """ + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0,) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = tuple(stop[i] - start[i] for i in range(dim)) + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=dtype, device=device)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [W, H, D] + grid = torch.stack(grid, dim=0) # [dim, W, H, D] + + return grid + + +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 + + +def get_nd_rotary_pos_embed( + rope_dim_list, + start, + *args, + theta=10000.0, + theta_rescale_factor: float | list[float] = 1.0, + interpolation_factor: float | list[float] = 1.0, + shard_dim: int = 0, + sp_rank: int = 0, + sp_world_size: int = 1, + dtype: torch.dtype = torch.float32, + start_frame: int = 0, + device: torch.device | str | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure. + Supports sequence parallelism by allowing sharding of a specific dimension. + + Args: + rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n. + sum(rope_dim_list) should equal to head_dim of attention layer. + start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start, + args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. + *args: See above. + theta (float): Scaling factor for frequency computation. Defaults to 10000.0. + theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0. + interpolation_factor (float): Factor to scale positions. Defaults to 1.0. + shard_dim (int): Which dimension to shard for sequence parallelism. Defaults to 0. + sp_rank (int): Rank in the sequence parallel group. Defaults to 0. + sp_world_size (int): World size of the sequence parallel group. Defaults to 1. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: (cos, sin) tensors of shape [HW, D/2] + """ + # Determine per-axis sizes for the (possibly sharded) grid without materializing it + ndim = len(rope_dim_list) + if len(args) == 0: + # start is grid_size + sizes = _to_tuple(start, dim=ndim) + starts = (0,) * ndim + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + starts = _to_tuple(start, dim=ndim) + stops = _to_tuple(args[0], dim=ndim) + sizes = tuple(stops[i] - starts[i] for i in range(ndim)) + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + starts = _to_tuple(start, dim=ndim) + _ = _to_tuple(args[0], dim=ndim) # stop, unused here + sizes = _to_tuple(args[1], dim=ndim) + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + assert ( + shard_dim < ndim + ), f"shard_dim {shard_dim} must be less than number of dimensions {ndim}" + + # Apply sequence parallel sharding to the sizes and compute shard offset + shard_sizes = list(sizes) + shard_offsets = [0] * 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 + + # Handle theta scaling/interpolation factor per-axis + if isinstance(theta_rescale_factor, int | float): + theta_rescale_factor = [theta_rescale_factor] * ndim + elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1: + theta_rescale_factor = [theta_rescale_factor[0]] * ndim + assert ( + len(theta_rescale_factor) == ndim + ), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + + if isinstance(interpolation_factor, int | float): + interpolation_factor = [interpolation_factor] * ndim + elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1: + interpolation_factor = [interpolation_factor[0]] * ndim + assert ( + len(interpolation_factor) == ndim + ), "len(interpolation_factor) should equal to len(rope_dim_list)" + + # 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(rope_dim_list) // 2 + cos = torch.empty((num_tokens, head_dim_half), device=device, dtype=dtype) + sin = torch.empty((num_tokens, head_dim_half), device=device, dtype=dtype) + # Compute per-axis 1D embeddings once and expand via repeats to [N, d_i/2] + col_offset = 0 + for i in range(ndim): + dim_i = int(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] + + pos_i = torch.arange(size_i, device=device, dtype=dtype) + base_offset + + cos_1d, sin_1d = get_1d_rotary_pos_embed( + dim_i, + pos_i, + theta=theta, + theta_rescale_factor=theta_rescale_factor[i], + interpolation_factor=interpolation_factor[i], + dtype=dtype, + device=device, + ) # [size_i, dim_i/2] + + # Expand to [num_tokens, dim_i/2] matching flatten order (last dims vary fastest) + repeats_per_entry = 1 + for j in range(i + 1, 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, sin + + +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 + + +_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) + key = ( + head_size, + rotary_dim, + max_position, + 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, base, is_neox_style, dtype + ) + else: + raise ValueError(f"Unknown RoPE scaling {rope_scaling}") + _ROPE_DICT[key] = rotary_emb + return rotary_emb diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/triton_ops.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/triton_ops.py new file mode 100644 index 00000000..2a8d96af --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/triton_ops.py @@ -0,0 +1,948 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# TODO: for temporary usage, expecting a refactor +from typing import Optional + +import torch +import triton # type: ignore +import triton.language as tl # type: ignore +from torch import Tensor + + +@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, + 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 and shift for 4D + 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_ptrs = shift_ptr + scale_row_idx * 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) + + one = tl.full([BLOCK_N], 1.0, dtype=scale.dtype) + output = normalized * (one + scale) + shift + + tl.store(out_ptrs, output, mask=mask) + + +@triton.jit +def fuse_scale_shift_kernel_blc_opt( + x_ptr, + shift_ptr, + scale_ptr, + 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 * (1 + scale) + shift + tl.store(y_ptr + x_off, y, mask=mask) + + +def fuse_scale_shift_kernel( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + 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 [B, F, C] without the singleton dim into [B*F, C] + scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous() + shift_reshaped = shift.squeeze(2).reshape(-1, C).contiguous() + + _fused_scale_shift_4d_kernel[grid]( + output_2d, + x_2d, + scale_reshaped, + shift_reshaped, + 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 (scale_blc.abs().max() == 0) and (shift_blc.abs().max() == 0): + 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, + 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 + + +@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 + + +# 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.cuda.get_device_properties(torch.cuda.current_device()), "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.cuda.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) + 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, + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/usp.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/usp.py new file mode 100644 index 00000000..4f3804c9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/usp.py @@ -0,0 +1,255 @@ +# 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 packaging.version import parse +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, +) + +_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 // world_size, 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 // world_size, 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}" + seq_dim = 1 if head_dim == 2 else 2 + + # Bring to canonical [b, h, s, d] + if head_dim == 1 and seq_dim == 2: + x_c = x + else: + x_c = x.permute(0, head_dim, seq_dim, 3).contiguous() + + b, h, s, d = x_c.shape + assert ( + h % world_size == 0 + ), f"h ({h}) must be divisible by world_size ({world_size})" + + # [b, h, s, d] -> [h, b, s, d] + x_c = x_c.permute(1, 0, 2, 3).contiguous() + # all-to-all along h + x_c = _usp_all_to_all_single(x_c) + # -> [b, h // world, s * world, d] + x_c = ( + x_c.reshape(world_size, h // world_size, b, -1, d) + .permute(2, 1, 0, 3, 4) + .reshape(b, h // world_size, -1, d) + ) + + if head_dim == 1 and seq_dim == 2: + return x_c + + # Map back to original ordering, preserving head/seq positions + new_order = [0, None, None, 3] + new_order[head_dim] = 1 + new_order[seq_dim] = 2 + return x_c.permute(tuple(new_order)).contiguous() + + +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 // world_size, s_global, 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}" + seq_dim = 1 if head_dim == 2 else 2 + + # Bring to canonical [b, h, s, d] + if head_dim == 1 and seq_dim == 2: + x_c = x + else: + x_c = x.permute(0, head_dim, seq_dim, 3).contiguous() + + b, h, s, d = x_c.shape + assert ( + s % world_size == 0 + ), f"s ({s}) must be divisible by world_size ({world_size})" + + # [b, h, s, d] -> [s, b, h, d] + x_c = x_c.permute(2, 0, 1, 3).contiguous() + x_c = _usp_all_to_all_single(x_c) + # -> [b, h * world, s // world, d] + x_c = ( + x_c.reshape(world_size, s // world_size, b, -1, d) + .permute(2, 0, 3, 1, 4) + .reshape(b, -1, s // world_size, d) + ) + + if head_dim == 1 and seq_dim == 2: + return x_c + + # Map back to original ordering, preserving head/seq positions + new_order = [0, None, None, 3] + new_order[head_dim] = 1 + new_order[seq_dim] = 2 + return x_c.permute(tuple(new_order)).contiguous() + + +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 = parse(torch.__version__).release >= parse("2.6.0").release + + attn_kwargs = dict( + mesh=ring_pg, + op=attn_callable_adapter, + dropout_p=dropout_p, + is_causal=is_causal, + query=query, + key=key, + value=value, + ) + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/utils.py new file mode 100644 index 00000000..615ebc38 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/utils.py @@ -0,0 +1,24 @@ +# 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 torch + + +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py new file mode 100644 index 00000000..d556ab58 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py @@ -0,0 +1,190 @@ +# 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 sglang.multimodal_gen.runtime.layers.activation import get_act_fn +from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.mlp import MLP + + +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 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 = ReplicatedLinear( + hidden_size, hidden_size * factor, bias=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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/vocab_parallel_embedding.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/vocab_parallel_embedding.py new file mode 100644 index 00000000..fbddaab4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/layers/vocab_parallel_embedding.py @@ -0,0 +1,480 @@ +# 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.nn.functional as F +from torch.nn.parameter import Parameter, UninitializedParameter + +from sglang.multimodal_gen.runtime.distributed import ( + divide, + get_tp_rank, + get_tp_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.multimodal_gen.runtime.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, + method_has_implemented_embedding, +) +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) +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 = "", + ): + super().__init__() + + # Keep the input dimensions. + tp_rank = get_tp_rank() + self.tp_size = get_tp_world_size() + 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) + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/component_loader.py new file mode 100644 index 00000000..3e93e80a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -0,0 +1,766 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +import glob +import json +import os +import time +from abc import ABC, abstractmethod +from collections.abc import Generator, Iterable +from copy import deepcopy +from typing import Any, cast + +import torch +import torch.distributed as dist +import torch.nn as nn +from safetensors.torch import load_file as safetensors_load_file +from torch.distributed import init_device_mesh +from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME + +from sglang.multimodal_gen.configs.models import EncoderConfig, ModelConfig +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.loader.fsdp_load import ( + maybe_load_fsdp_model, + shard_model, +) +from sglang.multimodal_gen.runtime.loader.utils import set_default_torch_dtype +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, + get_hf_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 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_module_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"))) + + +def load_native(library, component_module_path: str, server_args: ServerArgs): + if library == "transformers": + from transformers import AutoModel + + config = get_hf_config( + component_module_path, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + ) + return AutoModel.from_pretrained( + component_module_path, + config=config, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + ) + elif library == "diffusers": + import diffusers + + config = get_diffusers_component_config(model_path=component_module_path) + class_name = config.pop("_class_name", None) + if class_name: + cls = getattr(diffusers, class_name) + return cls.from_pretrained( + component_module_path, revision=server_args.revision, **config + ) + else: + raise ValueError("Cannot determine class name for generic diffusers loader") + else: + raise ValueError(f"Unsupported library: {library}") + + +class ComponentLoader(ABC): + """Base class for loading a specific type of model component.""" + + def __init__(self, device=None) -> None: + self.device = device + + def should_offload(self, server_args, model_config: ModelConfig | None = None): + raise NotImplementedError() + + 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, + module_name: str, + transformers_or_diffusers: str, + ): + """ + Template method that standardizes logging around the core load implementation. + The priority of loading method is: + 1. load customized module + 2. load native diffusers/transformers module + If all of the above methods failed, an error will be thrown + + """ + logger.info("Loading %s from %s", module_name, component_model_path) + try: + component = self.load_customized( + component_model_path, server_args, module_name + ) + source = "customized" + except Exception as _e: + # 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 module %s: %s is loaded, performance may be sub-optimal", + module_name, + component.__class__.__name__, + ) + + if component is None: + logger.warning("Loaded %s returned None", module_name) + else: + logger.info( + f"Loaded %s: %s from: {source}", + module_name, + component.__class__.__name__, + ) + return component + + def load_native( + self, + component_model_path: str, + server_args: ServerArgs, + transformers_or_diffusers: str, + ): + """ + Load the component using the native library (transformers/diffusers). + """ + return load_native(transformers_or_diffusers, component_model_path, server_args) + + def load_customized( + self, component_model_path: str, server_args: ServerArgs, module_name: str + ): + """ + Load the customized version component, implemented and optimized in SGL-diffusion + """ + raise NotImplementedError( + f"load_customized not implemented for {self.__class__.__name__}" + ) + + @abstractmethod + def load_customized( + self, model_path: str, server_args: ServerArgs, module_name: str + ) -> Any: + """Implement the minimal core load logic in subclasses.""" + raise NotImplementedError + + @classmethod + def for_module_type( + cls, module_type: str, transformers_or_diffusers: str + ) -> "ComponentLoader": + """ + Factory method to create a component loader for a specific module type. + + Args: + module_type: Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler") + transformers_or_diffusers: Whether the module is from transformers or diffusers + + Returns: + A component loader for the specified module type + """ + # Map of module types to their loader classes and expected library + module_type = _normalize_module_type(module_type) + module_loaders = { + "scheduler": (SchedulerLoader, "diffusers"), + "transformer": (TransformerLoader, "diffusers"), + "vae": (VAELoader, "diffusers"), + "text_encoder": (TextEncoderLoader, "transformers"), + "tokenizer": (TokenizerLoader, "transformers"), + "image_processor": (ImageProcessorLoader, "transformers"), + "image_encoder": (ImageEncoderLoader, "transformers"), + "processor": (AutoProcessorLoader, "transformers"), + } + + if module_type in module_loaders: + loader_cls, expected_library = module_loaders[module_type] + # Assert that the library matches what's expected for this module type + assert ( + transformers_or_diffusers == expected_library + ), f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" + return loader_cls() + + # For unknown module types, use a generic loader + logger.warning( + "No specific loader found for module type: %s. Using generic loader.", + module_type, + ) + return GenericComponentLoader(transformers_or_diffusers) + + +class TextEncoderLoader(ComponentLoader): + """Loader for text encoders.""" + + @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.""" + + counter_before_loading_weights: float = 0.0 + counter_after_loading_weights: float = 0.0 + + def should_offload(self, server_args, model_config: ModelConfig | None = None): + should_offload = server_args.text_encoder_cpu_offload + fsdp_shard_conditions = getattr(model_config, "_fsdp_shard_conditions", []) + 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}`" + ) + + 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) + + if self.counter_before_loading_weights == 0.0: + self.counter_before_loading_weights = time.perf_counter() + # 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, module_name: str + ): + """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, + # ) + diffusers_pretrained_config = get_config( + component_model_path, trust_remote_code=True + ) + model_config = get_diffusers_component_config(model_path=component_model_path) + _clean_hf_config_inplace(model_config) + logger.info("HF model config: %s", model_config) + + 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(module_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) + with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]): + with local_torch_device, skip_init_modules(): + architectures = getattr(model_config, "architectures", []) + model_cls, _ = ModelRegistry.resolve_model_cls(architectures) + 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) + ) + self.counter_after_loading_weights = time.perf_counter() + logger.info( + "Loading weights took %.2f seconds", + self.counter_after_loading_weights + - self.counter_before_loading_weights, + ) + + # Explicitly move model to target device after loading weights + 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" + ) + else: + mesh = init_device_mesh( + "cuda", + 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, + ) + # 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: + raise ValueError( + "Following model weights were not initialized from " + f"checkpoint: {weights_not_loaded}" + ) + + return model.eval() + + +class ImageEncoderLoader(TextEncoderLoader): + def should_offload(self, server_args, model_config: ModelConfig | None = None): + should_offload = server_args.image_encoder_cpu_offload + fsdp_shard_conditions = getattr(model_config, "_fsdp_shard_conditions", []) + 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, + # ) + with open(os.path.join(component_model_path, "config.json")) as f: + model_config = json.load(f) + _clean_hf_config_inplace(model_config) + logger.info("HF model config: %s", model_config) + + 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 + should_offload = self.should_offload(server_args) + # 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, + ) + + +class ImageProcessorLoader(ComponentLoader): + """Loader for image processor.""" + + def load_customized( + self, component_model_path: str, server_args: ServerArgs, module_name: str + ) -> Any: + return AutoImageProcessor.from_pretrained(component_model_path, use_fast=True) + + +class AutoProcessorLoader(ComponentLoader): + """Loader for auto processor.""" + + def load_customized( + self, component_model_path: str, server_args: ServerArgs, module_name: str + ) -> Any: + return AutoProcessor.from_pretrained(component_model_path) + + +class TokenizerLoader(ComponentLoader): + """Loader for tokenizers.""" + + def load_customized( + self, component_model_path: str, server_args: ServerArgs, module_name: str + ) -> Any: + return AutoTokenizer.from_pretrained( + component_model_path, + padding_size="right", + ) + + +class VAELoader(ComponentLoader): + """Loader for VAE.""" + + def should_offload(self, server_args, cpu_offload_flag, model_config): + return True + + def load_customized( + self, component_model_path: str, server_args: ServerArgs, *args + ): + """Load the VAE based on the model path, and inference args.""" + config = get_diffusers_component_config(model_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." + + server_args.model_paths["vae"] = component_model_path + + logger.info("HF model config: %s", config) + vae_config = server_args.pipeline_config.vae_config + vae_config.update_model_arch(config) + + # NOTE: some post init logics are only available after updated with config + vae_config.post_init() + + target_device = self.target_device(server_args.vae_cpu_offload) + + with set_default_torch_dtype( + PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + ), skip_init_modules(): + vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) + vae = vae_cls(vae_config).to(target_device) + + # Find all safetensors files + safetensors_list = _list_safetensors_files(component_model_path) + # TODO(PY) + assert ( + len(safetensors_list) == 1 + ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}" + loaded = safetensors_load_file(safetensors_list[0]) + vae.load_state_dict( + loaded, strict=False + ) # We might only load encoder or decoder + + return vae.eval() + + +class TransformerLoader(ComponentLoader): + """Loader for transformer.""" + + def load_customized( + self, component_model_path: str, server_args: ServerArgs, *args + ): + """Load the transformer based on the model path, and inference args.""" + config = get_diffusers_component_config(model_path=component_model_path) + hf_config = deepcopy(config) + cls_name = config.pop("_class_name") + if cls_name is None: + raise ValueError( + "Model config does not contain a _class_name attribute. " + "Only diffusers format is supported." + ) + + logger.info("transformer cls_name: %s", cls_name) + if server_args.override_transformer_cls_name is not None: + cls_name = server_args.override_transformer_cls_name + logger.info("Overriding transformer cls_name to %s", cls_name) + + server_args.model_paths["transformer"] = component_model_path + + # Config from Diffusers supersedes sgl_diffusion's model config + dit_config = server_args.pipeline_config.dit_config + dit_config.update_model_arch(config) + + model_cls, _ = ModelRegistry.resolve_model_cls(cls_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}") + + # Check if we should use custom initialization weights + custom_weights_path = getattr( + server_args, "init_weights_from_safetensors", None + ) + use_custom_weights = False + + if use_custom_weights: + logger.info( + "Using custom initialization weights from: %s", custom_weights_path + ) + assert ( + custom_weights_path is not None + ), "Custom initialization weights must be provided" + if os.path.isdir(custom_weights_path): + safetensors_list = _list_safetensors_files(custom_weights_path) + else: + assert custom_weights_path.endswith( + ".safetensors" + ), "Custom initialization weights must be a safetensors file" + safetensors_list = [custom_weights_path] + + logger.info( + "Loading model from %s safetensors files: %s", + len(safetensors_list), + safetensors_list, + ) + + default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] + + # Load the model using FSDP loader + logger.info("Loading %s, default_dtype: %s", cls_name, default_dtype) + assert server_args.hsdp_shard_dim is not None + model = maybe_load_fsdp_model( + model_cls=model_cls, + init_params={"config": dit_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, + # TODO(will): make these configurable + default_dtype=default_dtype, + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=None, + ) + + total_params = sum(p.numel() for p in model.parameters()) + logger.info("Loaded model with %.2fB parameters", total_params / 1e9) + + assert ( + next(model.parameters()).dtype == default_dtype + ), "Model dtype does not match default dtype" + + model = model.eval() + return model + + +class SchedulerLoader(ComponentLoader): + """Loader for scheduler.""" + + 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(model_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) + if server_args.pipeline_config.timesteps_scale is not None: + scheduler.set_timesteps_scale(server_args.pipeline_config.timesteps_scale) + return scheduler + + +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 pipeline components. + This replaces the chain of if-else statements in load_pipeline_module. + """ + + @staticmethod + def load_module( + module_name: str, + component_model_path: str, + transformers_or_diffusers: str, + server_args: ServerArgs, + ): + """ + Load a pipeline module. + + Args: + module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler") + component_model_path: Path to the component model + transformers_or_diffusers: Whether the module is from transformers or diffusers + + Returns: + The loaded module + """ + logger.info( + "Loading %s using %s from %s", + module_name, + transformers_or_diffusers, + component_model_path, + ) + + # Get the appropriate loader for this module type + loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers) + + try: + # Load the module + return loader.load( + component_model_path, + server_args, + module_name, + transformers_or_diffusers, + ) + except Exception as e: + logger.error( + f"Error while loading component: {module_name}, {component_model_path=}" + ) + raise e diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py new file mode 100644 index 00000000..38c73c90 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -0,0 +1,314 @@ +# 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. + +import contextlib +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.loader.utils import ( + get_param_names_mapping, + hf_to_custom_state_dict, +) +from sglang.multimodal_gen.runtime.loader.weight_utils import ( + safetensors_weights_iterator, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.utils import set_mixed_precision_policy + +logger = init_logger(__name__) + + +# TODO(PY): move this to utils elsewhere +@contextlib.contextmanager +def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]: + """ + Context manager to set torch's default dtype. + + Args: + dtype (torch.dtype): The desired default dtype inside the context manager. + + Returns: + ContextManager: context manager for setting default dtype. + + Example: + >>> with set_default_dtype(torch.bfloat16): + >>> x = torch.tensor([1, 2, 3]) + >>> x.dtype + torch.bfloat16 + + + """ + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + yield + finally: + torch.set_default_dtype(old_dtype) + + +# 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, + default_dtype: torch.dtype, + 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, +) -> torch.nn.Module: + """ + Load the model with FSDP if is training, else load the model without FSDP. + """ + # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are + # manually casting the inputs to the model + mp_policy = MixedPrecisionPolicy( + param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False + ) + + set_mixed_precision_policy( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + output_dtype=output_dtype, + mp_policy=mp_policy, + ) + + with set_default_dtype(default_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 + from sglang.multimodal_gen.runtime.platforms import current_platform + + 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( + "cuda", + # (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, + default_dtype, + strict=True, + cpu_offload=cpu_offload, + param_names_mapping=param_names_mapping_fn, + ) + 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. + + Raises: + ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered. + """ + 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]): + 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, + 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 + 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 + + Raises: + NotImplementedError: If got FSDP with more than 1D. + """ + meta_sd = model.state_dict() + sharded_sd = {} + custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( + full_sd_iterator, param_names_mapping + ) # type: ignore + for target_param_name, full_tensor in custom_param_sd.items(): + meta_sharded_param = meta_sd.get(target_param_name) + if meta_sharded_param is None: + raise ValueError( + f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect." + ) + if not hasattr(meta_sharded_param, "device_mesh"): + full_tensor = full_tensor.to(device=device, dtype=param_dtype) + # In cases where parts of the model aren't sharded, some parameters will be plain tensors + sharded_tensor = full_tensor + else: + full_tensor = full_tensor.to(device=device, dtype=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[target_param_name] = nn.Parameter(sharded_tensor) + + model.reverse_param_names_mapping = reverse_param_names_mapping + 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) + + # List of allowed parameter name patterns + ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress"] # Can be extended as needed + 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) + if not hasattr(meta_sharded_param, "device_mesh"): + # Initialize with zeros + sharded_tensor = torch.zeros_like( + meta_sharded_param, device=device, dtype=param_dtype + ) + else: + # Initialize with zeros and distribute + full_tensor = torch.zeros_like( + meta_sharded_param, device=device, dtype=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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/utils.py new file mode 100644 index 00000000..9f375b9a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/utils.py @@ -0,0 +1,102 @@ +# 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 re +from collections import defaultdict +from collections.abc import Callable, Iterator +from typing import Any + +import torch + +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) + yield + torch.set_default_dtype(old_dtype) + + +def get_param_names_mapping( + mapping_dict: dict[str, str] +) -> 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]: + # Try to match and transform the name using the regex patterns in mapping_dict + for pattern, replacement in mapping_dict.items(): + match = re.match(pattern, name) + if match: + merge_index = None + total_split_params = None + if isinstance(replacement, tuple): + merge_index = replacement[1] + total_split_params = replacement[2] + replacement = replacement[0] + name = re.sub(pattern, replacement, name) + return name, merge_index, total_split_params + + # If no pattern matches, return the original name + return name, None, None + + 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 + ) + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/weight_utils.py new file mode 100644 index 00000000..2bda6ee6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/loader/weight_utils.py @@ -0,0 +1,300 @@ +# 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 and initializing model weights.""" +import hashlib +import json +import os +import tempfile +from collections.abc import Generator +from pathlib import Path + +import filelock +import huggingface_hub.constants +import torch +from safetensors.torch import safe_open +from tqdm.auto import tqdm + +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() + + +def enable_hf_transfer() -> None: + """automatically activates hf_transfer""" + if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ: + try: + # enable hf hub transfer if available + import hf_transfer # type: ignore # noqa + + huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER = True + except ImportError: + pass + + +enable_hf_transfer() + + +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, +) -> 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." + ) + + 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 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 = torch.load(bin_file, map_location=device, weights_only=True) + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/managers/forward_context.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/managers/forward_context.py new file mode 100644 index 00000000..e506929c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/managers/forward_context.py @@ -0,0 +1,120 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py new file mode 100644 index 00000000..7fa2014f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -0,0 +1,191 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +import multiprocessing as mp +import os +import time +from typing import List + +import torch +from setproctitle import setproctitle + +from sglang.multimodal_gen.runtime.distributed import ( + get_sp_group, + maybe_init_distributed_environment_and_model_parallel, +) +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_cfg_group, + get_tp_group, +) +from sglang.multimodal_gen.runtime.pipelines_core import Req, build_pipeline +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch +from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs +from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch +from sglang.multimodal_gen.runtime.utils.logging_utils import ( + configure_logger, + init_logger, + suppress_other_loggers, +) +from sglang.multimodal_gen.runtime.utils.perf_logger import ( + PerformanceLogger, + RequestTimings, +) + +logger = init_logger(__name__) + +CYAN = "\033[1;36m" +RESET = "\033[0;0m" + + +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 = 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.""" + setproctitle(f"sgl_diffusion::scheduler_TP{self.local_rank}") + torch.cuda.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, + ) + + self.pipeline = build_pipeline(self.server_args) + + logger.info( + f"Worker {self.rank}: Initialized device, model, and distributed environment." + ) + + def execute_forward(self, batch: List[Req]) -> OutputBatch: + """ + Execute a forward pass. + """ + assert self.pipeline is not None + # TODO: dealing with first req for now + req = batch[0] + output_batch = None + try: + start_time = time.monotonic() + timings = RequestTimings(request_id=req.request_id) + req.timings = timings + + output_batch = self.pipeline.forward(req, self.server_args) + duration_ms = (time.monotonic() - start_time) * 1000 + + if output_batch.timings: + output_batch.timings.total_duration_ms = duration_ms + PerformanceLogger.log_request_summary(timings=output_batch.timings) + except Exception as e: + if output_batch is None: + from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import ( + OutputBatch, + ) + + output_batch = OutputBatch() + output_batch.error = f"Error executing request {req.request_id}: {e}" + finally: + return output_batch + + def set_lora(self, lora_nickname: str, lora_path: str | None = None) -> None: + """ + Set the LoRA adapter for the pipeline. + """ + assert self.pipeline is not None + self.pipeline.set_lora(lora_nickname, lora_path) + + def merge_lora_weights(self) -> None: + """ + Merge LoRA weights. + """ + assert self.pipeline is not None + self.pipeline.merge_lora_weights() + + def unmerge_lora_weights(self) -> None: + """ + Unmerge LoRA weights. + """ + assert self.pipeline is not None + self.pipeline.unmerge_lora_weights() + + +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) + suppress_other_loggers() + set_cuda_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 + + 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() + logger.info(f"Worker {rank}: Shutdown complete.") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/managers/scheduler.py new file mode 100644 index 00000000..35da11dc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -0,0 +1,230 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +from typing import Any, List + +import zmq + +from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( + MergeLoraWeightsReq, + SetLoraReq, + 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 init_logger + +logger = init_logger(__name__) + + +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: + self.receiver, actual_endpoint = get_zmq_socket( + self.context, zmq.REP, 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, + } + + def _handle_set_lora(self, reqs: List[Any]): + # TODO: return set status + req = reqs[0] + self.worker.set_lora(req.lora_nickname, req.lora_path) + return {"status": "ok"} + + def _handle_merge_lora(self, _reqs: List[Any]): + self.worker.merge_lora_weights() + return {"status": "ok"} + + def _handle_unmerge_lora(self, _reqs: List[Any]): + self.worker.unmerge_lora_weights() + return {"status": "ok"} + + def _handle_generation(self, reqs: List[Req]): + return self.worker.execute_forward(reqs) + + def return_result(self, output_batch: OutputBatch): + """ + replies to client, only on rank 0 + """ + if self.receiver is not None: + self.receiver.send_pyobj(output_batch) + + def recv_reqs(self) -> List[Any]: + """ + For non-main schedulers, reqs are broadcasted from main using broadcast_pyobj + """ + if self.receiver is not None: + try: + recv_reqs = self.receiver.recv_pyobj() + except zmq.ZMQError: + # re-raise or handle appropriately to let the outer loop continue + raise + + # Ensure recv_reqs is a list + if not isinstance(recv_reqs, list): + recv_reqs = [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 + + # TODO: queueing, cancellation + def event_loop(self) -> None: + """ + The main event loop that listens for ZMQ requests. + Handles abortion + """ + + logger.info( + f"Rank 0 scheduler listening on tcp://*:{self.server_args.scheduler_port}" + ) + + while self._running: + reqs = None + # 1: receive requests + try: + reqs = self.recv_reqs() + except Exception as e: + logger.error( + f"Error receiving requests in scheduler event loop: {e}", + exc_info=True, + ) + continue + + # 2: execute, make sure a reply is always sent + try: + first_req = reqs[0] if reqs else None + + handler = self.request_handlers.get(type(first_req)) + if handler: + output_batch = handler(reqs) + else: + output_batch = { + "status": "error", + "message": f"Unknown request type: {type(first_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 {"status": "error", "message": str(e)} + ) + + try: + self.return_result(output_batch) + 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 + + logger.info("Scheduler event loop terminated.") + if self.receiver is not None: + self.receiver.close() + self.context.term() + + 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 _execute_on_rank0(self, payload: dict[str, Any]) -> dict[str, Any]: + """Execute task locally on the rank 0 worker.""" + method = payload["method"] + kwargs = {k: v for k, v in payload.items() if k != "method"} + handler = getattr(self.worker, method, None) + if handler: + result = handler(**kwargs) + return {"status": "ok", "result": result} + return {"status": "error", "error": f"Unknown method: {method}"} + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/__init__.py new file mode 100644 index 00000000..af2eb7d1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/__init__.py @@ -0,0 +1 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/base.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/base.py new file mode 100644 index 00000000..886a6a33 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/base.py @@ -0,0 +1,134 @@ +# 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 +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(BaseDiT): + """ + An intermediate base class that adds TeaCache optimization functionality to DiT models. + TeaCache accelerates inference by selectively skipping redundant computation when consecutive + diffusion steps are similar enough. + """ + + # 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.cnt = 0 + self.teacache_thresh = 0 + self.coefficients: list[float] = [] + + # NOTE(will): Only wan2.1 needs these, so we are hardcoding it here + if self.config.prefix == "wan": + self.use_ret_steps = self.config.cache_config.use_ret_steps + self.is_even = False + self.previous_residual_even: torch.Tensor | None = None + self.previous_residual_odd: torch.Tensor | None = None + self.accumulated_rel_l1_distance_even = 0 + self.accumulated_rel_l1_distance_odd = 0 + self.should_calc_even = True + self.should_calc_odd = True + else: + self.accumulated_rel_l1_distance = 0 + self.previous_modulated_input = None + self.previous_resiual = None + self.previous_e0_even: torch.Tensor | None = None + self.previous_e0_odd: torch.Tensor | None = None + + def maybe_cache_states( + self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor + ) -> None: + pass + + def should_skip_forward_for_cached_states(self, **kwargs: dict[str, Any]) -> bool: + return False + + def retrieve_cached_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + raise NotImplementedError("maybe_retrieve_cached_states is not implemented") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py new file mode 100644 index 00000000..2789ebdf --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py @@ -0,0 +1,851 @@ +# 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, +) + +# 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.layernorm import ( + FP32LayerNorm, + LayerNormScaleShift, + RMSNorm, + ScaleResidual, + 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.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.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 = "", + ): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True) + self.to_k = ReplicatedLinear(dim, dim, bias=True) + self.to_v = ReplicatedLinear(dim, dim, bias=True) + + self.to_out = ReplicatedLinear(dim, dim, bias=True) + 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, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32, + compute_dtype=torch.float32, + ) + + # 2. Cross-attention + # Only T2V for now + self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32, + ) + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.mlp_residual = ScaleResidual() + + 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(hidden_states, ff_output, c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) + + return hidden_states + + +class CausalWanTransformer3DModel(BaseDiT): + _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]) -> 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}", + ) + for i in range(config.num_layers) + ] + ) + + # 4. Output norm & projection + self.norm_out = LayerNormScaleShift( + inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32, + compute_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__() + + @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() 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() 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/flux.py new file mode 100644 index 00000000..fbb75242 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -0,0 +1,558 @@ +# 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, FeedForward +from diffusers.models.embeddings import ( + CombinedTimestepGuidanceTextProjEmbeddings, + CombinedTimestepTextProjEmbeddings, +) +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 LayerNorm as LayerNorm +from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm +from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.mlp import MLP +from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( + NDRotaryEmbedding, + _apply_rotary_emb, +) +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.logging_utils import init_logger + +logger = init_logger(__name__) # pylint: disable=invalid-name + + +def _get_projections(attn: "FluxAttention", 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 + + +def _get_fused_projections( + attn: "FluxAttention", hidden_states, encoder_hidden_states=None +): + query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) + + encoder_query = encoder_key = encoder_value = None + if encoder_hidden_states is not None and hasattr(attn, "to_added_qkv"): + encoder_query, encoder_key, encoder_value = attn.to_added_qkv( + encoder_hidden_states + ).chunk(3, dim=-1) + + return query, key, value, encoder_query, encoder_key, encoder_value + + +def _get_qkv_projections( + attn: "FluxAttention", hidden_states, encoder_hidden_states=None +): + if attn.fused_projections: + return _get_fused_projections(attn, hidden_states, encoder_hidden_states) + return _get_projections(attn, hidden_states, encoder_hidden_states) + + +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, + ): + 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.norm_q = RMSNorm(dim_head, eps=eps) + + self.norm_k = RMSNorm(dim_head, eps=eps) + self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias) + self.to_k = ReplicatedLinear(query_dim, self.inner_dim, bias=bias) + self.to_v = ReplicatedLinear(query_dim, self.inner_dim, bias=bias) + + if not self.pre_only: + self.to_out = torch.nn.ModuleList([]) + self.to_out.append( + ReplicatedLinear(self.inner_dim, self.out_dim, bias=out_bias) + ) + 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) + self.add_q_proj = ReplicatedLinear( + added_kv_proj_dim, self.inner_dim, bias=added_proj_bias + ) + self.add_k_proj = ReplicatedLinear( + added_kv_proj_dim, self.inner_dim, bias=added_proj_bias + ) + self.add_v_proj = ReplicatedLinear( + added_kv_proj_dim, self.inner_dim, bias=added_proj_bias + ) + self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias) + + 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.TORCH_SDPA, + AttentionBackendEnum.SAGE_ATTN, + }, + ) + + 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 = self.norm_q(query) + key = self.norm_k(key) + + 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 = self.norm_added_q(encoder_query) + encoder_key = self.norm_added_k(encoder_key) + + 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 + query = _apply_rotary_emb( + query, cos, sin, is_neox_style=False, interleaved=False + ) + key = _apply_rotary_emb( + key, cos, sin, is_neox_style=False, interleaved=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, + ) + 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: + return x + + +class FluxSingleTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 4.0, + ): + super().__init__() + self.mlp_hidden_dim = int(dim * mlp_ratio) + + self.norm = AdaLayerNormZeroSingle(dim) + self.proj_mlp = ReplicatedLinear(dim, self.mlp_hidden_dim) + self.act_mlp = nn.GELU(approximate="tanh") + self.proj_out = ReplicatedLinear(dim + self.mlp_hidden_dim, dim) + + 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, + ) + + 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) + proj_hidden_states, _ = self.proj_mlp(norm_hidden_states) + mlp_hidden_states = self.act_mlp(proj_hidden_states) + joint_attention_kwargs = joint_attention_kwargs or {} + 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, + ): + 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, + ) + + self.norm2 = LayerNorm(dim, eps=1e-6, elementwise_affine=False) + self.ff = MLP( + input_dim=dim, mlp_hidden_dim=dim * 4, output_dim=dim, act_type="gelu" + ) + self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") + + self.norm2_context = LayerNorm(dim, eps=1e-6, elementwise_affine=False) + self.ff_context = MLP( + input_dim=dim, mlp_hidden_dim=dim * 4, output_dim=dim, act_type="gelu" + ) + + self.ff_context = FeedForward( + dim=dim, dim_out=dim, activation_fn="gelu-approximate" + ) + + 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) + 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) + 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() else torch.float64, + ) + + def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + pos = ids.float() + # freqs_cos, freqs_sin = self.rope.forward(positions=pos) + freqs_cos, freqs_sin = self.rope.forward_uncached(pos=pos) + return freqs_cos.contiguous().float(), freqs_sin.contiguous().float() + + +class FluxTransformer2DModel(CachableDiT): + """ + The Transformer model introduced in Flux. + + Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ + """ + + def __init__(self, config: FluxConfig, hf_config: dict[str, Any]) -> 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 = ReplicatedLinear( + self.config.joint_attention_dim, self.inner_dim + ) + self.x_embedder = ReplicatedLinear(self.config.in_channels, self.inner_dim) + 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, + ) + for _ 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, + ) + for _ 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 = ReplicatedLinear( + self.inner_dim, + self.config.patch_size * self.config.patch_size * self.out_channels, + bias=True, + ) + + 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) + + temb = ( + self.time_text_embed(timestep, pooled_projections) + if guidance is None + else self.time_text_embed(timestep, guidance, 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 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=temb, + freqs_cis=freqs_cis, + joint_attention_kwargs=joint_attention_kwargs, + ) + + for index_block, block in enumerate(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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py new file mode 100644 index 00000000..3ff593f8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -0,0 +1,814 @@ +# 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, Optional, Tuple + +import torch +import torch.nn as nn +from diffusers.models.attention import AttentionModuleMixin +from diffusers.models.embeddings import ( + TimestepEmbedding, + Timesteps, + get_1d_rotary_pos_embed, +) +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 +from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) # pylint: disable=invalid-name + + +def _get_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 + + +def _get_fused_projections( + attn: "Flux2Attention", hidden_states, encoder_hidden_states=None +): + query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) + + encoder_query = encoder_key = encoder_value = (None,) + if encoder_hidden_states is not None and hasattr(attn, "to_added_qkv"): + encoder_query, encoder_key, encoder_value = attn.to_added_qkv( + encoder_hidden_states + ).chunk(3, dim=-1) + + return query, key, value, encoder_query, encoder_key, encoder_value + + +def _get_qkv_projections( + attn: "Flux2Attention", hidden_states, encoder_hidden_states=None +): + if attn.fused_projections: + return _get_fused_projections(attn, hidden_states, encoder_hidden_states) + return _get_projections(attn, hidden_states, encoder_hidden_states) + + +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, + ): + 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 = nn.Linear(dim, inner_dim * 2, bias=bias) + self.act_fn = Flux2SwiGLU() + self.linear_out = nn.Linear(inner_dim, dim_out, bias=bias) + + 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, + ): + 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 = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) + + # 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(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) + 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 = torch.nn.Linear( + added_kv_proj_dim, self.inner_dim, bias=added_proj_bias + ) + self.add_k_proj = torch.nn.Linear( + added_kv_proj_dim, self.inner_dim, bias=added_proj_bias + ) + self.add_v_proj = torch.nn.Linear( + added_kv_proj_dim, self.inner_dim, bias=added_proj_bias + ) + self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias) + + 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.TORCH_SDPA, + AttentionBackendEnum.SAGE_ATTN, + }, + ) + + 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 = self.norm_q(query) + key = self.norm_k(key) + + 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 = self.norm_added_q(encoder_query) + encoder_key = self.norm_added_k(encoder_key) + + 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 + query = _apply_rotary_emb( + query, cos, sin, is_neox_style=False, interleaved=True + ) + key = _apply_rotary_emb( + key, cos, sin, is_neox_style=False, interleaved=True + ) + + hidden_states = self.attn(query, key, value) + + 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, + ): + 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 = torch.nn.Linear( + self.query_dim, + self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor, + bias=bias, + ) + 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 = torch.nn.Linear( + self.inner_dim + self.mlp_hidden_dim, self.out_dim, bias=out_bias + ) + + 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.TORCH_SDPA, + AttentionBackendEnum.SAGE_ATTN, + }, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **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 + query = _apply_rotary_emb( + query, cos, sin, is_neox_style=False, interleaved=True + ) + key = _apply_rotary_emb( + key, cos, sin, is_neox_style=False, interleaved=True + ) + hidden_states = self.attn(query, key, value) + 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, + ): + 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, + ) + + 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, + **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, + ): + 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, + ) + + self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.ff = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias) + + 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 + ) + + 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 + ): + 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 + ) + + self.guidance_embedder = TimestepEmbedding( + in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias + ) + + def forward(self, timestep: torch.Tensor, guidance: torch.Tensor) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder( + timesteps_proj.to(timestep.dtype) + ) # (N, D) + + 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 + + +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 = nn.Linear(dim, dim * 3 * self.mod_param_sets, bias=bias) + 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): + # 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.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: torch.Tensor) -> torch.Tensor: + # Expected ids shape: [S, len(self.axes_dim)] + cos_out = [] + sin_out = [] + pos = ids.float() + is_mps = ids.device.type == "mps" + is_npu = ids.device.type == "npu" + freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64 + # Unlike Flux 1, loop over len(self.axes_dim) rather than ids.shape[-1] + for i in range(len(self.axes_dim)): + cos, sin = get_1d_rotary_pos_embed( + self.axes_dim[i], + pos[..., i], + theta=self.theta, + repeat_interleave_real=True, + use_real=True, + freqs_dtype=freqs_dtype, + ) + cos_out.append(cos) + sin_out.append(sin) + freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) + freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) + return freqs_cos, freqs_sin + + +class Flux2Transformer2DModel(CachableDiT): + """ + The Transformer model introduced in Flux 2. + + Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ + + """ + + def __init__(self, config: FluxConfig, hf_config: dict[str, Any]): + 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 + self.out_channels = out_channels or in_channels + self.inner_dim = num_attention_heads * attention_head_dim + + # 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, + ) + + # 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 = nn.Linear(in_channels, self.inner_dim, bias=False) + self.context_embedder = nn.Linear( + joint_attention_dim, self.inner_dim, bias=False + ) + + # 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, + ) + 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, + ) + 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 = nn.Linear( + self.inner_dim, patch_size * patch_size * self.out_channels, bias=False + ) + + self.gradient_checkpointing = False + + 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) + 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, + ) + # 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py new file mode 100644 index 00000000..ad9e10dd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py @@ -0,0 +1,961 @@ +# 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.layernorm import ( + LayerNormScaleShift, + RMSNorm, + ScaleResidual, + 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.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 + + +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 = "", + ): + 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, norm_type="layer", elementwise_affine=False, dtype=dtype + ) + self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( + hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype + ) + self.img_mlp_residual = ScaleResidual() + + # Image attention components + self.img_attn_qkv = ReplicatedLinear( + hidden_size, + hidden_size * 3, + bias=True, + params_dtype=dtype, + prefix=f"{prefix}.img_attn_qkv", + ) + + 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", + ) + + self.img_mlp = MLP( + hidden_size, + mlp_hidden_dim, + bias=True, + dtype=dtype, + prefix=f"{prefix}.img_mlp", + ) + + # 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, norm_type="layer", elementwise_affine=False, dtype=dtype + ) + self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( + hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype + ) + self.txt_mlp_residual = ScaleResidual() + + # Text attention components + self.txt_attn_qkv = ReplicatedLinear( + hidden_size, hidden_size * 3, bias=True, params_dtype=dtype + ) + + # 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 + ) + + self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype) + + # 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_residual, img_mlp_out, img_mlp_gate) + + # 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_residual, txt_mlp_out, txt_mlp_gate) + + 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 = "", + ): + 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", + ) + + # Combined projection and MLP output + self.linear2 = ReplicatedLinear( + hidden_size + mlp_hidden_dim, + hidden_size, + bias=True, + params_dtype=dtype, + prefix=f"{prefix}.linear2", + ) + + # 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, + norm_type="layer", + eps=1e-6, + elementwise_affine=False, + dtype=dtype, + ) + self.output_residual = ScaleResidual() + + # 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(x, output, mod_gate) + + +class HunyuanVideoTransformer3DModel(CachableDiT): + """ + 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]): + 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}", + ) + 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}", + ) + 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__() + + # 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( + "cuda", 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( + "cuda", 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.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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py new file mode 100644 index 00000000..989d6d52 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -0,0 +1,650 @@ +# 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 numpy as np +import torch +import torch.nn as nn +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.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig +from sglang.multimodal_gen.runtime.layers.attention import USPAttention +from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm, RMSNorm +from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.triton_ops import ( + apply_rotary_embedding, + fuse_scale_shift_kernel, +) +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) # pylint: disable=invalid-name + + +class QwenTimestepProjEmbeddings(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, scale=1000 + ) + self.timestep_embedder = TimestepEmbedding( + in_channels=256, time_embed_dim=embedding_dim + ) + + def forward(self, timestep, hidden_states): + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder( + timesteps_proj.to(dtype=hidden_states.dtype) + ) # (N, D) + + conditioning = timesteps_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, + ) + + # 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() else torch.float64, + # ) + + # 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 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, + ) -> 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 + + # layers + self.to_q = ReplicatedLinear(dim, dim) + self.to_k = ReplicatedLinear(dim, dim) + self.to_v = ReplicatedLinear(dim, dim) + 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() + self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads + self.inner_kv_dim = self.inner_dim + if added_kv_proj_dim is not None: + self.add_k_proj = ReplicatedLinear( + added_kv_proj_dim, self.inner_kv_dim, bias=True + ) + self.add_v_proj = ReplicatedLinear( + added_kv_proj_dim, self.inner_kv_dim, bias=True + ) + if context_pre_only is not None: + self.add_q_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.TORCH_SDPA, + }, + ) + + 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] + + # Compute QKV for image stream (sample projections) + img_query, _ = self.to_q(hidden_states) + img_key, _ = self.to_k(hidden_states) + img_value, _ = self.to_v(hidden_states) + + # Compute QKV for text stream (context projections) + txt_query, _ = self.add_q_proj(encoder_hidden_states) + txt_key, _ = self.add_k_proj(encoder_hidden_states) + txt_value, _ = self.add_v_proj(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.norm_q is not None: + img_query = self.norm_q(img_query) + if self.norm_k is not None: + img_key = self.norm_k(img_key) + if self.norm_added_q is not None: + txt_query = self.norm_added_q(txt_query) + if self.norm_added_k is not None: + txt_key = self.norm_added_k(txt_key) + + # Apply RoPE + if image_rotary_emb is not None: + (img_cos, img_sin), (txt_cos, txt_sin) = image_rotary_emb + img_query = apply_rotary_embedding( + img_query, img_cos, img_sin, interleaved=True + ) + img_key = apply_rotary_embedding( + img_key, img_cos, img_sin, interleaved=True + ) + txt_query = apply_rotary_embedding( + txt_query, txt_cos, txt_sin, interleaved=True + ) + txt_key = apply_rotary_embedding( + txt_key, txt_cos, txt_sin, interleaved=True + ) + + # 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, + ): + super().__init__() + + self.dim = dim + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + + # 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 = LayerNorm(dim, elementwise_affine=False, eps=eps) + + self.attn = QwenImageCrossAttention( + dim=dim, + num_heads=num_attention_heads, + added_kv_proj_dim=dim, + context_pre_only=False, + head_dim=attention_head_dim, + ) + self.img_norm2 = LayerNorm(dim, eps=eps, elementwise_affine=False) + self.img_mlp = FeedForward( + dim=dim, dim_out=dim, activation_fn="gelu-approximate" + ) + + # 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 = LayerNorm(dim, elementwise_affine=False, eps=eps) + # Text doesn't need separate attention - it's handled by img_attn joint computation + self.txt_norm2 = LayerNorm(dim, elementwise_affine=False, eps=eps) + self.txt_mlp = FeedForward( + dim=dim, dim_out=dim, activation_fn="gelu-approximate" + ) + + def _modulate(self, x, mod_params): + """Apply modulation to input tensor""" + shift, scale, gate = mod_params.chunk(3, dim=-1) + return fuse_scale_shift_kernel(x, scale, shift), gate.unsqueeze(1) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Get modulation parameters for both streams + img_mod_params = self.img_mod(temb) # [B, 6*dim] + txt_mod_params = self.txt_mod(temb) # [B, 6*dim] + + # 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_normed = self.img_norm1(hidden_states) + + img_modulated, img_gate1 = self._modulate(img_normed, img_mod1) + + # Process text stream - norm1 + modulation + txt_normed = self.txt_norm1(encoder_hidden_states) + txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1) + + # 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( + hidden_states=img_modulated, # Image stream (will be processed as "sample") + encoder_hidden_states=txt_modulated, # Text stream (will be processed as "context") + 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 + + # Apply attention gates and add residual (like in Megatron) + hidden_states = hidden_states + img_gate1 * img_attn_output + + encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output + + # Process image stream - norm2 + MLP + img_normed2 = self.img_norm2(hidden_states) + img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2) + img_mlp_output = self.img_mlp(img_modulated2) + hidden_states = hidden_states + img_gate2 * img_mlp_output + + # Process text stream - norm2 + MLP + txt_normed2 = self.txt_norm2(encoder_hidden_states) + txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2) + txt_mlp_output = self.txt_mlp(txt_modulated2) + encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output + + # 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 + + +class QwenImageTransformer2DModel(CachableDiT): + """ + The Transformer model introduced in Qwen. + + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["QwenImageTransformerBlock"] + _skip_layerwise_casting_patterns = ["pos_embed", "norm"] + _repeated_blocks = ["QwenImageTransformerBlock"] + + def __init__( + self, + config: QwenImageDitConfig, + hf_config: dict[str, Any], + ): + 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.out_channels = out_channels or in_channels + self.inner_dim = num_attention_heads * attention_head_dim + + self.rotary_emb = QwenEmbedRope( + theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True + ) + + self.time_text_embed = QwenTimestepProjEmbeddings(embedding_dim=self.inner_dim) + + 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, + ) + for _ 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 + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + encoder_hidden_states_mask: torch.Tensor = None, + timestep: torch.LongTensor = None, + txt_seq_lens: Optional[List[int]] = None, + freqs_cis: tuple[torch.Tensor, 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) + 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) + + 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=temb, + image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=attention_kwargs, + ) + + # 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) + + output = self.proj_out(hidden_states) + return output + + +EntryClass = QwenImageTransformer2DModel diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/stepvideo.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/stepvideo.py new file mode 100644 index 00000000..529c4995 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/stepvideo.py @@ -0,0 +1,729 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# Copyright 2025 StepFun 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. +# ============================================================================== +from typing import Any + +import torch +from einops import rearrange, repeat +from torch import nn + +from sglang.multimodal_gen.configs.models.dits import StepVideoConfig +from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size +from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention +from sglang.multimodal_gen.runtime.layers.layernorm import LayerNormScaleShift +from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.mlp import MLP +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 TimestepEmbedder +from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum + + +class PatchEmbed2D(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.Conv2d( + 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 StepVideoRMSNorm(nn.Module): + + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x) -> torch.Tensor: + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +class SelfAttention(nn.Module): + + def __init__( + self, + hidden_dim, + head_dim, + rope_split: tuple[int, int, int] = (64, 32, 32), + bias: bool = False, + with_rope: bool = True, + with_qk_norm: bool = True, + attn_type: str = "torch", + supported_attention_backends=( + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, + ), + ): + super().__init__() + self.head_dim = head_dim + self.hidden_dim = hidden_dim + self.rope_split = list(rope_split) + self.n_heads = hidden_dim // head_dim + + self.wqkv = ReplicatedLinear(hidden_dim, hidden_dim * 3, bias=bias) + self.wo = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias) + + self.with_rope = with_rope + self.with_qk_norm = with_qk_norm + if self.with_qk_norm: + self.q_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True) + self.k_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True) + + # self.core_attention = self.attn_processor(attn_type=attn_type) + self.parallel = attn_type == "parallel" + self.attn = USPAttention( + num_heads=self.n_heads, + head_size=head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + ) + + def _apply_rope(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + """ + x: [B, S, H, D] + cos: [S, D/2] where D = head_dim = sum(self.rope_split) + sin: [S, D/2] + returns x with rotary applied exactly as v0 did + """ + B, S, H, D = x.shape + # 1) split cos/sin per chunk + half_splits = [c // 2 for c in self.rope_split] # [32,16,16] for [64,32,32] + cos_splits = cos.split(half_splits, dim=1) + sin_splits = sin.split(half_splits, dim=1) + + outs = [] + idx = 0 + for chunk_size, cos_i, sin_i in zip( + self.rope_split, cos_splits, sin_splits, strict=True + ): + # slice the corresponding channels + x_chunk = x[..., idx : idx + chunk_size] # [B,S,H,chunk_size] + idx += chunk_size + + # flatten to [S, B*H, chunk_size] + x_flat = rearrange(x_chunk, "b s h d -> s (b h) d") + + # apply rotary on *that* chunk + out_flat = _apply_rotary_emb(x_flat, cos_i, sin_i, is_neox_style=True) + + # restore [B,S,H,chunk_size] + out = rearrange(out_flat, "s (b h) d -> b s h d", b=B, h=H) + outs.append(out) + + # concatenate back to [B,S,H,D] + return torch.cat(outs, dim=-1) + + def forward( + self, + x, + cu_seqlens=None, + max_seqlen=None, + rope_positions=None, + cos_sin=None, + attn_mask=None, + mask_strategy=None, + ): + + B, S, _ = x.shape + xqkv, _ = self.wqkv(x) + xqkv = xqkv.view(*x.shape[:-1], self.n_heads, 3 * self.head_dim) + q, k, v = torch.split(xqkv, [self.head_dim] * 3, dim=-1) # [B,S,H,D] + + if self.with_qk_norm: + q = self.q_norm(q) + k = self.k_norm(k) + + if self.with_rope: + if rope_positions is not None: + F, Ht, W = rope_positions + assert F * Ht * W == S, "rope_positions mismatches sequence length" + + cos, sin = cos_sin + cos = cos.to(x.device, dtype=x.dtype) + sin = sin.to(x.device, dtype=x.dtype) + + q = self._apply_rope(q, cos, sin) + k = self._apply_rope(k, cos, sin) + + output = self.attn(q, k, v) # [B,heads,S,D] + + output = rearrange(output, "b s h d -> b s (h d)") + output, _ = self.wo(output) + + return output + + +class CrossAttention(nn.Module): + + def __init__( + self, + hidden_dim, + head_dim, + bias=False, + with_qk_norm=True, + supported_attention_backends=( + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, + ), + ) -> None: + super().__init__() + self.head_dim = head_dim + self.n_heads = hidden_dim // head_dim + + self.wq = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias) + self.wkv = ReplicatedLinear(hidden_dim, hidden_dim * 2, bias=bias) + self.wo = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias) + + self.with_qk_norm = with_qk_norm + if self.with_qk_norm: + self.q_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True) + self.k_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True) + + self.attn = LocalAttention( + num_heads=self.n_heads, + head_size=head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + ) + + def forward( + self, x: torch.Tensor, encoder_hidden_states: torch.Tensor, attn_mask=None + ) -> torch.Tensor: + + xq, _ = self.wq(x) + xq = xq.view(*xq.shape[:-1], self.n_heads, self.head_dim) + + xkv, _ = self.wkv(encoder_hidden_states) + xkv = xkv.view(*xkv.shape[:-1], self.n_heads, 2 * self.head_dim) + + xk, xv = torch.split(xkv, [self.head_dim] * 2, dim=-1) ## seq_len, n, dim + + if self.with_qk_norm: + xq = self.q_norm(xq) + xk = self.k_norm(xk) + + output = self.attn(xq, xk, xv) + + output = rearrange(output, "b s h d -> b s (h d)") + output, _ = self.wo(output) + + return output + + +class AdaLayerNormSingle(nn.Module): + r""" + Norm layer adaptive layer norm single (adaLN-single). + + As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3). + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + use_additional_conditions (`bool`): To use additional conditions for normalization or not. + """ + + def __init__(self, embedding_dim: int, time_step_rescale=1000): + super().__init__() + + self.emb = TimestepEmbedder(embedding_dim) + + self.silu = nn.SiLU() + self.linear = ReplicatedLinear(embedding_dim, 6 * embedding_dim, bias=True) + + self.time_step_rescale = time_step_rescale ## timestep usually in [0, 1], we rescale it to [0,1000] for stability + + def forward( + self, + timestep: torch.Tensor, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + embedded_timestep = self.emb(timestep * self.time_step_rescale) + + out, _ = self.linear(self.silu(embedded_timestep)) + + return out, embedded_timestep + + +class StepVideoTransformerBlock(nn.Module): + r""" + A basic Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + only_cross_attention (`bool`, *optional*): + Whether to use only cross-attention layers. In this case two cross attention layers are used. + double_self_attention (`bool`, *optional*): + Whether to use two self-attention layers. In this case no cross attention layers are used. + upcast_attention (`bool`, *optional*): + Whether to upcast the attention computation to float32. This is useful for mixed precision training. + norm_elementwise_affine (`bool`, *optional*, defaults to `True`): + Whether to use learnable elementwise affine parameters for normalization. + norm_type (`str`, *optional*, defaults to `"layer_norm"`): + The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. + final_dropout (`bool` *optional*, defaults to False): + Whether to apply a final dropout after the last feed-forward layer. + positional_embeddings (`str`, *optional*, defaults to `None`): + The type of positional embeddings to apply to. + num_positional_embeddings (`int`, *optional*, defaults to `None`): + The maximum number of positional embeddings to apply. + """ + + def __init__( + self, + dim: int, + attention_head_dim: int, + norm_eps: float = 1e-5, + ff_inner_dim: int | None = None, + ff_bias: bool = False, + attention_type: str = "torch", + ): + super().__init__() + self.dim = dim + self.norm1 = LayerNormScaleShift( + dim, norm_type="layer", elementwise_affine=True, eps=norm_eps + ) + self.attn1 = SelfAttention( + dim, + attention_head_dim, + bias=False, + with_rope=True, + with_qk_norm=True, + ) + + self.norm2 = LayerNormScaleShift( + dim, norm_type="layer", elementwise_affine=True, eps=norm_eps + ) + self.attn2 = CrossAttention( + dim, attention_head_dim, bias=False, with_qk_norm=True + ) + + self.ff = MLP( + input_dim=dim, + mlp_hidden_dim=dim * 4 if ff_inner_dim is None else ff_inner_dim, + act_type="gelu_pytorch_tanh", + bias=ff_bias, + ) + + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + + @torch.no_grad() + def forward( + self, + q: torch.Tensor, + kv: torch.Tensor, + t_expand: torch.LongTensor, + attn_mask=None, + rope_positions: list | None = None, + cos_sin=None, + mask_strategy=None, + ) -> torch.Tensor: + + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + torch.clone(chunk) + for chunk in ( + self.scale_shift_table[None] + t_expand.reshape(-1, 6, self.dim) + ).chunk(6, dim=1) + ) + + scale_shift_q = self.norm1( + q, scale=scale_msa.squeeze(1), shift=shift_msa.squeeze(1) + ) + + attn_q = self.attn1( + scale_shift_q, + rope_positions=rope_positions, + cos_sin=cos_sin, + mask_strategy=mask_strategy, + ) + + q = attn_q * gate_msa + q + + attn_q = self.attn2(q, kv, attn_mask) + + q = attn_q + q + + scale_shift_q = self.norm2( + q, scale=scale_mlp.squeeze(1), shift=shift_mlp.squeeze(1) + ) + + ff_output = self.ff(scale_shift_q) + + q = ff_output * gate_mlp + q + + return q + + +class StepVideoModel(BaseDiT): + # (Optional) Keep the same attribute for compatibility with splitting, etc. + _fsdp_shard_conditions = [ + lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit(), + # lambda n, m: "pos_embed" in n # If needed for the patch embedding. + ] + param_names_mapping = StepVideoConfig().param_names_mapping + reverse_param_names_mapping = StepVideoConfig().reverse_param_names_mapping + lora_param_names_mapping = StepVideoConfig().lora_param_names_mapping + _supported_attention_backends = StepVideoConfig()._supported_attention_backends + + def __init__(self, config: StepVideoConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + 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_layers = config.num_layers + self.dropout = config.dropout + self.patch_size = config.patch_size + self.norm_type = config.norm_type + self.norm_elementwise_affine = config.norm_elementwise_affine + self.norm_eps = config.norm_eps + self.use_additional_conditions = config.use_additional_conditions + self.caption_channels = config.caption_channels + self.attention_type = config.attention_type + self.num_channels_latents = config.num_channels_latents + # Compute inner dimension. + self.hidden_size = config.hidden_size + + # Image/video patch embedding. + self.pos_embed = PatchEmbed2D( + patch_size=self.patch_size, + in_chans=self.in_channels, + embed_dim=self.hidden_size, + ) + + self._rope_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {} + # Transformer blocks. + self.transformer_blocks = nn.ModuleList( + [ + StepVideoTransformerBlock( + dim=self.hidden_size, + attention_head_dim=self.attention_head_dim, + attention_type=self.attention_type, + ) + for _ in range(self.num_layers) + ] + ) + + # Output blocks. + self.norm_out = LayerNormScaleShift( + self.hidden_size, + norm_type="layer", + eps=self.norm_eps, + elementwise_affine=self.norm_elementwise_affine, + ) + self.scale_shift_table = nn.Parameter( + torch.randn(2, self.hidden_size) / (self.hidden_size**0.5) + ) + self.proj_out = ReplicatedLinear( + self.hidden_size, self.patch_size * self.patch_size * self.out_channels + ) + # Time modulation via adaptive layer norm. + self.adaln_single = AdaLayerNormSingle(self.hidden_size) + + # Set up caption conditioning. + if isinstance(self.caption_channels, int): + caption_channel = self.caption_channels + else: + caption_channel, clip_channel = self.caption_channels + self.clip_projection = ReplicatedLinear(clip_channel, self.hidden_size) + self.caption_norm = nn.LayerNorm( + caption_channel, + eps=self.norm_eps, + elementwise_affine=self.norm_elementwise_affine, + ) + self.caption_projection = MLP( + input_dim=caption_channel, + mlp_hidden_dim=self.hidden_size, + act_type="gelu_pytorch_tanh", + ) + + # Flag to indicate if using parallel attention. + self.parallel = self.attention_type == "parallel" + + self.__post_init__() + + def patchfy(self, hidden_states) -> torch.Tensor: + hidden_states = rearrange(hidden_states, "b f c h w -> (b f) c h w") + hidden_states = self.pos_embed(hidden_states) + return hidden_states + + def prepare_attn_mask( + self, encoder_attention_mask, encoder_hidden_states, q_seqlen + ) -> tuple[torch.Tensor, torch.Tensor]: + kv_seqlens = encoder_attention_mask.sum(dim=1).int() + mask = torch.zeros( + [len(kv_seqlens), q_seqlen, max(kv_seqlens)], + dtype=torch.bool, + device=encoder_attention_mask.device, + ) + encoder_hidden_states = encoder_hidden_states[:, : max(kv_seqlens)] + for i, kv_len in enumerate(kv_seqlens): + mask[i, :, :kv_len] = 1 + return encoder_hidden_states, mask + + def block_forward( + self, + hidden_states, + encoder_hidden_states=None, + t_expand=None, + rope_positions=None, + cos_sin=None, + attn_mask=None, + parallel=True, + mask_strategy=None, + ) -> torch.Tensor: + + for i, block in enumerate(self.transformer_blocks): + hidden_states = block( + hidden_states, + encoder_hidden_states, + t_expand=t_expand, + attn_mask=attn_mask, + rope_positions=rope_positions, + cos_sin=cos_sin, + mask_strategy=mask_strategy[i], + ) + + return hidden_states + + def _get_rope( + self, + rope_positions: tuple[int, int, int], + dtype: torch.dtype, + device: torch.device, + ): + F, Ht, W = rope_positions + key = (F, Ht, W, dtype) + if key not in self._rope_cache: + cos, sin = get_rotary_pos_embed( + rope_sizes=(F * get_sp_world_size(), Ht, W), + hidden_size=self.hidden_size, + heads_num=self.hidden_size // self.attention_head_dim, + rope_dim_list=(64, 32, 32), # same split you used + rope_theta=1.0e4, + dtype=torch.float32, # build once in fp32 + ) + # move & cast once + self._rope_cache[key] = ( + cos.to(device, dtype=dtype), + sin.to(device, dtype=dtype), + ) + return self._rope_cache[key] + + @torch.inference_mode() + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + t_expand: torch.LongTensor | None = None, + encoder_hidden_states_2: torch.Tensor | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + encoder_attention_mask: torch.Tensor | None = None, + fps: torch.Tensor | None = None, + return_dict: bool = True, + mask_strategy=None, + guidance=None, + ): + assert hidden_states.ndim == 5 + "hidden_states's shape should be (bsz, f, ch, h ,w)" + frame = hidden_states.shape[2] + hidden_states = rearrange(hidden_states, "b c f h w -> b f c h w", f=frame) + if mask_strategy is None: + mask_strategy = [None, None] + bsz, frame, _, height, width = hidden_states.shape + height, width = height // self.patch_size, width // self.patch_size + + hidden_states = self.patchfy(hidden_states) + len_frame = hidden_states.shape[1] + + t_expand, embedded_timestep = self.adaln_single(t_expand) + encoder_hidden_states = self.caption_projection( + self.caption_norm(encoder_hidden_states) + ) + + if encoder_hidden_states_2 is not None and hasattr(self, "clip_projection"): + clip_embedding, _ = self.clip_projection(encoder_hidden_states_2) + encoder_hidden_states = torch.cat( + [clip_embedding, encoder_hidden_states], dim=1 + ) + + hidden_states = rearrange( + hidden_states, "(b f) l d-> b (f l) d", b=bsz, f=frame, l=len_frame + ).contiguous() + encoder_hidden_states, attn_mask = self.prepare_attn_mask( + encoder_attention_mask, encoder_hidden_states, q_seqlen=frame * len_frame + ) + + cos_sin = self._get_rope( + (frame, height, width), hidden_states.dtype, hidden_states.device + ) + + hidden_states = self.block_forward( + hidden_states, + encoder_hidden_states, + t_expand=t_expand, + rope_positions=[frame, height, width], + cos_sin=cos_sin, + attn_mask=attn_mask, + parallel=self.parallel, + mask_strategy=mask_strategy, + ) + + hidden_states = rearrange( + hidden_states, "b (f l) d -> (b f) l d", b=bsz, f=frame, l=len_frame + ) + + embedded_timestep = repeat( + embedded_timestep, "b d -> (b f) d", f=frame + ).contiguous() + + shift, scale = ( + self.scale_shift_table[None] + embedded_timestep[:, None] + ).chunk(2, dim=1) + hidden_states = self.norm_out( + hidden_states, shift=shift.squeeze(1), scale=scale.squeeze(1) + ) + # Modulation + hidden_states, _ = self.proj_out(hidden_states) + + # unpatchify + hidden_states = hidden_states.reshape( + shape=( + -1, + height, + width, + self.patch_size, + self.patch_size, + self.out_channels, + ) + ) + + hidden_states = rearrange(hidden_states, "n h w p q c -> n c h p w q") + output = hidden_states.reshape( + shape=( + -1, + self.out_channels, + height * self.patch_size, + width * self.patch_size, + ) + ) + + output = rearrange(output, "(b f) c h w -> b c f h w", f=frame) + return output + + +EntryClass = StepVideoModel diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py new file mode 100644 index 00000000..cb674e49 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -0,0 +1,945 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any + +import numpy as np +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.parallel_state import get_sp_world_size +from sglang.multimodal_gen.runtime.layers.attention import ( + UlyssesAttention_VSA, + USPAttention, +) +from sglang.multimodal_gen.runtime.layers.layernorm import ( + FP32LayerNorm, + LayerNormScaleShift, + RMSNorm, + ScaleResidual, + 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.rotary_embedding import ( + NDRotaryEmbedding, + _apply_rotary_emb, +) +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.logging_utils import init_logger + +logger = init_logger(__name__) + + +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, + supported_attention_backends: set[AttentionBackendEnum] | 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 + + # layers + self.to_q = ReplicatedLinear(dim, dim) + self.to_k = ReplicatedLinear(dim, dim) + self.to_v = ReplicatedLinear(dim, dim) + self.to_out = ReplicatedLinear(dim, dim) + 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() + + # 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=supported_attention_backends, + ) + + def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int): + 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] + """ + pass + + +class WanT2VCrossAttention(WanSelfAttention): + + def forward(self, x, context, context_lens, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + crossattn_cache["k"] = k + crossattn_cache["v"] = v + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + + # 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, + supported_attention_backends: set[AttentionBackendEnum] | None = None, + ) -> None: + # VSA should not be in supported_attention_backends + super().__init__( + dim, + num_heads, + window_size, + qk_norm, + eps, + supported_attention_backends=supported_attention_backends, + ) + + self.add_k_proj = ReplicatedLinear(dim, dim) + self.add_v_proj = ReplicatedLinear(dim, dim) + self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_added_q = 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:] + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(b, -1, n, d) + v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) + img_x = self.attn(q, k_img, v_img) + # compute attention + 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 = "", + ): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True) + self.to_k = ReplicatedLinear(dim, dim, bias=True) + self.to_v = ReplicatedLinear(dim, dim, bias=True) + + self.to_out = ReplicatedLinear(dim, dim, bias=True) + self.attn1 = USPAttention( + num_heads=num_heads, + head_size=dim // num_heads, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn1", + ) + + 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, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32, + compute_dtype=torch.float32, + ) + + # 2. Cross-attention + if added_kv_proj_dim is not None: + # I2V + self.attn2 = WanI2VCrossAttention( + dim, + num_heads, + qk_norm=qk_norm, + eps=eps, + supported_attention_backends=supported_attention_backends, + ) + else: + # T2V + self.attn2 = WanT2VCrossAttention( + dim, + num_heads, + qk_norm=qk_norm, + eps=eps, + supported_attention_backends=supported_attention_backends, + ) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32, + ) + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.mlp_residual = ScaleResidual() + + 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 + norm1 = self.norm1(hidden_states.float()) + norm_hidden_states = (norm1 * (1 + scale_msa) + shift_msa).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)) + + # Apply rotary embeddings + cos, sin = freqs_cis + 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(hidden_states, ff_output, c_gate_msa) + 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 = "", + ): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True) + self.to_k = ReplicatedLinear(dim, dim, bias=True) + self.to_v = ReplicatedLinear(dim, dim, bias=True) + self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True) + + self.to_out = ReplicatedLinear(dim, dim, bias=True) + self.attn1 = UlyssesAttention_VSA( + num_heads=num_heads, + head_size=dim // num_heads, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn1", + ) + 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, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32, + compute_dtype=torch.float32, + ) + + if AttentionBackendEnum.VIDEO_SPARSE_ATTN in supported_attention_backends: + supported_attention_backends.remove(AttentionBackendEnum.VIDEO_SPARSE_ATTN) + # 2. Cross-attention + if added_kv_proj_dim is not None: + # I2V + self.attn2 = WanI2VCrossAttention( + dim, + num_heads, + qk_norm=qk_norm, + eps=eps, + supported_attention_backends=supported_attention_backends, + ) + else: + # T2V + self.attn2 = WanT2VCrossAttention( + dim, + num_heads, + qk_norm=qk_norm, + eps=eps, + supported_attention_backends=supported_attention_backends, + ) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32, + ) + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.mlp_residual = ScaleResidual() + + 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.float()) * (1 + scale_msa) + shift_msa + ).to(orig_dtype) + 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 + 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(hidden_states, ff_output, c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) + + return hidden_states + + +class WanTransformer3DModel(CachableDiT): + _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]) -> 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"{config.prefix}.blocks.{i}", + ) + for i in range(config.num_layers) + ] + ) + + # 4. Output norm & projection + self.norm_out = LayerNormScaleShift( + inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32, + compute_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 + ) + + # For type checking + self.previous_e0_even = None + self.previous_e0_odd = None + self.previous_residual_even = None + self.previous_residual_odd = None + self.is_even = True + 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.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() else torch.float64, + ) + + 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 + 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 + + # 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) + # 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 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 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 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 enable_teacache: + self.maybe_cache_states(hidden_states, original_hidden_states) + # 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: + if self.is_even: + self.previous_residual_even = ( + hidden_states.squeeze(0) - original_hidden_states + ) + else: + self.previous_residual_odd = ( + hidden_states.squeeze(0) - 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 or not forward_batch.enable_teacache: + return False + teacache_params = forward_batch.teacache_params + assert teacache_params is not None, "teacache_params is not initialized" + assert isinstance( + teacache_params, WanTeaCacheParams + ), "teacache_params is not a WanTeaCacheParams" + current_timestep = forward_context.current_timestep + num_inference_steps = forward_batch.num_inference_steps + + # initialize the coefficients, cutoff_steps, and ret_steps + coefficients = teacache_params.coefficients + use_ret_steps = teacache_params.use_ret_steps + cutoff_steps = teacache_params.get_cutoff_steps(num_inference_steps) + ret_steps = teacache_params.ret_steps + teacache_thresh = teacache_params.teacache_thresh + + if current_timestep == 0: + self.cnt = 0 + + timestep_proj = kwargs["timestep_proj"] + temb = kwargs["temb"] + modulated_inp = timestep_proj if use_ret_steps else temb + + if self.cnt % 2 == 0: # even -> condition + self.is_even = True + if self.cnt < ret_steps or self.cnt >= cutoff_steps: + self.should_calc_even = True + self.accumulated_rel_l1_distance_even = 0 + else: + assert ( + self.previous_e0_even is not None + ), "previous_e0_even is not initialized" + assert ( + self.accumulated_rel_l1_distance_even is not None + ), "accumulated_rel_l1_distance_even is not initialized" + rescale_func = np.poly1d(coefficients) + self.accumulated_rel_l1_distance_even += rescale_func( + ( + (modulated_inp - self.previous_e0_even).abs().mean() + / self.previous_e0_even.abs().mean() + ) + .cpu() + .item() + ) + if self.accumulated_rel_l1_distance_even < teacache_thresh: + self.should_calc_even = False + else: + self.should_calc_even = True + self.accumulated_rel_l1_distance_even = 0 + self.previous_e0_even = modulated_inp.clone() + + else: # odd -> unconditon + self.is_even = False + if self.cnt < ret_steps or self.cnt >= cutoff_steps: + self.should_calc_odd = True + self.accumulated_rel_l1_distance_odd = 0 + else: + assert ( + self.previous_e0_odd is not None + ), "previous_e0_odd is not initialized" + assert ( + self.accumulated_rel_l1_distance_odd is not None + ), "accumulated_rel_l1_distance_odd is not initialized" + rescale_func = np.poly1d(coefficients) + self.accumulated_rel_l1_distance_odd += rescale_func( + ( + (modulated_inp - self.previous_e0_odd).abs().mean() + / self.previous_e0_odd.abs().mean() + ) + .cpu() + .item() + ) + if self.accumulated_rel_l1_distance_odd < teacache_thresh: + self.should_calc_odd = False + else: + self.should_calc_odd = True + self.accumulated_rel_l1_distance_odd = 0 + self.previous_e0_odd = modulated_inp.clone() + self.cnt += 1 + should_skip_forward = False + if self.is_even: + if not self.should_calc_even: + should_skip_forward = True + else: + if not self.should_calc_odd: + should_skip_forward = True + + return should_skip_forward + + def retrieve_cached_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.is_even: + return hidden_states + self.previous_residual_even + else: + return hidden_states + self.previous_residual_odd + + +EntryClass = WanTransformer3DModel diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/zimage.py new file mode 100644 index 00000000..bdc38d52 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -0,0 +1,605 @@ +import math +from typing import Any, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig +from sglang.multimodal_gen.runtime.layers.attention import USPAttention +from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm +from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +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( + [ + ReplicatedLinear(frequency_embedding_size, mid_size, bias=True), + nn.SiLU(), + ReplicatedLinear(mid_size, out_size, bias=True), + ] + ) + + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + with torch.amp.autocast("cuda", 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__() + self.w1 = ReplicatedLinear(dim, hidden_dim, bias=False) + self.w2 = ReplicatedLinear(hidden_dim, dim, bias=False) + self.w3 = ReplicatedLinear(dim, hidden_dim, bias=False) + + def _forward_silu_gating(self, x1, x3): + return F.silu(x1) * x3 + + def forward(self, x): + x1, _ = self.w1(x) + x3, _ = self.w3(x) + out, _ = self.w2(self._forward_silu_gating(x1, x3)) + 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, + ) -> None: + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = dim // num_heads + self.qk_norm = qk_norm + + self.to_q = ReplicatedLinear(dim, dim, bias=False) + self.to_k = ReplicatedLinear(dim, self.head_dim * num_kv_heads, bias=False) + self.to_v = ReplicatedLinear(dim, self.head_dim * num_kv_heads, bias=False) + + 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([ReplicatedLinear(dim, dim, bias=False)]) + + self.attn = USPAttention( + num_heads=num_heads, + head_size=self.head_dim, + num_kv_heads=num_kv_heads, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends={ + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, + }, + ) + + def forward( + self, + hidden_states: torch.Tensor, + freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ): + 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.num_heads, self.head_dim) + k = k.view(*k.shape[:-1], self.num_kv_heads, self.head_dim) + v = v.view(*v.shape[:-1], self.num_kv_heads, self.head_dim) + + if self.norm_q is not None: + q = self.norm_q(q) + if self.norm_k is not None: + k = self.norm_k(k) + + # Apply RoPE + def apply_rotary_emb( + x_in: torch.Tensor, + freqs_cis: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + cos, sin = freqs_cis + x_out = _apply_rotary_emb(x_in, cos, sin, is_neox_style=False) + return x_out + + if freqs_cis is not None: + q = apply_rotary_emb(q, freqs_cis) + k = apply_rotary_emb(k, freqs_cis) + + 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, + ): + 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, + ) + + self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8)) + + 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 = ReplicatedLinear(hidden_size, out_channels, bias=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): + _supports_gradient_checkpointing = True + _no_split_modules = ["ZImageTransformerBlock"] + + def __init__( + self, + config: ZImageDitConfig, + hf_config: dict[str, Any], + ) -> 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 = ReplicatedLinear( + f_patch_size * patch_size * patch_size * self.in_channels, + self.dim, + bias=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, + ) + 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, + ) + 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, + ) + 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 + ) + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/base.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/base.py new file mode 100644 index 00000000..7b3eece9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/base.py @@ -0,0 +1,71 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/bert.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/bert.py new file mode 100644 index 00000000..5a423e51 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/bert.py @@ -0,0 +1,46 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/clip.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/clip.py new file mode 100644 index 00000000..ec80e387 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/clip.py @@ -0,0 +1,700 @@ +# 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.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=False, + 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, + ): + """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, + ) + 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) -> torch.Tensor: + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, _ = self.self_attn(hidden_states=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 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 + ) -> 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) + 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: + r""" + Returns: + + """ + 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, + # attention_mask=attention_mask, + # causal_attention_mask=causal_attention_mask, + # output_attentions=output_attentions, + return_all_hidden_states=output_hidden_states, + # return_dict=return_dict, + ) + + 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] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/llama.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/llama.py new file mode 100644 index 00000000..ea208f12 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/llama.py @@ -0,0 +1,459 @@ +# 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 = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py new file mode 100644 index 00000000..fef6ece6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py @@ -0,0 +1,459 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py new file mode 100644 index 00000000..c354f923 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py @@ -0,0 +1,1180 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +from types import SimpleNamespace + +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): + super().__init__() + self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config( + config.vision_config + ) + self.language_model = Qwen2_5_VLTextModel(config.text_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) + 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 DotDict(dict): + def __init__(self, mapping): + super().__init__() + for key, value in mapping.items(): + if isinstance(value, dict): + value = DotDict(value) # 递归转换 + elif isinstance(value, list): + # 如果是 list，且元素是 dict 也递归转换 + value = [ + DotDict(item) if isinstance(item, dict) else item for item in value + ] + self[key] = value + + def __getattr__(self, item): + try: + return self[item] + except KeyError: + raise AttributeError(f"No attribute '{item}'") + + def __setattr__(self, key, value): + self[key] = value + + def __delattr__(self, key): + del self[key] + + +def dict_to_namespace(d): + for k, v in d.items(): + if isinstance(v, dict): + d[k] = dict_to_namespace(v) + elif isinstance(v, list): + d[k] = [dict_to_namespace(i) if isinstance(i, dict) else i for i in v] + return SimpleNamespace(**d) + + +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) + config = config.arch_config + self.model = Qwen2_5_VLModel(config) + self.lm_head = nn.Linear( + config.text_config.hidden_size, config.text_config.vocab_size, bias=False + ) + + 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: + 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: + print(params_dict.keys()) + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py new file mode 100644 index 00000000..dae653ea --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py @@ -0,0 +1,614 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +# type: ignore +# Copyright 2025 StepFun 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. +# ============================================================================== +import os +from functools import wraps + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers import PretrainedConfig, PreTrainedModel + +from sglang.multimodal_gen.runtime.models.dits.stepvideo import StepVideoRMSNorm + + +class EmptyInitOnDevice(torch.overrides.TorchFunctionMode): + + def __init__(self, device=None): + self.device = device + + def __torch_function__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + if getattr(func, "__module__", None) == "torch.nn.init": + if "tensor" in kwargs: + return kwargs["tensor"] + else: + return args[0] + if ( + self.device is not None + and func in torch.utils._device._device_constructors() + and kwargs.get("device") is None + ): + kwargs["device"] = self.device + return func(*args, **kwargs) + + +def with_empty_init(func): + + @wraps(func) + def wrapper(*args, **kwargs): + with EmptyInitOnDevice("cpu"): + return func(*args, **kwargs) + + return wrapper + + +class LLaMaEmbedding(nn.Module): + """Language model embeddings. + + Arguments: + hidden_size: hidden size + vocab_size: vocabulary size + max_sequence_length: maximum size of sequence. This + is used for positional embedding + embedding_dropout_prob: dropout probability for embeddings + init_method: weight initialization method + num_tokentypes: size of the token-type embeddings. 0 value + will ignore this embedding + """ + + def __init__( + self, + cfg, + ): + super().__init__() + self.hidden_size = cfg.hidden_size + self.params_dtype = cfg.params_dtype + self.fp32_residual_connection = cfg.fp32_residual_connection + self.embedding_weights_in_fp32 = cfg.embedding_weights_in_fp32 + self.word_embeddings = torch.nn.Embedding( + cfg.padded_vocab_size, + self.hidden_size, + ) + self.embedding_dropout = torch.nn.Dropout(cfg.hidden_dropout) + + def forward(self, input_ids): + # Embeddings. + if self.embedding_weights_in_fp32: + self.word_embeddings = self.word_embeddings.to(torch.float32) + embeddings = self.word_embeddings(input_ids) + if self.embedding_weights_in_fp32: + embeddings = embeddings.to(self.params_dtype) + self.word_embeddings = self.word_embeddings.to(self.params_dtype) + + # Data format change to avoid explicit transposes : [b s h] --> [s b h]. + embeddings = embeddings.transpose(0, 1).contiguous() + + # If the input flag for fp32 residual connection is set, convert for float. + if self.fp32_residual_connection: + embeddings = embeddings.float() + + # Dropout. + embeddings = self.embedding_dropout(embeddings) + + return embeddings + + +class StepChatTokenizer: + """Step Chat Tokenizer""" + + def __init__( + self, + model_file, + name="StepChatTokenizer", + bot_token="<|BOT|>", # Begin of Turn + eot_token="<|EOT|>", # End of Turn + call_start_token="<|CALL_START|>", # Call Start + call_end_token="<|CALL_END|>", # Call End + think_start_token="<|THINK_START|>", # Think Start + think_end_token="<|THINK_END|>", # Think End + mask_start_token="<|MASK_1e69f|>", # Mask start + mask_end_token="<|UNMASK_1e69f|>", # Mask end + ): + import sentencepiece + + self._tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file) + + self._vocab = {} + self._inv_vocab = {} + + self._special_tokens = {} + self._inv_special_tokens = {} + + self._t5_tokens = [] + + for idx in range(self._tokenizer.get_piece_size()): + text = self._tokenizer.id_to_piece(idx) + self._inv_vocab[idx] = text + self._vocab[text] = idx + + if self._tokenizer.is_control(idx) or self._tokenizer.is_unknown(idx): + self._special_tokens[text] = idx + self._inv_special_tokens[idx] = text + + self._unk_id = self._tokenizer.unk_id() + self._bos_id = self._tokenizer.bos_id() + self._eos_id = self._tokenizer.eos_id() + + for token in [ + bot_token, + eot_token, + call_start_token, + call_end_token, + think_start_token, + think_end_token, + ]: + assert token in self._vocab, f"Token '{token}' not found in tokenizer" + assert ( + token in self._special_tokens + ), f"Token '{token}' is not a special token" + + for token in [mask_start_token, mask_end_token]: + assert token in self._vocab, f"Token '{token}' not found in tokenizer" + + self._bot_id = self._tokenizer.piece_to_id(bot_token) + self._eot_id = self._tokenizer.piece_to_id(eot_token) + self._call_start_id = self._tokenizer.piece_to_id(call_start_token) + self._call_end_id = self._tokenizer.piece_to_id(call_end_token) + self._think_start_id = self._tokenizer.piece_to_id(think_start_token) + self._think_end_id = self._tokenizer.piece_to_id(think_end_token) + self._mask_start_id = self._tokenizer.piece_to_id(mask_start_token) + self._mask_end_id = self._tokenizer.piece_to_id(mask_end_token) + + self._underline_id = self._tokenizer.piece_to_id("\u2581") + + @property + def vocab(self): + return self._vocab + + @property + def inv_vocab(self): + return self._inv_vocab + + @property + def vocab_size(self): + return self._tokenizer.vocab_size() + + def tokenize(self, text: str) -> list[int]: + return self._tokenizer.encode_as_ids(text) + + def detokenize(self, token_ids: list[int]) -> str: + return self._tokenizer.decode_ids(token_ids) + + +class Tokens: + + def __init__( + self, input_ids, cu_input_ids, attention_mask, cu_seqlens, max_seq_len + ) -> None: + self.input_ids = input_ids + self.attention_mask = attention_mask + self.cu_input_ids = cu_input_ids + self.cu_seqlens = cu_seqlens + self.max_seq_len = max_seq_len + + def to(self, device): + self.input_ids = self.input_ids.to(device) + self.attention_mask = self.attention_mask.to(device) + self.cu_input_ids = self.cu_input_ids.to(device) + self.cu_seqlens = self.cu_seqlens.to(device) + return self + + +class Wrapped_StepChatTokenizer(StepChatTokenizer): + + def __call__( + self, + text, + max_length=320, + padding="max_length", + truncation=True, + return_tensors="pt", + ): + # [bos, ..., eos, pad, pad, ..., pad] + self.BOS = 1 + self.EOS = 2 + self.PAD = 2 + out_tokens = [] + attn_mask = [] + if len(text) == 0: + part_tokens = [self.BOS] + [self.EOS] + valid_size = len(part_tokens) + if len(part_tokens) < max_length: + part_tokens += [self.PAD] * (max_length - valid_size) + out_tokens.append(part_tokens) + attn_mask.append([1] * valid_size + [0] * (max_length - valid_size)) + else: + for part in text: + part_tokens = self.tokenize(part) + part_tokens = part_tokens[ + : (max_length - 2) + ] # leave 2 space for bos and eos + part_tokens = [self.BOS] + part_tokens + [self.EOS] + valid_size = len(part_tokens) + if len(part_tokens) < max_length: + part_tokens += [self.PAD] * (max_length - valid_size) + out_tokens.append(part_tokens) + attn_mask.append([1] * valid_size + [0] * (max_length - valid_size)) + + out_tokens = torch.tensor(out_tokens, dtype=torch.long) + attn_mask = torch.tensor(attn_mask, dtype=torch.long) + + # padding y based on tp size + padded_len = 0 + padded_flag = False + if padded_len > 0: + padded_flag = True + if padded_flag: + pad_tokens = torch.tensor( + [[self.PAD] * max_length], device=out_tokens.device + ) + pad_attn_mask = torch.tensor( + [[1] * padded_len + [0] * (max_length - padded_len)], + device=attn_mask.device, + ) + out_tokens = torch.cat([out_tokens, pad_tokens], dim=0) + attn_mask = torch.cat([attn_mask, pad_attn_mask], dim=0) + + # cu_seqlens + cu_out_tokens = out_tokens.masked_select(attn_mask != 0).unsqueeze(0) + seqlen = attn_mask.sum(dim=1).tolist() + cu_seqlens = torch.cumsum(torch.tensor([0] + seqlen), 0).to( + device=out_tokens.device, dtype=torch.int32 + ) + max_seq_len = max(seqlen) + return Tokens(out_tokens, cu_out_tokens, attn_mask, cu_seqlens, max_seq_len) + + +def flash_attn_func( + q, + k, + v, + dropout_p=0.0, + softmax_scale=None, + causal=True, + return_attn_probs=False, + tp_group_rank=0, + tp_group_size=1, +): + softmax_scale = q.size(-1) ** (-0.5) if softmax_scale is None else softmax_scale + return torch.ops.Optimus.fwd( + q, + k, + v, + None, + dropout_p, + softmax_scale, + causal, + return_attn_probs, + None, + tp_group_rank, + tp_group_size, + )[0] + + +class FlashSelfAttention(torch.nn.Module): + + def __init__( + self, + attention_dropout=0.0, + ): + super().__init__() + self.dropout_p = attention_dropout + + def forward(self, q, k, v, cu_seqlens=None, max_seq_len=None): + if cu_seqlens is None: + output = flash_attn_func(q, k, v, dropout_p=self.dropout_p) + else: + raise ValueError("cu_seqlens is not supported!") + + return output + + +def safediv(n, d): + q, r = divmod(n, d) + assert r == 0 + return q + + +class MultiQueryAttention(nn.Module): + + def __init__(self, cfg, layer_id=None): + super().__init__() + + self.head_dim = cfg.hidden_size // cfg.num_attention_heads + self.max_seq_len = cfg.seq_length + self.use_flash_attention = cfg.use_flash_attn + assert self.use_flash_attention, "FlashAttention is required!" + + self.n_groups = cfg.num_attention_groups + self.tp_size = 1 + self.n_local_heads = cfg.num_attention_heads + self.n_local_groups = self.n_groups + + self.wqkv = nn.Linear( + cfg.hidden_size, + cfg.hidden_size + self.head_dim * 2 * self.n_groups, + bias=False, + ) + self.wo = nn.Linear( + cfg.hidden_size, + cfg.hidden_size, + bias=False, + ) + + # assert self.use_flash_attention, 'non-Flash attention not supported yet.' + self.core_attention = FlashSelfAttention( + attention_dropout=cfg.attention_dropout + ) + # self.core_attention = LocalAttention( + # num_heads = self.n_local_heads, + # head_size = self.head_dim, + # # num_kv_heads = self.n_local_groups, + # casual = True, + # supported_attention_backends = [_Backend.FLASH_ATTN, _Backend.TORCH_SDPA], # RIVER TODO + # ) + self.layer_id = layer_id + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + max_seq_len: torch.Tensor | None, + ): + seqlen, bsz, dim = x.shape + xqkv = self.wqkv(x) + + xq, xkv = torch.split( + xqkv, + (dim // self.tp_size, self.head_dim * 2 * self.n_groups // self.tp_size), + dim=-1, + ) + + # gather on 1st dimension + xq = xq.view(seqlen, bsz, self.n_local_heads, self.head_dim) + xkv = xkv.view(seqlen, bsz, self.n_local_groups, 2 * self.head_dim) + xk, xv = xkv.chunk(2, -1) + + # rotary embedding + flash attn + xq = rearrange(xq, "s b h d -> b s h d") + xk = rearrange(xk, "s b h d -> b s h d") + xv = rearrange(xv, "s b h d -> b s h d") + + # q_per_kv = self.n_local_heads // self.n_local_groups + # if q_per_kv > 1: + # b, s, h, d = xk.size() + # if h == 1: + # xk = xk.expand(b, s, q_per_kv, d) + # xv = xv.expand(b, s, q_per_kv, d) + # else: + # ''' To cover the cases where h > 1, we have + # the following implementation, which is equivalent to: + # xk = xk.repeat_interleave(q_per_kv, dim=-2) + # xv = xv.repeat_interleave(q_per_kv, dim=-2) + # but can avoid calling aten::item() that involves cpu. + # ''' + # idx = torch.arange(q_per_kv * h, device=xk.device).reshape(q_per_kv, -1).permute(1, 0).flatten() + # xk = torch.index_select(xk.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() + # xv = torch.index_select(xv.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() + if self.use_flash_attention: + output = self.core_attention(xq, xk, xv) + # reduce-scatter only support first dimension now + output = rearrange(output, "b s h d -> s b (h d)").contiguous() + else: + xq, xk, xv = [ + rearrange(x, "b s ... -> s b ...").contiguous() for x in (xq, xk, xv) + ] + output = self.core_attention(xq, xk, xv) # , mask) + output = self.wo(output) + return output + + +class FeedForward(nn.Module): + + def __init__( + self, + cfg, + dim: int, + hidden_dim: int, + layer_id: int, + multiple_of: int = 256, + ): + super().__init__() + + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + + def swiglu(x): + x = torch.chunk(x, 2, dim=-1) + return F.silu(x[0]) * x[1] + + self.swiglu = swiglu + + self.w1 = nn.Linear( + dim, + 2 * hidden_dim, + bias=False, + ) + self.w2 = nn.Linear( + hidden_dim, + dim, + bias=False, + ) + + def forward(self, x): + x = self.swiglu(self.w1(x)) + output = self.w2(x) + return output + + +class TransformerBlock(nn.Module): + + def __init__(self, cfg, layer_id: int): + super().__init__() + + self.n_heads = cfg.num_attention_heads + self.dim = cfg.hidden_size + self.head_dim = cfg.hidden_size // cfg.num_attention_heads + self.attention = MultiQueryAttention( + cfg, + layer_id=layer_id, + ) + + self.feed_forward = FeedForward( + cfg, + dim=cfg.hidden_size, + hidden_dim=cfg.ffn_hidden_size, + layer_id=layer_id, + ) + self.layer_id = layer_id + self.attention_norm = StepVideoRMSNorm( + cfg.hidden_size, + eps=cfg.layernorm_epsilon, + ) + self.ffn_norm = StepVideoRMSNorm( + cfg.hidden_size, + eps=cfg.layernorm_epsilon, + ) + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + max_seq_len: torch.Tensor | None, + ): + residual = self.attention.forward( + self.attention_norm(x), mask, cu_seqlens, max_seq_len + ) + h = x + residual + ffn_res = self.feed_forward.forward(self.ffn_norm(h)) + out = h + ffn_res + return out + + +class Transformer(nn.Module): + + def __init__( + self, + config, + max_seq_size=8192, + ): + super().__init__() + self.num_layers = config.num_layers + self.layers = self._build_layers(config) + + def _build_layers(self, config): + layers = torch.nn.ModuleList() + for layer_id in range(self.num_layers): + layers.append( + TransformerBlock( + config, + layer_id=layer_id + 1, + ) + ) + return layers + + def forward( + self, + hidden_states, + attention_mask, + cu_seqlens=None, + max_seq_len=None, + ): + + if max_seq_len is not None and not isinstance(max_seq_len, torch.Tensor): + max_seq_len = torch.tensor(max_seq_len, dtype=torch.int32, device="cpu") + + for lid, layer in enumerate(self.layers): + hidden_states = layer( + hidden_states, + attention_mask, + cu_seqlens, + max_seq_len, + ) + return hidden_states + + +class Step1Model(PreTrainedModel): + config_class = PretrainedConfig + + @with_empty_init + def __init__( + self, + config, + ): + super().__init__(config) + self.tok_embeddings = LLaMaEmbedding(config) + self.transformer = Transformer(config) + + def forward( + self, + input_ids=None, + attention_mask=None, + ): + + hidden_states = self.tok_embeddings(input_ids) + + hidden_states = self.transformer( + hidden_states, + attention_mask, + ) + return hidden_states + + +class STEP1TextEncoder(torch.nn.Module): + + def __init__(self, model_dir, max_length=320): + super().__init__() + self.max_length = max_length + self.text_tokenizer = Wrapped_StepChatTokenizer( + os.path.join(model_dir, "step1_chat_tokenizer.model") + ) + text_encoder = Step1Model.from_pretrained(model_dir) + self.text_encoder = text_encoder.eval().to(torch.bfloat16) + + @torch.no_grad + def forward(self, prompts, with_mask=True, max_length=None): + self.device = next(self.text_encoder.parameters()).device + + with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16): + if type(prompts) is str: + prompts = [prompts] + txt_tokens = self.text_tokenizer( + prompts, + max_length=max_length or self.max_length, + padding="max_length", + truncation=True, + return_tensors="pt", + ) + y = self.text_encoder( + txt_tokens.input_ids.to(self.device), + attention_mask=( + txt_tokens.attention_mask.to(self.device) if with_mask else None + ), + ) + y_mask = txt_tokens.attention_mask + return y.transpose(0, 1), y_mask + + +EntryClass = STEP1TextEncoder diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/t5.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/t5.py new file mode 100644 index 00000000..048308ad --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/t5.py @@ -0,0 +1,716 @@ +# 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_tp_rank, get_tp_world_size +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.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" + + +_seen_keys = set() # 用集合记录已经出现过的 key + + +@dataclass +class AttentionMetadata: + attn_bias: torch.Tensor + + +class T5DenseActDense(nn.Module): + + def __init__( + self, config: T5Config, quant_config: QuantizationConfig | None = None + ): + super().__init__() + self.wi = MergedColumnParallelLinear(config.d_model, [config.d_ff], bias=False) + self.wo = RowParallelLinear( + config.d_ff, config.d_model, bias=False, quant_config=quant_config + ) + 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__() + self.wi_0 = MergedColumnParallelLinear( + config.d_model, [config.d_ff], bias=False, quant_config=quant_config + ) + self.wi_1 = MergedColumnParallelLinear( + config.d_model, [config.d_ff], bias=False, quant_config=quant_config + ) + # 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 + ) + 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. + tp_world_size = get_tp_world_size() + assert config.num_heads % tp_world_size == 0 + self.n_heads = config.num_heads // 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.d_model // self.total_num_heads, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + 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, + ) + self.o = RowParallelLinear( + self.d_model, + self.d_model, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + @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.d_model // self.total_num_heads + 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 get_tp_world_size() > 1: + rank = get_tp_rank() + 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 + + self.shared = VocabParallelEmbedding( + config.vocab_size, config.d_model, org_num_embeddings=config.vocab_size + ) + + 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 + + self.shared = VocabParallelEmbedding( + config.vocab_size, config.d_model, org_num_embeddings=config.vocab_size + ) + + 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] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/vision.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/vision.py new file mode 100644 index 00000000..3150abf1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/encoders/vision.py @@ -0,0 +1,96 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/parameter.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/parameter.py new file mode 100644 index 00000000..ba9b42c6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/parameter.py @@ -0,0 +1,423 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/registry.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/registry.py new file mode 100644 index 00000000..2ee26257 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/registry.py @@ -0,0 +1,365 @@ +# 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"), +} + + +@lru_cache(maxsize=None) +def _discover_and_register_models() -> dict[str, tuple[str, str, str]]: + discovered_models = _IMAGE_ENCODER_MODELS + 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 + + for node in ast.walk(tree): + 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 first_class_def is None and isinstance(node, ast.ClassDef): + first_class_def = node + 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, + ) + + except Exception as e: + logger.warning(f"Could not parse {filepath} to find models: {e}") + + 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 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: + raise Exception( + f"Unsupported model architecture: {arch}. Registered architectures: {self.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() + } +) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/base.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/base.py new file mode 100644 index 00000000..eb4e3bdd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/base.py @@ -0,0 +1,37 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py new file mode 100644 index 00000000..d184802b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -0,0 +1,698 @@ +# 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 + + Args: + sample (`torch.FloatTensor`): + The input sample. + timestep (`int`, *optional*): + The current timestep in the diffusion chain. + + Returns: + `torch.FloatTensor`: + A scaled input sample. + """ + # 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py new file mode 100644 index 00000000..1e6b84e0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py @@ -0,0 +1,853 @@ +# 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 + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + 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`", + ) + + 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": + 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 + + def __len__(self): + return self.config.num_train_timesteps + + +EntryClass = FlowUniPCMultistepScheduler diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_self_forcing_flow_match.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_self_forcing_flow_match.py new file mode 100644 index 00000000..9a4749a6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_self_forcing_flow_match.py @@ -0,0 +1,142 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py new file mode 100644 index 00000000..df5e9b83 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py @@ -0,0 +1,1207 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/utils.py new file mode 100644 index 00000000..73a2f062 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/utils.py @@ -0,0 +1,148 @@ +# 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 + + Args: + x (torch.Tensor): input tensor. + shift (torch.Tensor, optional): shift tensor. Defaults to None. + scale (torch.Tensor, optional): scale tensor. Defaults to None. + + Returns: + torch.Tensor: the output tensor after modulate. + """ + 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder.py new file mode 100644 index 00000000..91fa447e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder.py @@ -0,0 +1,585 @@ +# 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. + 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 + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py new file mode 100644 index 00000000..9df72c09 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py @@ -0,0 +1,524 @@ +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(nn.Module, 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__() + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py new file mode 100644 index 00000000..e860b313 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py @@ -0,0 +1,1170 @@ +# 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.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", + ): + 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(3, 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", + ): + 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, 3, 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(nn.Module): + 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__() + 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.encoder = QwenImageEncoder3d( + base_dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout + ) + 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 + ) + + self.spatial_compression_ratio = 2 ** len(self.temperal_downsample) + + # 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) + ) + latents_std_tensor = torch.tensor(config.arch_config.latents_std, dtype=dtype, device=cuda_device) + self.scaling_factor = (1.0 / latents_std_tensor).view(1, latent_channels, 1, 1, 1) + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/common.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/common.py new file mode 100644 index 00000000..58dad250 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/common.py @@ -0,0 +1,648 @@ +# 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 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): + 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: + 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 + + def to(self, device) -> "ParallelTiledVAE": + # TODO: implement this + return self + + @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 + ) -> 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=[1, 2, 3], + ) + 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=[1, 2, 3], + ) + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py new file mode 100644 index 00000000..a9de61a5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py @@ -0,0 +1,852 @@ +# 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(nn.Module, 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/stepvideovae.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/stepvideovae.py new file mode 100644 index 00000000..d202b7a6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/stepvideovae.py @@ -0,0 +1,1184 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 StepFun 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. +# ============================================================================== +from typing import Any + +import torch +from einops import rearrange +from torch import nn +from torch.nn import functional as F + +from sglang.multimodal_gen.configs.models.vaes import StepVideoVAEConfig +from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE + + +def base_group_norm(x, norm_layer, act_silu=False, channel_last=False) -> torch.Tensor: + if hasattr(base_group_norm, "spatial") and base_group_norm.spatial: + assert channel_last + x_shape = x.shape + x = x.flatten(0, 1) + if channel_last: + # Permute to NCHW format + x = x.permute(0, 3, 1, 2) + + out = F.group_norm( + x.contiguous(), + norm_layer.num_groups, + norm_layer.weight, + norm_layer.bias, + norm_layer.eps, + ) + if act_silu: + out = F.silu(out) + + if channel_last: + # Permute back to NHWC format + out = out.permute(0, 2, 3, 1) + + out = out.view(x_shape) + else: + if channel_last: + # Permute to NCHW format + x = x.permute(0, 3, 1, 2) + out = F.group_norm( + x.contiguous(), + norm_layer.num_groups, + norm_layer.weight, + norm_layer.bias, + norm_layer.eps, + ) + if act_silu: + out = F.silu(out) + if channel_last: + # Permute back to NHWC format + out = out.permute(0, 2, 3, 1) + return out + + +def base_conv2d(x, conv_layer, channel_last=False, residual=None) -> torch.Tensor: + if channel_last: + x = x.permute(0, 3, 1, 2) # NHWC to NCHW + out = F.conv2d( + x, + conv_layer.weight, + conv_layer.bias, + stride=conv_layer.stride, + padding=conv_layer.padding, + ) + if residual is not None: + if channel_last: + residual = residual.permute(0, 3, 1, 2) # NHWC to NCHW + out += residual + if channel_last: + out = out.permute(0, 2, 3, 1) # NCHW to NHWC + return out + + +def base_conv3d( + x, conv_layer, channel_last=False, residual=None, only_return_output=False +) -> torch.Tensor: + if only_return_output: + size = cal_outsize( + x.shape, conv_layer.weight.shape, conv_layer.stride, conv_layer.padding + ) + return torch.empty(size, device=x.device, dtype=x.dtype) + if channel_last: + x = x.permute(0, 4, 1, 2, 3) # NDHWC to NCDHW + out = F.conv3d( + x, + conv_layer.weight, + conv_layer.bias, + stride=conv_layer.stride, + padding=conv_layer.padding, + ) + if residual is not None: + if channel_last: + residual = residual.permute(0, 4, 1, 2, 3) # NDHWC to NCDHW + out += residual + if channel_last: + out = out.permute(0, 2, 3, 4, 1) # NCDHW to NDHWC + return out + + +def cal_outsize(input_sizes, kernel_sizes, stride, padding) -> list: + stride_d, stride_h, stride_w = stride + padding_d, padding_h, padding_w = padding + dilation_d, dilation_h, dilation_w = 1, 1, 1 + + in_d = input_sizes[1] + in_h = input_sizes[2] + in_w = input_sizes[3] + + kernel_d = kernel_sizes[2] + kernel_h = kernel_sizes[3] + kernel_w = kernel_sizes[4] + out_channels = kernel_sizes[0] + + out_d = calc_out_(in_d, padding_d, dilation_d, kernel_d, stride_d) + out_h = calc_out_(in_h, padding_h, dilation_h, kernel_h, stride_h) + out_w = calc_out_(in_w, padding_w, dilation_w, kernel_w, stride_w) + size = [input_sizes[0], out_d, out_h, out_w, out_channels] + return size + + +def calc_out_( + in_size: int, padding: int, dilation: int, kernel: int, stride: int +) -> int: + return (in_size + 2 * padding - dilation * (kernel - 1) - 1) // stride + 1 + + +def base_conv3d_channel_last(x, conv_layer, residual=None) -> torch.Tensor: + in_numel = x.numel() + out_numel = int(x.numel() * conv_layer.out_channels / conv_layer.in_channels) + if (in_numel >= 2**30) or (out_numel >= 2**30): + assert conv_layer.stride[0] == 1, "time split asks time stride = 1" + + B, T, H, W, C = x.shape + K = conv_layer.kernel_size[0] + + chunks = 4 + chunk_size = T // chunks + + if residual is None: + out_nhwc = base_conv3d( + x, + conv_layer, + channel_last=True, + residual=residual, + only_return_output=True, + ) + else: + out_nhwc = residual + + assert B == 1 + for i in range(chunks): + if i == chunks - 1: + xi = x[:1, chunk_size * i :] + out_nhwci = out_nhwc[:1, chunk_size * i :] + else: + xi = x[:1, chunk_size * i : chunk_size * (i + 1) + K - 1] + out_nhwci = out_nhwc[:1, chunk_size * i : chunk_size * (i + 1)] + if residual is not None: + if i == chunks - 1: + ri = residual[:1, chunk_size * i :] + else: + ri = residual[:1, chunk_size * i : chunk_size * (i + 1)] + else: + ri = None + out_nhwci.copy_(base_conv3d(xi, conv_layer, channel_last=True, residual=ri)) + else: + out_nhwc = base_conv3d(x, conv_layer, channel_last=True, residual=residual) + return out_nhwc + + +class Upsample2D(nn.Module): + + def __init__( + self, channels, use_conv=False, use_conv_transpose=False, out_channels=None + ) -> None: + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + + if use_conv: + self.conv = nn.Conv2d(self.channels, self.out_channels, 3, padding=1) + else: + assert "Not Supported" + self.conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1) + + def forward(self, x, output_size=None) -> torch.Tensor: + assert x.shape[-1] == self.channels + + if self.use_conv_transpose: + return self.conv(x) + + if output_size is None: + x = ( + F.interpolate( + x.permute(0, 3, 1, 2).to(memory_format=torch.channels_last), + scale_factor=2.0, + mode="nearest", + ) + .permute(0, 2, 3, 1) + .contiguous() + ) + else: + x = ( + F.interpolate( + x.permute(0, 3, 1, 2).to(memory_format=torch.channels_last), + size=output_size, + mode="nearest", + ) + .permute(0, 2, 3, 1) + .contiguous() + ) + + # x = self.conv(x) + x = base_conv2d(x, self.conv, channel_last=True) + return x + + +class Downsample2D(nn.Module): + + def __init__(self, channels, use_conv=False, out_channels=None, padding=1) -> None: + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + stride = 2 + + if use_conv: + self.conv = nn.Conv2d( + self.channels, self.out_channels, 3, stride=stride, padding=padding + ) + else: + assert self.channels == self.out_channels + self.conv = nn.AvgPool2d(kernel_size=stride, stride=stride) + + def forward(self, x) -> torch.Tensor: + assert x.shape[-1] == self.channels + if self.use_conv and self.padding == 0: + pad = (0, 0, 0, 1, 0, 1) + x = F.pad(x, pad, mode="constant", value=0) + + assert x.shape[-1] == self.channels + # x = self.conv(x) + x = base_conv2d(x, self.conv, channel_last=True) + return x + + +class CausalConv(nn.Module): + + def __init__(self, chan_in, chan_out, kernel_size, **kwargs) -> None: + super().__init__() + + if isinstance(kernel_size, int): + kernel_size = ( + kernel_size if isinstance(kernel_size, tuple) else ((kernel_size,) * 3) + ) + time_kernel_size, height_kernel_size, width_kernel_size = kernel_size + + self.dilation = kwargs.pop("dilation", 1) + self.stride = kwargs.pop("stride", 1) + if isinstance(self.stride, int): + self.stride = (self.stride, 1, 1) + time_pad = self.dilation * (time_kernel_size - 1) + max((1 - self.stride[0]), 0) + height_pad = height_kernel_size // 2 + width_pad = width_kernel_size // 2 + self.time_causal_padding = ( + width_pad, + width_pad, + height_pad, + height_pad, + time_pad, + 0, + ) + self.time_uncausal_padding = ( + width_pad, + width_pad, + height_pad, + height_pad, + 0, + 0, + ) + + self.conv = nn.Conv3d( + chan_in, + chan_out, + kernel_size, + stride=self.stride, + dilation=self.dilation, + **kwargs, + ) + self.chan_in = chan_in + self.chan_out = chan_out + self.is_first_run = True + + def forward(self, x, is_init=True, residual=None) -> torch.Tensor: + x = nn.functional.pad( + x, self.time_causal_padding if is_init else self.time_uncausal_padding + ) + x = self.conv(x) + if residual is not None: + x.add_(residual) + return x + + +class ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + factor: int, + ) -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor = factor + assert out_channels * factor**3 % in_channels == 0 + self.repeats = out_channels * factor**3 // in_channels + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor, + self.factor, + self.factor, + 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, + x.size(4) * self.factor, + x.size(6) * self.factor, + ) + x = x[:, :, self.factor - 1 :, :, :] + return x + + +class ConvPixelShuffleUpSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + factor: int, + ) -> None: + super().__init__() + self.factor = factor + out_ratio = factor**3 + self.conv = CausalConv( + in_channels, out_channels * out_ratio, kernel_size=kernel_size + ) + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = self.conv(x, is_init) + x = self.pixel_shuffle_3d(x, self.factor) + return x + + @staticmethod + def pixel_shuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: + batch_size, channels, depth, height, width = x.size() + new_channels = channels // (factor**3) + new_depth = depth * factor + new_height = height * factor + new_width = width * factor + + x = x.view( + batch_size, new_channels, factor, factor, factor, depth, height, width + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view(batch_size, new_channels, new_depth, new_height, new_width) + x = x[:, :, factor - 1 :, :, :] + return x + + +class ConvPixelUnshuffleDownSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + factor: int, + ) -> None: + super().__init__() + self.factor = factor + out_ratio = factor**3 + assert out_channels % out_ratio == 0 + self.conv = CausalConv( + in_channels, out_channels // out_ratio, kernel_size=kernel_size + ) + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = self.conv(x, is_init) + x = self.pixel_unshuffle_3d(x, self.factor) + return x + + @staticmethod + def pixel_unshuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: + pad = (0, 0, 0, 0, factor - 1, 0) # (left, right, top, bottom, front, back) + x = F.pad(x, pad) + B, C, D, H, W = x.shape + x = x.view(B, C, D // factor, factor, H // factor, factor, W // factor, factor) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view(B, C * factor**3, D // factor, H // factor, W // factor) + return x + + +class PixelUnshuffleChannelAveragingDownSampleLayer3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + factor: int, + ) -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor = factor + assert in_channels * factor**3 % out_channels == 0 + self.group_size = in_channels * factor**3 // out_channels + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + pad = ( + 0, + 0, + 0, + 0, + self.factor - 1, + 0, + ) # (left, right, top, bottom, front, back) + x = F.pad(x, pad) + B, C, D, H, W = x.shape + x = x.view( + B, + C, + D // self.factor, + self.factor, + H // self.factor, + self.factor, + W // self.factor, + self.factor, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, C * self.factor**3, D // self.factor, H // self.factor, W // self.factor + ) + x = x.view( + B, + self.out_channels, + self.group_size, + D // self.factor, + H // self.factor, + W // self.factor, + ) + x = x.mean(dim=2) + return x + + +def base_group_norm_with_zero_pad( + x, norm_layer, act_silu=True, pad_size=2 +) -> torch.Tensor: + out_shape = list(x.shape) + out_shape[1] += pad_size + out = torch.empty(out_shape, dtype=x.dtype, device=x.device) + out[:, pad_size:] = base_group_norm( + x, norm_layer, act_silu=act_silu, channel_last=True + ) + out[:, :pad_size] = 0 + return out + + +class CausalConvChannelLast(CausalConv): + time_causal_padding: tuple[Any, ...] + time_uncausal_padding: tuple[Any, ...] + + def __init__(self, chan_in, chan_out, kernel_size, **kwargs) -> None: + super().__init__(chan_in, chan_out, kernel_size, **kwargs) + + self.time_causal_padding = (0, 0) + self.time_causal_padding + self.time_uncausal_padding = (0, 0) + self.time_uncausal_padding + + def forward(self, x, is_init=True, residual=None) -> torch.Tensor: + if self.is_first_run: + self.is_first_run = False + # self.conv.weight = nn.Parameter(self.conv.weight.permute(0,2,3,4,1).contiguous()) + + x = nn.functional.pad( + x, self.time_causal_padding if is_init else self.time_uncausal_padding + ) + + x = base_conv3d_channel_last(x, self.conv, residual=residual) + return x + + +class CausalConvAfterNorm(CausalConv): + + def __init__(self, chan_in, chan_out, kernel_size, **kwargs) -> None: + super().__init__(chan_in, chan_out, kernel_size, **kwargs) + + if self.time_causal_padding == (1, 1, 1, 1, 2, 0): + self.conv = nn.Conv3d( + chan_in, + chan_out, + kernel_size, + stride=self.stride, + dilation=self.dilation, + padding=(0, 1, 1), + **kwargs, + ) + else: + self.conv = nn.Conv3d( + chan_in, + chan_out, + kernel_size, + stride=self.stride, + dilation=self.dilation, + **kwargs, + ) + self.is_first_run = True + + def forward(self, x, is_init=True, residual=None) -> torch.Tensor: + if self.is_first_run: + self.is_first_run = False + + if self.time_causal_padding == (1, 1, 1, 1, 2, 0): + pass + else: + x = nn.functional.pad(x, self.time_causal_padding).contiguous() + + x = base_conv3d_channel_last(x, self.conv, residual=residual) + return x + + +class AttnBlock(nn.Module): + + def __init__(self, in_channels) -> None: + super().__init__() + + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels) + self.q = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.k = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.v = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.proj_out = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + + def attention(self, x, is_init=True) -> torch.Tensor: + x = base_group_norm(x, self.norm, act_silu=False, channel_last=True) + q = self.q(x, is_init) + k = self.k(x, is_init) + v = self.v(x, is_init) + + b, t, h, w, c = q.shape + q, k, v = map(lambda x: rearrange(x, "b t h w c -> b 1 (t h w) c"), (q, k, v)) + x = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True) + x = rearrange(x, "b 1 (t h w) c -> b t h w c", t=t, h=h, w=w) + + return x + + def forward(self, x): + x = x.permute(0, 2, 3, 4, 1).contiguous() + h = self.attention(x) + x = self.proj_out(h, residual=x) + x = x.permute(0, 4, 1, 2, 3) + return x + + +class Resnet3DBlock(nn.Module): + + def __init__( + self, + in_channels, + out_channels=None, + temb_channels=512, + conv_shortcut=False, + ) -> None: + super().__init__() + + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels) + self.conv1 = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3) + if temb_channels > 0: + self.temb_proj = nn.Linear(temb_channels, out_channels) + + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels) + self.conv2 = CausalConvAfterNorm(out_channels, out_channels, kernel_size=3) + + assert conv_shortcut is False + self.use_conv_shortcut = conv_shortcut + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = CausalConvAfterNorm( + in_channels, out_channels, kernel_size=3 + ) + else: + self.nin_shortcut = CausalConvAfterNorm( + in_channels, out_channels, kernel_size=1 + ) + + def forward(self, x, temb=None, is_init=True) -> torch.Tensor: + x = x.permute(0, 2, 3, 4, 1).contiguous() + + h = base_group_norm_with_zero_pad(x, self.norm1, act_silu=True, pad_size=2) + h = self.conv1(h) + if temb is not None: + h = h + self.temb_proj(nn.functional.silu(temb))[:, :, None, None] + + x = self.nin_shortcut(x) if self.in_channels != self.out_channels else x + + h = base_group_norm_with_zero_pad(h, self.norm2, act_silu=True, pad_size=2) + x = self.conv2(h, residual=x) + + x = x.permute(0, 4, 1, 2, 3) + return x + + +class Downsample3D(nn.Module): + + def __init__(self, in_channels, with_conv, stride) -> None: + super().__init__() + + self.with_conv = with_conv + if with_conv: + self.conv = CausalConv( + in_channels, in_channels, kernel_size=3, stride=stride + ) + + def forward(self, x, is_init=True) -> torch.Tensor: + if self.with_conv: + x = self.conv(x, is_init) + else: + x = nn.functional.avg_pool3d(x, kernel_size=2, stride=2) + return x + + +class VideoEncoder(nn.Module): + + def __init__( + self, + ch=32, + ch_mult=(4, 8, 16, 16), + num_res_blocks=2, + in_channels=3, + z_channels=16, + double_z=True, + down_sampling_layer=(1, 2), + resamp_with_conv=True, + version=1, + ) -> None: + super().__init__() + + temb_ch = 0 + + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + + # downsampling + self.conv_in = CausalConv(in_channels, ch, kernel_size=3) + self.down_sampling_layer = down_sampling_layer + + in_ch_mult = (1,) + tuple(ch_mult) + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append( + Resnet3DBlock( + in_channels=block_in, + out_channels=block_out, + temb_channels=temb_ch, + ) + ) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + if i_level in self.down_sampling_layer: + down.downsample = Downsample3D( + block_in, resamp_with_conv, stride=(2, 2, 2) + ) + else: + down.downsample = Downsample2D( + block_in, resamp_with_conv, padding=0 + ) # DIFF + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = Resnet3DBlock( + in_channels=block_in, out_channels=block_in, temb_channels=temb_ch + ) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = Resnet3DBlock( + in_channels=block_in, out_channels=block_in, temb_channels=temb_ch + ) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) + self.version = version + if version == 2: + channels = 4 * z_channels * 2**3 + self.conv_patchify = ConvPixelUnshuffleDownSampleLayer3D( + block_in, channels, kernel_size=3, factor=2 + ) + self.shortcut_pathify = PixelUnshuffleChannelAveragingDownSampleLayer3D( + block_in, channels, 2 + ) + self.shortcut_out = PixelUnshuffleChannelAveragingDownSampleLayer3D( + channels, 2 * z_channels if double_z else z_channels, 1 + ) + self.conv_out = CausalConvChannelLast( + channels, 2 * z_channels if double_z else z_channels, kernel_size=3 + ) + else: + self.conv_out = CausalConvAfterNorm( + block_in, 2 * z_channels if double_z else z_channels, kernel_size=3 + ) + + @torch.inference_mode() + def forward(self, x, video_frame_num, is_init=True) -> torch.Tensor: + # timestep embedding + temb = None + + t = video_frame_num + + # downsampling + h = self.conv_in(x, is_init) + + # make it real channel last, but behave like normal layout + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](h, temb, is_init) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + + if i_level != self.num_resolutions - 1: + if isinstance(self.down[i_level].downsample, Downsample2D): + _, _, t, _, _ = h.shape + h = rearrange(h, "b c t h w -> (b t) h w c", t=t) + h = self.down[i_level].downsample(h) + h = rearrange(h, "(b t) h w c -> b c t h w", t=t) + else: + h = self.down[i_level].downsample(h, is_init) + + h = self.mid.block_1(h, temb, is_init) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb, is_init) + + h = h.permute(0, 2, 3, 4, 1).contiguous() # b c l h w -> b l h w c + if self.version == 2: + h = base_group_norm(h, self.norm_out, act_silu=True, channel_last=True) + h = h.permute(0, 4, 1, 2, 3).contiguous() + shortcut = self.shortcut_pathify(h, is_init) + h = self.conv_patchify(h, is_init) + h = h.add_(shortcut) + shortcut = self.shortcut_out(h, is_init).permute(0, 2, 3, 4, 1) + h = self.conv_out(h.permute(0, 2, 3, 4, 1).contiguous(), is_init) + h = h.add_(shortcut) + else: + h = base_group_norm_with_zero_pad( + h, self.norm_out, act_silu=True, pad_size=2 + ) + h = self.conv_out(h, is_init) + h = h.permute(0, 4, 1, 2, 3) # b l h w c -> b c l h w + + h = rearrange(h, "b c t h w -> b t c h w") + return h + + +class Res3DBlockUpsample(nn.Module): + + def __init__( + self, input_filters, num_filters, down_sampling_stride, down_sampling=False + ) -> None: + super().__init__() + + self.input_filters = input_filters + self.num_filters = num_filters + + self.act_ = nn.SiLU(inplace=True) + + self.conv1 = CausalConvChannelLast( + num_filters, num_filters, kernel_size=[3, 3, 3] + ) + self.norm1 = nn.GroupNorm(32, num_filters) + + self.conv2 = CausalConvChannelLast( + num_filters, num_filters, kernel_size=[3, 3, 3] + ) + self.norm2 = nn.GroupNorm(32, num_filters) + + self.down_sampling = down_sampling + if down_sampling: + self.down_sampling_stride = down_sampling_stride + else: + self.down_sampling_stride = [1, 1, 1] + + if num_filters != input_filters or down_sampling: + self.conv3 = CausalConvChannelLast( + input_filters, + num_filters, + kernel_size=[1, 1, 1], + stride=self.down_sampling_stride, + ) + self.norm3 = nn.GroupNorm(32, num_filters) + + def forward(self, x, is_init=False) -> torch.Tensor: + x = x.permute(0, 2, 3, 4, 1).contiguous() + + residual = x + + h = self.conv1(x, is_init) + h = base_group_norm(h, self.norm1, act_silu=True, channel_last=True) + + h = self.conv2(h, is_init) + h = base_group_norm(h, self.norm2, act_silu=False, channel_last=True) + + if self.down_sampling or self.num_filters != self.input_filters: + x = self.conv3(x, is_init) + x = base_group_norm(x, self.norm3, act_silu=False, channel_last=True) + + h.add_(x) + h = self.act_(h) + if residual is not None: + h.add_(residual) + + h = h.permute(0, 4, 1, 2, 3) + return h + + +class Upsample3D(nn.Module): + + def __init__(self, in_channels, scale_factor=2) -> None: + super().__init__() + + self.scale_factor = scale_factor + self.conv3d = Res3DBlockUpsample( + input_filters=in_channels, + num_filters=in_channels, + down_sampling_stride=(1, 1, 1), + down_sampling=False, + ) + + def forward(self, x, is_init=True, is_split=True) -> torch.Tensor: + b, c, t, h, w = x.shape + + # x = x.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3).to(memory_format=torch.channels_last_3d) + if is_split: + split_size = c // 8 + x_slices = torch.split(x, split_size, dim=1) + x = [ + nn.functional.interpolate(x, scale_factor=self.scale_factor) + for x in x_slices + ] + x = torch.cat(x, dim=1) + else: + x = nn.functional.interpolate(x, scale_factor=self.scale_factor) + + x = self.conv3d(x, is_init) + return x + + +class VideoDecoder(nn.Module): + + def __init__( + self, + ch=128, + z_channels=16, + out_channels=3, + ch_mult=(1, 2, 4, 4), + num_res_blocks=2, + temporal_up_layers=(2, 3), + temporal_downsample=4, + resamp_with_conv=True, + version=1, + ) -> None: + super().__init__() + + temb_ch = 0 + + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.temporal_downsample = temporal_downsample + + block_in = ch * ch_mult[self.num_resolutions - 1] + self.version = version + if version == 2: + channels = 4 * z_channels * 2**3 + self.conv_in = CausalConv(z_channels, channels, kernel_size=3) + self.shortcut_in = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D( + z_channels, channels, 1 + ) + self.conv_unpatchify = ConvPixelShuffleUpSampleLayer3D( + channels, block_in, kernel_size=3, factor=2 + ) + self.shortcut_unpathify = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D( + channels, block_in, 2 + ) + else: + self.conv_in = CausalConv(z_channels, block_in, kernel_size=3) + + # middle + self.mid = nn.Module() + self.mid.block_1 = Resnet3DBlock( + in_channels=block_in, out_channels=block_in, temb_channels=temb_ch + ) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = Resnet3DBlock( + in_channels=block_in, out_channels=block_in, temb_channels=temb_ch + ) + + # upsampling + self.up_id = len(temporal_up_layers) + self.video_frame_num = 1 + self.cur_video_frame_num = self.video_frame_num // 2**self.up_id + 1 + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + block.append( + Resnet3DBlock( + in_channels=block_in, + out_channels=block_out, + temb_channels=temb_ch, + ) + ) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + if i_level in temporal_up_layers: + up.upsample = Upsample3D(block_in) + self.cur_video_frame_num = self.cur_video_frame_num * 2 + else: + up.upsample = Upsample2D(block_in, resamp_with_conv) + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) + self.conv_out = CausalConvAfterNorm(block_in, out_channels, kernel_size=3) + + @torch.inference_mode() + def forward(self, z, is_init=True) -> torch.Tensor: + z = rearrange(z, "b t c h w -> b c t h w") + h = self.conv_in(z, is_init=is_init) + if self.version == 2: + shortcut = self.shortcut_in(z, is_init=is_init) + h = h.add_(shortcut) + shortcut = self.shortcut_unpathify(h, is_init=is_init) + h = self.conv_unpatchify(h, is_init=is_init) + h = h.add_(shortcut) + + temb = None + + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + h = self.mid.block_1(h, temb, is_init=is_init) + h = self.mid.attn_1(h) + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + h = self.mid.block_2(h, temb, is_init=is_init) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) + h = self.up[i_level].block[i_block](h, temb, is_init=is_init) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + if isinstance(self.up[i_level].upsample, Upsample2D): + B = h.size(0) + h = h.permute(0, 2, 3, 4, 1).flatten(0, 1) + h = self.up[i_level].upsample(h) + h = h.unflatten(0, (B, -1)).permute(0, 4, 1, 2, 3) + else: + h = self.up[i_level].upsample(h, is_init=is_init) + + # end + h = h.permute(0, 2, 3, 4, 1) # b c l h w -> b l h w c + h = base_group_norm_with_zero_pad(h, self.norm_out, act_silu=True, pad_size=2) + h = self.conv_out(h) + h = h.permute(0, 4, 1, 2, 3) + + if is_init: + h = h[:, :, (self.temporal_downsample - 1) :] + return h + + +def rms_norm(input, normalized_shape, eps=1e-6) -> torch.Tensor: + dtype = input.dtype + input = input.to(torch.float32) + variance = ( + input.pow(2) + .flatten(-len(normalized_shape)) + .mean(-1)[(...,) + (None,) * len(normalized_shape)] + ) + input = input * torch.rsqrt(variance + eps) + return input.to(dtype) + + +class DiagonalGaussianDistribution: + + def __init__( + self, + parameters, + deterministic=False, + rms_norm_mean=False, + only_return_mean=False, + ) -> None: + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=-3) # N,[X],C,H,W + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + self.deterministic = deterministic + if self.deterministic: + self.var = self.std = torch.zeros_like( + self.mean, device=self.parameters.device, dtype=self.parameters.dtype + ) + if rms_norm_mean: + self.mean = rms_norm(self.mean, self.mean.size()[1:]) + self.only_return_mean = only_return_mean + + def sample(self, generator=None) -> torch.Tensor: + # make sure sample is on the same device + # as the parameters and has same dtype + sample = torch.randn( + self.mean.shape, generator=generator, device=self.parameters.device + ) + sample = sample.to(dtype=self.parameters.dtype) + x = self.mean + self.std * sample + if self.only_return_mean: + return self.mean + else: + return x + + +class AutoencoderKLStepvideo(nn.Module, ParallelTiledVAE): + + def __init__( + self, + config: StepVideoVAEConfig, + ) -> None: + nn.Module.__init__(self) + ParallelTiledVAE.__init__(self, config) + + self.frame_len = config.frame_len + + if config.version == 2: + self.latent_len = 3 + base_group_norm.spatial = True # type: ignore[attr-defined] + else: + self.latent_len = 5 + base_group_norm.spatial = False # type: ignore[attr-defined] + + self.encoder = VideoEncoder( + in_channels=config.in_channels, + z_channels=config.z_channels, + num_res_blocks=config.num_res_blocks, + version=config.version, + ) + + self.decoder = VideoDecoder( + z_channels=config.z_channels, + out_channels=config.out_channels, + num_res_blocks=config.num_res_blocks, + version=config.version, + ) + + self.world_size = config.world_size + # self.is_init = True + + def load_state_dict(self, state_dict, strict=True): + remapped = {} + for key, value in state_dict.items(): + if key.startswith("decoder.conv_out."): + # move “decoder.conv_out.weight” → “decoder.conv_out.conv.weight” + suffix = key[len("decoder.conv_out.") :] + remapped[f"decoder.conv_out.conv.{suffix}"] = value + else: + remapped[key] = value + super().load_state_dict(remapped, strict=strict) + + def _encode(self, x, is_init_image=True) -> torch.Tensor: + # b, len, c, h, w = x.size() + b, c, len, h, w = x.size() + # x = rearrange(x, 'b l c h w -> b c l h w').contiguous() + z = self.encoder(x, len, True) # 下采样[1, 4, 8, 16, 16] + return z + + @torch.inference_mode() + def encode(self, x): + # b (nc cf) c h w -> (b nc) cf c h w -> encode -> (b nc) cf c h w -> b (nc cf) c h w + chunks = list(x.split(self.frame_len, dim=1)) + for i in range(len(chunks)): + chunks[i] = self._encode(chunks[i], True) + z = torch.cat(chunks, dim=1) + + posterior = DiagonalGaussianDistribution(z) + return posterior.sample() + + def _decode(self, z) -> torch.Tensor: + + chunks = list(z.split(self.latent_len, dim=2)) + for i in range(len(chunks)): + chunks[i] = chunks[i].permute(0, 2, 1, 3, 4) + chunks[i] = chunks[i].to(next(self.decoder.parameters()).dtype) + chunks[i] = self.decoder(chunks[i], is_init=True) + x = torch.cat(chunks, dim=2) + return x + + def decode(self, z) -> torch.Tensor: + num_frames = z.size(2) + dec = ParallelTiledVAE.decode(self, z).permute(0, 2, 1, 3, 4) + dec = self.mix(dec).permute(0, 2, 1, 3, 4) + num_sample_frames = num_frames // 3 * 17 + return dec[:, :, :num_sample_frames] + + def mix(self, x) -> torch.Tensor: + remain_scale = 0.6 + mix_scale = 1.0 - remain_scale + front = slice(self.frame_len - 1, x.size(1) - 1, self.frame_len) + back = slice(self.frame_len, x.size(1), self.frame_len) + x[:, back] = x[:, back] * remain_scale + x[:, front] * mix_scale + x[:, front] = x[:, front] * remain_scale + x[:, back] * mix_scale + return x + + 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 = AutoencoderKLStepvideo diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py new file mode 100644 index 00000000..1018d43b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py @@ -0,0 +1,1343 @@ +# 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.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig +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.platforms import current_platform + +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) + + +@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 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: + 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. + + 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: 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. + + 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 WanUpsample(nn.Upsample): + r""" + Perform upsampling while ensuring the output tensor has the same data type as the input. + + Args: + x (torch.Tensor): Input tensor to be upsampled. + + 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 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): + 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 + + +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): + # 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 + + +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): + 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 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): + # 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 + + +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): + x_copy = x.clone() + 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) + + +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 + ): + 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) + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv_in = WanCausalConv3d(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( + WanResidualDownBlock( + 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(WanResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + self.down_blocks.append(WanAttentionBlock(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(WanResample(out_dim, mode=mode)) + scale /= 2.0 + + # middle blocks + self.mid_block = WanMidBlock(out_dim, dropout, non_linearity, num_layers=1) + + # output blocks + self.norm_out = WanRMS_norm(out_dim, images=False) + self.conv_out = WanCausalConv3d(out_dim, z_dim, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, 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_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 + 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) + 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): + """ + 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 + """ + if self.avg_shortcut is not None: + x_copy = x.clone() + + 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 + + +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): + """ + 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: + x = resnet(x) + + if self.upsamplers is not None: + x = self.upsamplers[0](x) + return 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, + ): + 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) + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + + # init block + self.conv_in = WanCausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.mid_block = WanMidBlock(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:], 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 = WanResidualUpBlock( + 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 = WanUpBlock( + 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) + + # output blocks + self.norm_out = WanRMS_norm(out_dim, images=False) + self.conv_out = WanCausalConv3d(out_dim, out_channels, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x): + ## 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) + 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(nn.Module, 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 + + 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, + ) + 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, + ) + + 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): + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vision_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vision_utils.py new file mode 100644 index 00000000..c50eefc2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/models/vision_utils.py @@ -0,0 +1,304 @@ +# 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 + +from sglang.multimodal_gen.runtime.utils.logging_utils import suppress_other_loggers + +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". + + Returns: + `PIL.Image.Image`: + A PIL Image. + """ + if isinstance(image, str): + if image.startswith("http://") or image.startswith("https://"): + with suppress_other_loggers(not_suppress_on_main_rank=True): + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/__init__.py new file mode 100644 index 00000000..af2eb7d1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/__init__.py @@ -0,0 +1 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/flux.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/flux.py new file mode 100644 index 00000000..d738c8ba --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/flux.py @@ -0,0 +1,124 @@ +# 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 ( + ConditioningStage, + DecodingStage, + DenoisingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, +) +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): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage_primary", + 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"), + ], + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"), + prepare_extra_set_timesteps_kwargs=[prepare_mu], + ), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +EntryClass = FluxPipeline diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/flux_2.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/flux_2.py new file mode 100644 index 00000000..58ce257a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/flux_2.py @@ -0,0 +1,124 @@ +# 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.pipelines_core.stages import ( + ConditioningStage, + DecodingStage, + DenoisingStage, + ImageVAEEncodingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, +) +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): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", + stage=InputValidationStage( + vae_image_processor=VaeImageProcessor( + vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor + * 2 + ), + ), + ) + + self.add_stage( + stage_name="prompt_encoding_stage_primary", + stage=TextEncodingStage( + text_encoders=[ + self.get_module("text_encoder"), + ], + tokenizers=[ + self.get_module("tokenizer"), + ], + ), + ) + + self.add_stage( + stage_name="image_encoding_stage_primary", + stage=ImageVAEEncodingStage( + vae_image_processor=VaeImageProcessor( + vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor + * 2 + ), + vae=self.get_module("vae"), + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"), + prepare_extra_set_timesteps_kwargs=[compute_empirical_mu], + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +EntryClass = Flux2Pipeline diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py new file mode 100644 index 00000000..0be68fc4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py @@ -0,0 +1,95 @@ +# 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 ( + ConditioningStage, + DecodingStage, + DenoisingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, +) +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): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage_primary", + 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"), + ], + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +EntryClass = HunyuanVideoPipeline diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py new file mode 100644 index 00000000..fabaae54 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py @@ -0,0 +1,187 @@ +# 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 ( + DecodingStage, + DenoisingStage, + ImageEncodingStage, + ImageVAEEncodingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.conditioning import ( + ConditioningStage, +) +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) * (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 QwenImagePipeline(LoRAPipeline, ComposedPipelineBase): + pipeline_name = "QwenImagePipeline" + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "vae", + "transformer", + "scheduler", + ] + + def create_pipeline_stages(self, server_args: ServerArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage_primary", + stage=TextEncodingStage( + text_encoders=[ + self.get_module("text_encoder"), + ], + tokenizers=[ + self.get_module("tokenizer"), + ], + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"), + prepare_extra_set_timesteps_kwargs=[prepare_mu], + ), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +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): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage_primary", + stage=ImageEncodingStage( + image_processor=self.get_module("processor"), + text_encoder=self.get_module("text_encoder"), + ), + ) + + self.add_stage( + stage_name="image_encoding_stage_primary", + stage=ImageVAEEncodingStage( + vae=self.get_module("vae"), + ), + ) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"), + prepare_extra_set_timesteps_kwargs=[prepare_mu], + ), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +EntryClass = [QwenImagePipeline, QwenImageEditPipeline] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/stepvideo_pipeline.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/stepvideo_pipeline.py new file mode 100644 index 00000000..9d2e80c2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/stepvideo_pipeline.py @@ -0,0 +1,182 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +# type: ignore +# 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. +""" + +import os +from typing import Any + +import torch +from huggingface_hub import hf_hub_download + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.loader.component_loader import ( + PipelineComponentLoader, +) +from sglang.multimodal_gen.runtime.models.encoders.bert import ( + HunyuanClip, # type: ignore +) +from sglang.multimodal_gen.runtime.models.encoders.stepllm import STEP1TextEncoder +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 ( + DecodingStage, + DenoisingStage, + InputValidationStage, + LatentPreparationStage, + StepvideoPromptEncodingStage, + TimestepPreparationStage, +) +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 StepVideoPipeline(LoRAPipeline, ComposedPipelineBase): + pipeline_name = "StepVideoPipeline" + + _required_config_modules = ["transformer", "scheduler", "vae"] + + def create_pipeline_stages(self, server_args: ServerArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage", + stage=StepvideoPromptEncodingStage( + stepllm=self.get_module("text_encoder"), + clip=self.get_module("text_encoder_2"), + ), + ) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + def build_llm(self, model_dir, device) -> torch.nn.Module: + text_encoder = ( + STEP1TextEncoder(model_dir, max_length=320).to(torch.bfloat16).eval() + ) + return text_encoder + + def build_clip(self, model_dir, device) -> HunyuanClip: + clip = HunyuanClip(model_dir, max_length=77).eval() + return clip + + def initialize_pipeline(self, server_args: ServerArgs): + """ + Initialize the pipeline. + """ + target_device = get_local_torch_device() + llm_dir = os.path.join(self.model_path, "step_llm") + clip_dir = os.path.join(self.model_path, "hunyuan_clip") + text_enc = self.build_llm(llm_dir, target_device) + clip_enc = self.build_clip(clip_dir, target_device) + self.add_module("text_encoder", text_enc) + self.add_module("text_encoder_2", clip_enc) + lib_path = ( + os.path.join( + server_args.model_path, + "lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so", + ) + if os.path.isdir(server_args.model_path) # local checkout + else hf_hub_download( + repo_id=server_args.model_path, + filename="lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so", + ) + ) + torch.ops.load_library(lib_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. + """ + 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") + + # some sanity checks + assert ( + len(model_index) > 1 + ), "model_index.json must contain at least one pipeline module" + + required_modules = ["transformer", "scheduler", "vae"] + for module_name in required_modules: + if module_name not in model_index: + raise ValueError( + f"model_index.json must contain a {module_name} module" + ) + logger.info("Diffusers config passed sanity checks") + + # all the component models used by the pipeline + modules = {} + for module_name, ( + transformers_or_diffusers, + architecture, + ) in model_index.items(): + component_model_path = os.path.join(self.model_path, module_name) + module = PipelineComponentLoader.load_module( + module_name=module_name, + component_model_path=component_model_path, + transformers_or_diffusers=transformers_or_diffusers, + server_args=server_args, + ) + logger.info("Loaded module %s from %s", module_name, component_model_path) + + if module_name in modules: + logger.warning("Overwriting module %s", module_name) + modules[module_name] = module + + required_modules = self.required_config_modules + # Check if all required modules were loaded + for module_name in required_modules: + if module_name not in modules or modules[module_name] is None: + raise ValueError( + f"Required module {module_name} was not loaded properly" + ) + + return modules + + +EntryClass = StepVideoPipeline diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_causal_dmd_pipeline.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_causal_dmd_pipeline.py new file mode 100644 index 00000000..b103ee0a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_causal_dmd_pipeline.py @@ -0,0 +1,81 @@ +# 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 ( + ConditioningStage, + DecodingStage, + CausalDMDDenoisingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, +) +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: + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer", None), + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=CausalDMDDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +EntryClass = WanCausalDMDPipeline diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_dmd_pipeline.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_dmd_pipeline.py new file mode 100644 index 00000000..3c973834 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_dmd_pipeline.py @@ -0,0 +1,101 @@ +# 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 ( + ConditioningStage, + DecodingStage, + DmdDenoisingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, +) + +# 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: + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer", None), + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DmdDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +EntryClass = WanDMDPipeline diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py new file mode 100644 index 00000000..7a49eef9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py @@ -0,0 +1,113 @@ +# 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.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 ( + ImageEncodingStage, + ConditioningStage, + DecodingStage, + DmdDenoisingStage, + ImageVAEEncodingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, +) + +# isort: on +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) + +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): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + + self.add_stage( + stage_name="image_encoding_stage", + stage=ImageEncodingStage( + image_encoder=self.get_module("image_encoder"), + image_processor=self.get_module("image_processor"), + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + + self.add_stage( + stage_name="image_latent_preparation_stage", + stage=ImageVAEEncodingStage(vae=self.get_module("vae")), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DmdDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +EntryClass = WanImageToVideoDmdPipeline diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_pipeline.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_pipeline.py new file mode 100644 index 00000000..93a19687 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_pipeline.py @@ -0,0 +1,118 @@ +# 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.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 ( + ImageEncodingStage, + ConditioningStage, + DecodingStage, + DenoisingStage, + ImageVAEEncodingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, +) + +# isort: on +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, +) + +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): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + + if ( + self.get_module("image_encoder") is not None + and self.get_module("image_processor") is not None + ): + self.add_stage( + stage_name="image_encoding_stage", + stage=ImageEncodingStage( + image_encoder=self.get_module("image_encoder"), + image_processor=self.get_module("image_processor"), + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + + self.add_stage( + stage_name="image_latent_preparation_stage", + stage=ImageVAEEncodingStage(vae=self.get_module("vae")), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +EntryClass = WanImageToVideoPipeline diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py new file mode 100644 index 00000000..8f1cbfc2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py @@ -0,0 +1,101 @@ +# 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.pipelines_core.stages import ( + ConditioningStage, + DecodingStage, + DenoisingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, +) +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: + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer", None), + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2", None), + scheduler=self.get_module("scheduler"), + vae=self.get_module("vae"), + pipeline=self, + ), + ) + + self.add_stage( + stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"), pipeline=self), + ) + + +EntryClass = WanPipeline diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py new file mode 100644 index 00000000..f8fd441d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py @@ -0,0 +1,116 @@ +# 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.pipelines_core.stages import ( + ConditioningStage, + DecodingStage, + DenoisingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, +) +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): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + + self.add_stage( + stage_name="prompt_encoding_stage_primary", + stage=TextEncodingStage( + text_encoders=[ + self.get_module("text_encoder"), + ], + tokenizers=[ + self.get_module("tokenizer"), + ], + ), + ) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"), + prepare_extra_set_timesteps_kwargs=[prepare_mu], + ), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) + ) + + +EntryClass = ZImagePipeline diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py new file mode 100644 index 00000000..32243af3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py @@ -0,0 +1,62 @@ +# 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 + model_info = get_model_info(model_path) + + pipeline_cls = model_info.pipeline_cls + + # instantiate the pipelines + pipeline = pipeline_cls(model_path, server_args) + + logger.info("Pipelines instantiated") + + return cast(PipelineWithLoRA, pipeline) + + +__all__ = [ + "build_pipeline", + "ComposedPipelineBase", + "Req", + "LoRAPipeline", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py new file mode 100644 index 00000000..175c14fe --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -0,0 +1,365 @@ +# 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 argparse +import os +from abc import ABC, abstractmethod +from typing import Any, cast + +import torch +from tqdm import tqdm + +from sglang.multimodal_gen.configs.pipeline_configs import PipelineConfig +from sglang.multimodal_gen.runtime.loader.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 Req +from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage +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] = {} + post_init_called: bool = False + 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") + # temp disable for duplicate initialing tp + # maybe_init_distributed_environment_and_model_parallel( + # server_args.tp_size, server_args.sp_size + # ) + + # Load modules directly in initialization + logger.info("Loading pipeline modules...") + self.modules = self.load_modules(server_args, loaded_modules) + + 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" + if self.post_init_called: + return + self.post_init_called = True + + self.initialize_pipeline(self.server_args) + if self.server_args.enable_torch_compile: + self.modules["transformer"] = torch.compile(self.modules["transformer"]) + logger.info("Torch Compile enabled for DiT") + + logger.info("Creating pipeline stages...") + self.create_pipeline_stages(self.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, + ) -> "ComposedPipelineBase": + """ + Load a pipeline from a pretrained model. + loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, + If provided, loaded_modules will be used instead of loading from config/pretrained weights. + """ + kwargs["model_path"] = model_path + server_args = ServerArgs.from_kwargs(**kwargs) + + logger.info("server_args in from_pretrained: %s", server_args) + + 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: + if module_name not in self.modules: + return default_value + return self.modules[module_name] + + 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) + self.model_path = model_path + # server_args.downloaded_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 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 + ): + logger.info( + "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..." + ) + self.required_config_modules.append("transformer_2") + logger.info( + "MoE pipeline detected. 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) + + 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) + 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 = os.path.join(self.model_path, load_module_name) + module = PipelineComponentLoader.load_module( + module_name=load_module_name, + component_model_path=component_model_path, + transformers_or_diffusers=transformers_or_diffusers, + server_args=server_args, + ) + logger.info("Loaded module %s from %s", module_name, component_model_path) + + if module_name in components: + logger.warning("Overwriting module %s", module_name) + components[module_name] = module + + # Check if all required modules were loaded + for module_name in required_modules: + if module_name not in components or components[module_name] is None: + raise ValueError( + f"Required module key: {module_name} value: {components.get(module_name)} was not found in loaded modules {components.keys()}" + ) + + return components + + def add_stage(self, stage_name: str, stage: PipelineStage): + assert self.modules is not None, "No modules are registered" + self._stages.append(stage) + self._stage_name_mapping[stage_name] = stage + setattr(self, stage_name, stage) + + # TODO(will): don't hardcode no_grad + @torch.no_grad() + def forward( + self, + batch: Req, + server_args: ServerArgs, + ) -> Req: + """ + 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 not self.post_init_called: + self.post_init() + + 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 + logger.info( + "Running pipeline stages: %s", + list(self._stage_name_mapping.keys()), + main_process_only=True, + ) + return self.executor.execute(self.stages, batch, server_args) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py new file mode 100644 index 00000000..dabe531c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py @@ -0,0 +1,92 @@ +# 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, +) +from sglang.multimodal_gen.runtime.pipelines_core import Req +from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import ( + PipelineExecutor, + Timer, +) +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 + + +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, + ) -> Req: + rank = get_classifier_free_guidance_rank() + cfg_rank = get_classifier_free_guidance_rank() + cfg_group = get_cfg_group() + + # TODO: decide when to gather on main when CFG_PARALLEL -> MAIN_RANK_ONLY + for stage in stages: + with Timer(stage.__class__.__name__): + paradigm = stage.parallelism_type + + if paradigm == StageParallelismType.MAIN_RANK_ONLY: + if rank == 0: + batch = stage(batch, server_args) + # 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] + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py new file mode 100644 index 00000000..917af320 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py @@ -0,0 +1,59 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +""" +Base class for all pipeline executors. +""" + +from abc import ABC, abstractmethod +from typing import List + +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.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 + +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, timings=None, simple_log=True, logger=logger) + + +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 + + @abstractmethod + def execute( + self, + stages: List[PipelineStage], + batch: Req, + server_args: ServerArgs, + ) -> Req: + """ + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py new file mode 100644 index 00000000..5e429b3e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py @@ -0,0 +1,39 @@ +# 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, + Timer, + logger, +) +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.server_args import ServerArgs + + +class SyncExecutor(PipelineExecutor): + """ + A simple synchronous executor that runs stages sequentially. + """ + + def execute( + self, + stages: List[PipelineStage], + batch: Req, + server_args: ServerArgs, + ) -> Req: + """ + Execute the pipeline stages sequentially. + """ + logger.info("Running pipeline stages sequentially with SyncExecutor.") + + for stage in stages: + with Timer(stage.__class__.__name__): + batch = stage(batch, server_args) + + return batch diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py new file mode 100644 index 00000000..33f8ad6c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py @@ -0,0 +1,284 @@ +# 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 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.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. + """ + + # [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]] = defaultdict( + dict + ) # state dicts of loaded lora adapters + loaded_adapter_paths: dict[str, str] = {} # nickname -> lora_path + cur_adapter_name: str = "" + cur_adapter_path: str = "" + # [dit_layer_name] = wrapped_lora_layer + lora_layers: dict[str, BaseLayerWithLoRA] = {} + lora_layers_critic: dict[str, BaseLayerWithLoRA] = {} + server_args: ServerArgs + exclude_lora_layers: list[str] = [] + device: torch.device = get_local_torch_device() + lora_target_modules: list[str] | None = None + lora_path: str | None = None + lora_nickname: str = "default" + lora_rank: int | None = None + lora_alpha: int | None = None + lora_initialized: bool = False + is_lora_merged: bool = False + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + 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 # 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 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 + converted_count = 0 + for name, layer in self.modules["transformer"].named_modules(): + if not self.is_target_layer(name): + continue + + 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: + self.lora_layers[name] = lora_layer + replace_submodule(self.modules["transformer"], name, lora_layer) + converted_count += 1 + logger.info("Converted %d layers to LoRA layers", converted_count) + + if "fake_score_transformer" in self.modules: + for name, layer in self.modules["fake_score_transformer"].named_modules(): + if not self.is_target_layer(name): + continue + layer = wrap_with_lora_layer( + layer, + lora_rank=self.lora_rank, + lora_alpha=self.lora_alpha, + ) + if layer is not None: + self.lora_layers_critic[name] = layer + replace_submodule( + self.modules["fake_score_transformer"], name, layer + ) + converted_count += 1 + logger.info( + "Converted %d layers to LoRA layers in the critic model", + converted_count, + ) + + def is_lora_effective(self): + return self.is_lora_merged + + def is_lora_set(self): + return self.lora_initialized and self.cur_adapter_name is not None + + 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 + lora_local_path = maybe_download_lora(lora_path) + lora_state_dict = load_file(lora_local_path) + + 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 (in_dim, r) @ (r, out_dim), we only merge (r, out_dim * n) where n is the number of linear layers to fuse + # see param mapping in HunyuanVideoArchConfig + if merge_index is not None and "lora_B" in name: + to_merge_params[target_name][merge_index] = weight + if len(to_merge_params[target_name]) == num_params_to_merge: + # cat at output dim according to the merge_index order + sorted_tensors = [ + to_merge_params[target_name][i] + for i in range(num_params_to_merge) + ] + weight = torch.cat(sorted_tensors, dim=1) + del to_merge_params[target_name] + else: + continue + + print(f"{name} -> {target_name}") + 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.cur_adapter_path = lora_path + 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, lora_path: str | None = None + ): # type: ignore + """ + Load a LoRA adapter into the pipeline and merge it into the transformer. + Args: + lora_nickname: The "nick name" of the adapter when referenced in the pipeline. + lora_path: The path to the adapter, either a local path or a Hugging Face repo id. + """ + if self.is_lora_merged and self.cur_adapter_name != lora_nickname: + raise ValueError( + f"LoRA '{self.cur_adapter_name}' is currently merged. " + "Please call 'unmerge_lora_weights' before setting a new LoRA." + ) + + if lora_nickname not in self.lora_adapters and lora_path is None: + raise ValueError( + f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it." + ) + if not self.lora_initialized: + self.convert_to_lora_layers() + + adapter_updated = False + rank = dist.get_rank() + + should_load = False + if lora_path is not None: + if lora_nickname not in self.loaded_adapter_paths: + should_load = True + elif self.loaded_adapter_paths[lora_nickname] != lora_path: + should_load = True + + if should_load: + adapter_updated = True + self.load_lora_adapter(lora_path, lora_nickname, rank) + + if ( + not adapter_updated + and self.cur_adapter_name == lora_nickname + and self.is_lora_merged + ): + return + self.cur_adapter_name = lora_nickname + + # Merge the new adapter + adapted_count = 0 + for name, layer in self.lora_layers.items(): + lora_A_name = name + ".lora_A" + lora_B_name = name + ".lora_B" + if ( + lora_A_name in self.lora_adapters[lora_nickname] + and lora_B_name in self.lora_adapters[lora_nickname] + ): + layer.set_lora_weights( + self.lora_adapters[lora_nickname][lora_A_name], + self.lora_adapters[lora_nickname][lora_B_name], + lora_path=lora_path, + ) + adapted_count += 1 + else: + if rank == 0: + logger.warning( + "LoRA adapter %s does not contain the weights for layer '%s'. LoRA will not be applied to it.", + lora_path, + name, + ) + layer.disable_lora = True + self.is_lora_merged = True + logger.info( + "Rank %d: LoRA adapter %s applied to %d layers", + rank, + lora_path, + adapted_count, + ) + + def merge_lora_weights(self) -> None: + if self.is_lora_merged: + logger.warning("LoRA weights are already merged") + return + + for name, layer in self.lora_layers.items(): + layer.merge_lora_weights() + logger.info("LoRA weights merged") + self.is_lora_merged = True + + def unmerge_lora_weights(self) -> None: + if not self.is_lora_merged: + logger.warning("LoRA weights are not merged.") + return + + for name, layer in self.lora_layers.items(): + layer.unmerge_lora_weights() + self.is_lora_merged = False diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py new file mode 100644 index 00000000..781d31e5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -0,0 +1,241 @@ +# 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 os +import pprint +from dataclasses import asdict, dataclass, field +from typing import TYPE_CHECKING, Any, Optional + +import PIL.Image +import torch + +from sglang.multimodal_gen.configs.sample.sampling_params import DataType +from sglang.multimodal_gen.configs.sample.teacache import ( + TeaCacheParams, + WanTeaCacheParams, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs + +if TYPE_CHECKING: + + from sglang.multimodal_gen.runtime.utils.perf_logger import RequestTimings + + +@dataclass +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. + """ + + # TODO(will): double check that args are separate from server_args + # properly. Also maybe think about providing an abstraction for pipeline + # specific arguments. + data_type: DataType + + request_id: str | None = None + + generator: torch.Generator | list[torch.Generator] | None = None + + # Image inputs + image_path: str | 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 + pixel_values: torch.Tensor | PIL.Image.Image | None = None + preprocessed_image: torch.Tensor | None = None + + # Text inputs + prompt: str | list[str] | None = None + negative_prompt: str | list[str] | None = None + prompt_path: str | None = None + output_path: str = "outputs/" + # without extension + output_file_name: str | 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 + + # Batch info + num_outputs_per_prompt: int = 1 + seed: int | None = None + seeds: list[int] | None = None + + # Tracking if embeddings are already processed + is_prompt_processed: bool = False + + # Latent tensors + latents: torch.Tensor | None = None + # Flux-2 + latent_ids: torch.Tensor | None = None + + 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 + + # Latent dimensions + height_latents: list[int] | int | None = None + width_latents: list[int] | int | None = None + num_frames: list[int] | 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 + ) + + # Original dimensions (before VAE scaling) + height: list[int] | int | None = None + width: list[int] | int | None = None + fps: list[int] | int | None = None + height_not_provided: bool = False + width_not_provided: bool = False + + # Timesteps + timesteps: torch.Tensor | None = None + timestep: torch.Tensor | float | int | None = None + step_index: int | None = None + boundary_ratio: float | None = None + + # Scheduler parameters + num_inference_steps: int = 50 + guidance_scale: float = 1.0 + guidance_scale_2: float | None = None + guidance_rescale: float = 0.0 + 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) + + return_trajectory_latents: bool = False + return_trajectory_decoded: bool = False + trajectory_timesteps: list[torch.Tensor] | None = None + trajectory_latents: torch.Tensor | None = None + + # Extra parameters that might be needed by specific pipeline implementations + extra: dict[str, Any] = field(default_factory=dict) + + # Misc + save_output: bool = True + return_frames: bool = False + + # TeaCache parameters + enable_teacache: bool = False + teacache_params: TeaCacheParams | WanTeaCacheParams | None = None + + # 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 + timings: Optional["RequestTimings"] = None + + # profile + profile: bool = False + num_profiled_timesteps: int = 8 + + # debugging + debug: bool = False + # dummy for now + perf_dump_path: str | None = None + + # results + output: torch.Tensor | None = None + + @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, output_idx): + 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}" + + return ( + os.path.join(self.output_path, output_file_name) + if output_file_name + else None + ) + + def __post_init__(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 + + def adjust_size(self, server_args: ServerArgs): + if self.height is None or self.width is None: + self.width = 1280 + self.height = 720 + + def __str__(self): + return pprint.pformat(asdict(self), indent=2, width=120) + + +@dataclass +class OutputBatch: + """ + Final output (after pipeline completion) + """ + + output: torch.Tensor | 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 + + # logged timings info, directly from Req.timings + timings: Optional["RequestTimings"] = None diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py new file mode 100644 index 00000000..3c962b0b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py @@ -0,0 +1,59 @@ +# 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.conditioning import ( + ConditioningStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_dmd import ( + DmdDenoisingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.encoding import EncodingStage +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.stepvideo_encoding import ( + StepvideoPromptEncodingStage, +) +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", + "ConditioningStage", + "DenoisingStage", + "DmdDenoisingStage", + "CausalDMDDenoisingStage", + "EncodingStage", + "DecodingStage", + "ImageEncodingStage", + "ImageVAEEncodingStage", + "TextEncodingStage", + "StepvideoPromptEncodingStage", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py new file mode 100644 index 00000000..e79ea187 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py @@ -0,0 +1,233 @@ +# 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.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.""" + 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 + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + A VerificationResult containing the verification status. + + """ + # Default implementation - no verification + return VerificationResult() + + def maybe_free_model_hooks(self): + 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. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + 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("cuda" if torch.cuda.is_available() else "cpu") + + 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. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + The updated batch information after this stage's processing. + """ + stage_name = self.__class__.__name__ + # Check if verification is enabled (simple approach for prototype) + enable_verification = getattr(server_args, "enable_stage_verification", False) + + if enable_verification: + # 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, timings=batch.timings): + result = self.forward(batch, server_args) + + if enable_verification: + # 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. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + The updated batch information after this stage's processing. + """ + raise NotImplementedError + + def backward( + self, + batch: Req, + server_args: ServerArgs, + ) -> Req: + raise NotImplementedError diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py new file mode 100644 index 00000000..d06dc142 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py @@ -0,0 +1,508 @@ +# 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.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +try: + from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import ( + SlidingTileAttentionBackend, + ) + + st_attn_available = True +except ImportError: + st_attn_available = False + SlidingTileAttentionBackend = None # type: ignore + +try: + from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import ( + VideoSparseAttentionBackend, + ) + + vsa_available = True +except ImportError: + vsa_available = False + VideoSparseAttentionBackend = None # type: ignore + +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 st_attn_available and self.attn_backend == SlidingTileAttentionBackend: + 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="cuda", 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="cuda", 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 ( + vsa_available + and self.attn_backend == VideoSparseAttentionBackend + ): + 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.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="cuda", + 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 + ), + ).to(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="cuda", 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py new file mode 100644 index 00000000..f0cb4569 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py @@ -0,0 +1,107 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +""" +Conditioning stage for diffusion pipelines. +""" + +import torch + +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 ConditioningStage(PipelineStage): + """ + Stage for applying conditioning to the diffusion process. + + This stage handles the application of conditioning, such as classifier-free guidance, + to the diffusion process. + """ + + @torch.no_grad() + def forward( + self, + batch: Req, + server_args: ServerArgs, + ) -> Req: + """ + Apply conditioning to the diffusion process. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + The batch with applied conditioning. + """ + # TODO!! + if not batch.do_classifier_free_guidance: + return batch + else: + return batch + + logger.info("batch.negative_prompt_embeds: %s", batch.negative_prompt_embeds) + logger.info( + "do_classifier_free_guidance: %s", batch.do_classifier_free_guidance + ) + logger.info("cfg_scale: %s", batch.guidance_scale) + + # Ensure negative prompt embeddings are available + assert ( + batch.negative_prompt_embeds is not None + ), "Negative prompt embeddings are required for classifier-free guidance" + + # Concatenate primary embeddings and masks + batch.prompt_embeds = torch.cat( + [batch.negative_prompt_embeds, batch.prompt_embeds] + ) + if batch.attention_mask is not None: + batch.attention_mask = torch.cat( + [batch.negative_attention_mask, batch.attention_mask] + ) + + # Concatenate secondary embeddings and masks if present + if batch.prompt_embeds_2 is not None: + batch.prompt_embeds_2 = torch.cat( + [batch.negative_prompt_embeds_2, batch.prompt_embeds_2] + ) + if batch.attention_mask_2 is not None: + batch.attention_mask_2 = torch.cat( + [batch.negative_attention_mask_2, batch.attention_mask_2] + ) + + return batch + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + """Verify conditioning stage inputs.""" + result = VerificationResult() + result.add_check( + "do_classifier_free_guidance", + batch.do_classifier_free_guidance, + V.bool_value, + ) + result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) + 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), + ) + return result + + def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + """Verify conditioning stage outputs.""" + result = VerificationResult() + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + return result diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py new file mode 100644 index 00000000..2d9cf7f0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -0,0 +1,213 @@ +# 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_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.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__) + + +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) -> None: + self.vae: ParallelTiledVAE = vae + self.pipeline = weakref.ref(pipeline) if pipeline else None + + @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 + """ + self.vae = self.vae.to(get_local_torch_device()) + latents = latents.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 + + # scale and shift + latents = self.scale_and_shift(latents, server_args) + latents = server_args.pipeline_config.preprocess_decoding(latents) + + # Decode latents + with torch.autocast( + device_type="cuda", 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) + image = self.vae.decode(latents) + + # De-normalize image to [0, 1] range + image = (image / 2 + 0.5).clamp(0, 1) + return image + + @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. + + Args: + batch: The current batch containing: + - latents: Tensor to decode (batch, channels, frames, height_latents, width_latents) + - return_trajectory_decoded (optional): Flag to decode trajectory latents + - trajectory_latents (optional): Latents at different timesteps + - trajectory_timesteps (optional): Corresponding timesteps + server_args: Configuration containing: + - output_type: "latent" to skip decoding, otherwise decode to pixels + - vae_cpu_offload: Whether to offload VAE to CPU after decoding + - model_loaded: Track VAE loading state + - model_paths: Path to VAE model if loading needed + + Returns: + Modified batch with: + - output: Decoded frames (batch, channels, frames, height, width) as CPU float32 + - trajectory_decoded (if requested): List of decoded frames per timestep + """ + # load vae if not already loaded (used for memory constrained devices) + pipeline = self.pipeline() if self.pipeline else None + if not server_args.model_loaded["vae"]: + loader = VAELoader() + self.vae = loader.load(server_args.model_paths["vae"], server_args) + if pipeline: + pipeline.add_module("vae", self.vae) + server_args.model_loaded["vae"] = True + + if server_args.output_type == "latent": + frames = batch.latents + else: + frames = self.decode(batch.latents, server_args) + + # decode trajectory latents if needed + if batch.return_trajectory_decoded: + trajectory_decoded = [] + assert ( + batch.trajectory_latents is not None + ), "batch should have trajectory latents" + for idx in range(batch.trajectory_latents.shape[1]): + # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width] + cur_latent = batch.trajectory_latents[:, idx, :, :, :, :] + cur_timestep = batch.trajectory_timesteps[idx] + logger.info("decoding trajectory latent for timestep: %s", cur_timestep) + decoded_frames = self.decode(cur_latent, server_args) + trajectory_decoded.append(decoded_frames.cpu().float()) + else: + trajectory_decoded = None + + # Convert to CPU float32 for compatibility + frames = frames.cpu().float() + + # Update batch with decoded image + output_batch = OutputBatch( + output=frames, + trajectory_timesteps=batch.trajectory_timesteps, + trajectory_latents=batch.trajectory_latents, + trajectory_decoded=trajectory_decoded, + timings=batch.timings, + ) + + # Offload models if needed + self.maybe_free_model_hooks() + + if server_args.vae_cpu_offload: + self.vae.to("cpu") + + if torch.backends.mps.is_available(): + del self.vae + if pipeline is not None and "vae" in pipeline.modules: + del pipeline.modules["vae"] + server_args.model_loaded["vae"] = False + + return output_batch diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py new file mode 100644 index 00000000..aae4af50 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -0,0 +1,1387 @@ +# 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.profiler +from einops import rearrange +from tqdm.auto import tqdm + +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.distributed import ( + cfg_model_parallel_all_reduce, + get_local_torch_device, + get_sp_parallel_rank, + get_sp_world_size, + get_world_group, +) +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, + get_world_rank, +) +from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( + FlashAttentionBackend, +) +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_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.interface import AttentionBackendEnum +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, masks_like + +try: + from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import ( + SlidingTileAttentionBackend, + ) + + st_attn_available = True +except ImportError: + st_attn_available = False + +try: + from sglang.multimodal_gen.runtime.layers.attention.backends.vmoba import ( + VMOBAAttentionBackend, + ) + from sglang.multimodal_gen.utils import is_vmoba_available + + vmoba_attn_available = is_vmoba_available() +except ImportError: + vmoba_attn_available = False + +try: + from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import ( + VideoSparseAttentionBackend, + ) + + vsa_available = True +except ImportError: + vsa_available = False + +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 + if self.server_args.enable_torch_compile: + full_graph = False + self.transformer = torch.compile( + self.transformer, mode="max-autotune", fullgraph=full_graph + ) + self.transformer_2 = ( + torch.compile( + self.transformer_2, mode="max-autotune", fullgraph=full_graph + ) + if transformer_2 is not None + else None + ) + + self.scheduler = scheduler + self.vae = vae + self.pipeline = weakref.ref(pipeline) if pipeline else None + + self.attn_backend = get_attn_backend( + head_size=attn_head_size, + dtype=torch.float16, # TODO(will): hack + supported_attention_backends={ + AttentionBackendEnum.SLIDING_TILE_ATTN, + AttentionBackendEnum.VIDEO_SPARSE_ATTN, + AttentionBackendEnum.VMOBA_ATTN, + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.SAGE_ATTN_THREE, + }, # hack + ) + + # cfg + self.guidance = None + + # misc + self.profiler = None + + @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=torch.float32, + device=device, + ).to(target_dtype) + * 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 _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): + """ + Prepare all necessary invariant variables for the denoising loop. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + A dictionary containing all the prepared variables for the denoising loop. + """ + pipeline = self.pipeline() if self.pipeline else None + if not server_args.model_loaded["transformer"]: + loader = TransformerLoader() + self.transformer = loader.load( + server_args.model_paths["transformer"], server_args + ) + if self.server_args.enable_torch_compile: + self.transformer = torch.compile( + self.transformer, mode="max-autotune", fullgraph=True + ) + if pipeline: + pipeline.add_module("transformer", self.transformer) + server_args.model_loaded["transformer"] = True + + # 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 + + # Get timesteps and calculate warmup steps + timesteps = batch.timesteps + if timesteps is None: + raise ValueError("Timesteps must be provided") + num_inference_steps = batch.num_inference_steps + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + + # 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 st_attn_available and self.attn_backend == SlidingTileAttentionBackend: + 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 + + # (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 + + # 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( + self.transformer.forward, + { + # 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( + self.transformer.forward, + { + "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, + ), + ) + + if batch.do_classifier_free_guidance: + neg_cond_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "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, + ), + ) + 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, + ): + # 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 + ) + + 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 ( + st_attn_available + and self.attn_backend == SlidingTileAttentionBackend + and server_args.STA_mode == STA_Mode.STA_SEARCHING + ): + self.save_sta_search_results(batch) + + # deallocate transformer if on mps + pipeline = self.pipeline() if self.pipeline else None + if torch.backends.mps.is_available(): + 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(), + ) + + 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 + + # For I2I tasks like QwenImageEdit, the image_latent (input image) should be + # replicated on all SP ranks, not sharded, as it provides global context. + if ( + server_args.pipeline_config.task_type != ModelTaskType.I2I + and 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 start_profile(self, batch: Req): + if not batch.profile: + return + + logger.info("Starting Profiler...") + # Build activities dynamically to avoid CUDA hangs when CUDA is unavailable + activities = [torch.profiler.ProfilerActivity.CPU] + if torch.cuda.is_available(): + activities.append(torch.profiler.ProfilerActivity.CUDA) + + self.profiler = torch.profiler.profile( + activities=activities, + schedule=torch.profiler.schedule( + skip_first=0, + wait=0, + warmup=1, + active=batch.num_profiled_timesteps, + repeat=5, + ), + on_trace_ready=lambda _: torch.profiler.tensorboard_trace_handler( + f"./logs" + ), + record_shapes=True, + with_stack=True, + ) + self.profiler.start() + + def step_profile(self): + if self.profiler: + if torch.cuda.is_available(): + torch.cuda.synchronize() + self.profiler.step() + + def stop_profile(self, batch: Req): + try: + if self.profiler: + logger.info("Stopping Profiler...") + if torch.cuda.is_available(): + torch.cuda.synchronize() + self.profiler.stop() + request_id = batch.request_id if batch.request_id else "profile_trace" + log_dir = f"./logs" + os.makedirs(log_dir, exist_ok=True) + + rank = get_world_rank() + trace_path = os.path.abspath( + os.path.join(log_dir, f"{request_id}-rank{rank}.trace.json.gz") + ) + logger.info(f"Saving profiler traces to: {trace_path}") + self.profiler.export_chrome_trace(trace_path) + torch.distributed.barrier() + except Exception as e: + logger.error(f"{e}") + + 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 _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. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + The batch with denoised latents. + """ + # 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() + + self.start_profile(batch=batch) + + # to avoid device-sync caused by timestep comparison + timesteps_cpu = timesteps.cpu() + num_timesteps = timesteps_cpu.shape[0] + with torch.autocast( + device_type=("cuda" if torch.cuda.is_available() else "cpu"), + 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): + # Skip if interrupted + if hasattr(self, "interrupt") and self.interrupt: + continue + + with StageProfiler( + f"denoising_step_{i}", logger=logger, timings=batch.timings + ): + 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) + 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, + ) + + # 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() + + self.step_profile() + + self.stop_profile(batch) + + denoising_end_time = time.time() + + if num_timesteps > 0: + 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, + ) + 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. + + Returns: + The prepared kwargs. + """ + extra_step_kwargs = {} + for k, v in kwargs.items(): + accepts = k in set(inspect.signature(func).parameters.keys()) + if accepts: + extra_step_kwargs[k] = v + return extra_step_kwargs + + def progress_bar( + self, iterable: Iterable | None = None, total: int | None = None + ) -> tqdm: + """ + Create a progress bar for the denoising process. + + Args: + iterable: The iterable to iterate over. + total: The total number of items. + + Returns: + A tqdm progress bar. + """ + local_rank = get_world_group().local_rank + if local_rank == 0: + return tqdm(iterable=iterable, total=total) + else: + return tqdm(iterable=iterable, total=total, disable=True) + + 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 + ) -> Any | None: + """ + Build attention metadata for custom attention backends. + + Args: + i: The current timestep index. + batch: The current batch information. + server_args: The inference arguments. + + Returns: + The attention metadata, or None if not applicable. + """ + attn_metadata = None + self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls() + if self.attn_metadata_builder_cls: + self.attn_metadata_builder = self.attn_metadata_builder_cls() + if (st_attn_available and self.attn_backend == SlidingTileAttentionBackend) or ( + vsa_available and self.attn_backend == VideoSparseAttentionBackend + ): + 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.VSA_sparsity, + device=get_local_torch_device(), + ) + elif vmoba_attn_available and self.attn_backend == VMOBAAttentionBackend: + moba_params = server_args.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 == FlashAttentionBackend: + attn_metadata = self.attn_metadata_builder.build( + raw_latent_shape=batch.raw_latent_shape + ) + else: + return None + + assert attn_metadata is not None, "attn_metadata cannot be None" + + return attn_metadata + + def _predict_noise( + self, + current_model, + latent_model_input, + timestep, + prompt_embeds, + target_dtype, + guidance: torch.Tensor, + **kwargs, + ): + return current_model( + hidden_states=latent_model_input, + encoder_hidden_states=prompt_embeds, + timestep=timestep, + guidance=guidance, + **kwargs, + ) + + def _predict_noise_with_cfg( + self, + current_model: torch.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, + prompt_embeds=server_args.pipeline_config.get_pos_prompt_embeds( + batch + ), + 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, + prompt_embeds=server_args.pipeline_config.get_neg_prompt_embeds( + batch + ), + 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) + + # 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.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. + + Args: + batch: The current batch information. + server_args: The inference arguments. + """ + # TODO(kevin): STA mask search, currently only support Wan2.1 with 69x768x1280 + STA_mode = server_args.STA_mode + skip_time_steps = server_args.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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py new file mode 100644 index 00000000..1ee183b0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py @@ -0,0 +1,204 @@ +# 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.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) -> None: + super().__init__(transformer, scheduler) + 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"] + prompt_embeds = prepared_vars["prompt_embeds"] + + denoising_loop_start_time = time.time() + self.start_profile(batch=batch) + + 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, timings=batch.timings + ): + # 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="cuda", + 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 = self.transformer( + latent_model_input.permute(0, 2, 1, 3, 4), + prompt_embeds, + 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], + ).to(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() + + self.stop_profile(batch) + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py new file mode 100644 index 00000000..6fe00884 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py @@ -0,0 +1,105 @@ +# 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.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: + 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. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + 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="cuda", 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py new file mode 100644 index 00000000..5264c1b0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py @@ -0,0 +1,357 @@ +# 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 PIL +import torch +from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution + +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.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 move_to_device(self, device): + 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. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + The batch with encoded prompt embeddings. + """ + + if batch.condition_image is None: + return batch + cuda_device = get_local_torch_device() + self.move_to_device(cuda_device) + + image = batch.condition_image + + image_processor_kwargs = ( + server_args.pipeline_config.prepare_image_processor_kwargs(batch) + ) + + image_inputs = self.image_processor( + images=image, return_tensors="pt", **image_processor_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.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.negative_prompt) + neg_image_processor_kwargs = dict(text=[txt], padding=True) + else: + neg_image_processor_kwargs = {} + + neg_image_inputs = self.image_processor( + images=image, return_tensors="pt", **neg_image_processor_kwargs + ).to(get_local_torch_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, + ) + 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, + ) + batch.prompt_embeds.append( + self.encoding_qwen_image_edit(outputs, image_inputs) + ) + + batch.negative_prompt_embeds.append( + self.encoding_qwen_image_edit(neg_outputs, neg_image_inputs) + ) + + self.move_to_device("cpu") + + 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 forward( + self, + batch: Req, + server_args: ServerArgs, + ) -> Req: + """ + Encode pixel representations into latent space. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + The batch with encoded outputs. + """ + + if batch.condition_image is None: + return batch + + num_frames = batch.num_frames + + self.vae = self.vae.to(get_local_torch_device()) + + image = batch.condition_image + 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="cuda", 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) + encoder_output: DiagonalGaussianDistribution = self.vae.encode( + video_condition + ) + + 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( + encoder_output, 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 + + batch.image_latent = server_args.pipeline_config.postprocess_image_latent( + latent_condition, batch + ) + + self.maybe_free_model_hooks() + + self.vae.to("cpu") + + 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) + ) + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py new file mode 100644 index 00000000..8ad061d9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py @@ -0,0 +1,250 @@ +# 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.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.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 + + 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 + # Peiyuan: using GPU seed will cause A100 and H100 to generate different results... + # FIXME: the generator's in latent preparation stage seems to be different from seeds + batch.generator = [torch.Generator("cpu").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 server_args.pipeline_config.task_type == ModelTaskType.I2I: + # calculate new condition image size + calculated_size = ( + server_args.pipeline_config.calculate_condition_image_size( + batch.condition_image, + condition_image_width, + condition_image_height, + ) + ) + + # preprocess condition image if necessary + if calculated_size is not None: + calculated_width, calculated_height = calculated_size + condition_image, calculated_size = ( + server_args.pipeline_config.preprocess_condition_image( + batch.condition_image, + calculated_width, + calculated_height, + self.vae_image_processor, + ) + ) + batch.condition_image = condition_image + + # adjust output image size + 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: + # duplicate with vae_image_processor + # further processing for ti2v task + 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.info("resized img height: %s, img width: %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 + max_area = 720 * 1280 + 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] + ) + height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value + width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value + + batch.condition_image = batch.condition_image.resize((width, height)) + batch.height = height + batch.width = width + + def forward( + self, + batch: Req, + server_args: ServerArgs, + ) -> Req: + """ + Validate and prepare inputs. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + The validated batch information. + """ + + self._generate_seeds(batch, server_args) + + # Ensure prompt is properly formatted + if 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 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 + + self.preprocess_condition_image( + batch, server_args, condition_image_width, condition_image_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 + ) + 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( + "num_inference_steps", batch.num_inference_steps, V.positive_int + ) + 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) + # Validate height and width + if batch.height % 8 != 0 or batch.width % 8 != 0: + raise ValueError( + f"Height and width must be divisible by 8 but are {batch.height} and {batch.width}." + ) + result.add_check("seeds", batch.seeds, V.list_not_empty) + result.add_check("generator", batch.generator, V.generator_or_list_generators) + return result diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py new file mode 100644 index 00000000..acd901b3 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py @@ -0,0 +1,163 @@ +# 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. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + The batch with prepared latent variables. + """ + + latent_num_frames = None + # Adjust video length based on VAE version if needed + if hasattr(self, "adjust_video_length"): + 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. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + The batch with adjusted video length. + """ + + video_length = batch.num_frames + 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 + else: # stepvideo only + latent_num_frames = video_length // 17 * 3 + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/stepvideo_encoding.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/stepvideo_encoding.py new file mode 100644 index 00000000..eed0edd9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/stepvideo_encoding.py @@ -0,0 +1,99 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 + +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.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__) + + +# The dedicated stepvideo prompt encoding stage. +class StepvideoPromptEncodingStage(PipelineStage): + """ + Stage for encoding prompts using the remote caption API. + + This stage applies the magic string transformations and calls + the remote caption service asynchronously to get: + - primary prompt embeddings, + - an attention mask, + - and a clip embedding. + """ + + def __init__(self, stepllm, clip) -> None: + super().__init__() + # self.caption_client = caption_client # This should have a call_caption(prompts: List[str]) method. + self.stepllm = stepllm + self.clip = clip + + @torch.no_grad() + def forward(self, batch: Req, server_args) -> Req: + + prompts = [batch.prompt + server_args.pipeline_config.pos_magic] + bs = len(prompts) + prompts += [server_args.pipeline_config.neg_magic] * bs + with set_forward_context(current_timestep=0, attn_metadata=None): + y, y_mask = self.stepllm(prompts) + clip_emb, _ = self.clip(prompts) + len_clip = clip_emb.shape[1] + y_mask = torch.nn.functional.pad(y_mask, (len_clip, 0), value=1) + pos_clip, neg_clip = clip_emb[:bs], clip_emb[bs:] + + # split positive vs negative text + batch.prompt_embeds = y[:bs] # [bs, seq_len, dim] + batch.negative_prompt_embeds = y[bs : 2 * bs] # [bs, seq_len, dim] + batch.prompt_attention_mask = y_mask[:bs] # [bs, seq_len] + batch.negative_attention_mask = y_mask[bs : 2 * bs] # [bs, seq_len] + batch.clip_embedding_pos = pos_clip + batch.clip_embedding_neg = neg_clip + return batch + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + """Verify stepvideo encoding stage inputs.""" + result = VerificationResult() + result.add_check("prompt", batch.prompt, V.string_not_empty) + return result + + def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + """Verify stepvideo encoding stage outputs.""" + result = VerificationResult() + result.add_check( + "prompt_embeds", batch.prompt_embeds, [V.is_tensor, V.with_dims(3)] + ) + result.add_check( + "negative_prompt_embeds", + batch.negative_prompt_embeds, + [V.is_tensor, V.with_dims(3)], + ) + result.add_check( + "prompt_attention_mask", + batch.prompt_attention_mask, + [V.is_tensor, V.with_dims(2)], + ) + result.add_check( + "negative_attention_mask", + batch.negative_attention_mask, + [V.is_tensor, V.with_dims(2)], + ) + result.add_check( + "clip_embedding_pos", + batch.clip_embedding_pos, + [V.is_tensor, V.with_dims(2)], + ) + result.add_check( + "clip_embedding_neg", + batch.clip_embedding_neg, + [V.is_tensor, V.with_dims(2)], + ) + return result diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py new file mode 100644 index 00000000..6423e45d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py @@ -0,0 +1,336 @@ +# 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. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + Returns: + The batch with encoded prompt embeddings. + """ + 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 not None: + 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 not None: + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py new file mode 100644 index 00000000..ca0e65f6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py @@ -0,0 +1,148 @@ +# 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 + +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: + self.scheduler = scheduler + self.prepare_extra_set_timesteps_kwargs = prepare_extra_set_timesteps_kwargs + + @property + def parallelism_type(self) -> StageParallelismType: + return StageParallelismType.REPLICATED + + def forward( + self, + batch: Req, + server_args: ServerArgs, + ) -> Req: + """ + Prepare timesteps for the diffusion process. + + Args: + batch: The current batch information. + server_args: The inference arguments. + + 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) + + # 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 + + # 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 + self.log_debug(f"timesteps: {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.""" + result = VerificationResult() + result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.with_dims(1)]) + return result diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/validators.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/validators.py new file mode 100644 index 00000000..1ca9e992 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/pipelines_core/stages/validators.py @@ -0,0 +1,522 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/__init__.py new file mode 100644 index 00000000..c87fb3aa --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/__init__.py @@ -0,0 +1,172 @@ +# 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.info("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.info("MPS (Metal Performance Shaders) is available") + except Exception as e: + logger.info("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.info("ROCm platform is available") + finally: + amdsmi.amdsmi_shut_down() + except Exception as e: + logger.info("ROCm platform is unavailable: %s", e) + + return ( + "sglang.multimodal_gen.runtime.platforms.rocm.RocmPlatform" if is_rocm else None + ) + + +builtin_platform_plugins = { + "cuda": cuda_platform_plugin, + "rocm": rocm_platform_plugin, + "mps": mps_platform_plugin, + "cpu": cpu_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 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 = "" + +if TYPE_CHECKING: + 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`. + # 2. when users use out-of-tree platform plugins, they might run + # `import sgl_diffusion`, some sgl_diffusion internal code might access + # `current_platform` during the import, and we need to make sure + # `current_platform` is only resolved after the plugins are loaded + # (we have tests for this, if any developer violate this, they will + # see the test failures). + 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"] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/cpu.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/cpu.py new file mode 100644 index 00000000..5186d248 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/cpu.py @@ -0,0 +1,61 @@ +# 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 + +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 + def get_device_total_memory(cls, device_id: int = 0) -> int: + # This is a rough estimate for CPU memory + # In practice, you might want to use psutil or similar + return 0 + + @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_device_communicator_cls(cls) -> str: + return "sglang.multimodal_gen.runtime.distributed.device_communicators.cpu_communicator.CpuCommunicator" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/cuda.py new file mode 100644 index 00000000..90320dae --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -0,0 +1,438 @@ +# 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 TypeVar + +import torch +from typing_extensions import ParamSpec + +from sglang.multimodal_gen.runtime.platforms.interface import ( + AttentionBackendEnum, + DeviceCapability, + Platform, + PlatformEnum, +) +from sglang.multimodal_gen.runtime.utils.common import is_blackwell +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_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_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_attn_backend_cls_str( + cls, + selected_backend: AttentionBackendEnum | None, + head_size: int, + dtype: torch.dtype, + ) -> str: + # 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. Fall back to Flash Attention." + ) + elif selected_backend == AttentionBackendEnum.SAGE_ATTN_THREE: + try: + from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import ( # noqa: F401 + SageAttention3Backend, + ) + from sglang.multimodal_gen.runtime.layers.attention.backends.sageattn.api import ( # noqa: F401 + sageattn_blackwell, + ) + + 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. Fall back to Flash Attention." + ) + 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.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 in [ + AttentionBackendEnum.FA, + ]: + if is_blackwell(): + from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( + set_fa_ver, + ) + + set_fa_ver(4) + target_backend = AttentionBackendEnum.FA + elif selected_backend: + raise ValueError(f"Invalid attention backend for {cls.device_name}") + else: + if is_blackwell(): + from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( + set_fa_ver, + ) + + set_fa_ver(4) + target_backend = AttentionBackendEnum.FA + + 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 + 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/interface.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/interface.py new file mode 100644 index 00000000..660f8e1c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/interface.py @@ -0,0 +1,252 @@ +# 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 typing import TYPE_CHECKING, 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): + FA = enum.auto() + SLIDING_TILE_ATTN = enum.auto() + TORCH_SDPA = enum.auto() + SAGE_ATTN = enum.auto() + SAGE_ATTN_THREE = enum.auto() + VIDEO_SPARSE_ATTN = enum.auto() + VMOBA_ATTN = enum.auto() + AITER = enum.auto() + NO_ATTENTION = enum.auto() + + def __str__(self): + return self.name.lower() + + +class PlatformEnum(enum.Enum): + CUDA = enum.auto() + ROCM = enum.auto() + TPU = enum.auto() + CPU = enum.auto() + MPS = 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 + + # 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] = [] + + def is_cuda(self) -> bool: + return self._enum == PlatformEnum.CUDA + + def is_rocm(self) -> bool: + return self._enum == PlatformEnum.ROCM + + def is_tpu(self) -> bool: + return self._enum == PlatformEnum.TPU + + def is_cpu(self) -> bool: + return self._enum == PlatformEnum.CPU + + def is_out_of_tree(self) -> bool: + return self._enum == PlatformEnum.OOT + + def is_cuda_alike(self) -> bool: + """Stateless version of :func:`torch.cuda.is_available`.""" + return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM) + + def is_mps(self) -> bool: + return self._enum == PlatformEnum.MPS + + @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 + def get_device_total_memory(cls, device_id: int = 0) -> int: + """Get the total memory of a device in bytes.""" + raise NotImplementedError + + @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.cuda.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_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 + + 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 = "" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/mps.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/mps.py new file mode 100644 index 00000000..2312ec05 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/mps.py @@ -0,0 +1,88 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.platforms.interface import ( + DeviceCapability, + Platform, + PlatformEnum, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +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 + 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 + 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 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_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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/rocm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/rocm.py new file mode 100644 index 00000000..1e5c370d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/platforms/rocm.py @@ -0,0 +1,138 @@ +# 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`. +""" + +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_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 + 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_attn_backend_cls_str( + cls, + selected_backend: AttentionBackendEnum | None, + head_size: int, + dtype: torch.dtype, + ) -> str: + logger.info( + "Trying SGLANG_DIFFUSION_ATTENTION_BACKEND=%s", + envs.SGLANG_DIFFUSION_ATTENTION_BACKEND, + ) + + 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 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/scheduler_client.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/scheduler_client.py new file mode 100644 index 00000000..c0c0a00a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/scheduler_client.py @@ -0,0 +1,149 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +import asyncio +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__) + + +# Using a singleton pattern to hold the ZMQ context and the socket connected to the scheduler +class SchedulerClient: + """ + A gateway for Scheduler, forwarding the ForwardBatch from http endpoints (or somewhere else) to background scheduler, with TCP socket + """ + + _instance = None + + def __new__(cls, *args, **kwargs): + if not cls._instance: + cls._instance = super(SchedulerClient, cls).__new__(cls) + return cls._instance + + def __init__(self, *args, **kwargs): + # Ensure the initialization runs only once for the singleton instance + if getattr(self, "_init_done", False): + return + # Queue + worker to strictly serialize ZeroMQ REQ/REP interactions + self._request_queue = asyncio.Queue() + self._worker_task = None + self._closing = False + self._init_done = True + + def initialize(self, server_args: ServerArgs): + self.server_args = server_args + self.context = zmq.asyncio.Context() + # This is the REQ socket used to connect to the backend Scheduler + self.scheduler_socket = self.context.socket(zmq.REQ) + scheduler_endpoint = server_args.scheduler_endpoint() + self.scheduler_socket.connect(scheduler_endpoint) + logger.info( + f"Scheduler client connected to backend scheduler at {scheduler_endpoint}" + ) + # Worker will be lazily started on the first forward call to ensure a running loop exists + + async def forward(self, batch: Any) -> Any: + """Enqueue a request to the backend Scheduler and await the reply.""" + if self._closing: + raise RuntimeError( + "SchedulerClient is closing; cannot forward new requests" + ) + + await self._ensure_worker_started() + + loop = asyncio.get_running_loop() + future = loop.create_future() + await self._request_queue.put((batch, future)) + return await future + + async def _ensure_worker_started(self): + # Start the worker only once and only when an event loop is running + if self._worker_task is None or self._worker_task.done(): + self._worker_task = asyncio.create_task(self._worker_loop()) + + async def _worker_loop(self): + while True: + try: + item = await self._request_queue.get() + try: + batch, future = item + except Exception: + # Malformed queue item; skip + self._request_queue.task_done() + continue + + try: + await self.scheduler_socket.send_pyobj(batch) + response = await self.scheduler_socket.recv_pyobj() + if not future.done(): + future.set_result(response) + except Exception as e: + if not future.done(): + future.set_exception(e) + finally: + self._request_queue.task_done() + except asyncio.CancelledError: + # Drain remaining items with cancellation error to avoid hanging waiters + while True: + try: + batch, future = self._request_queue.get_nowait() + except asyncio.QueueEmpty: + break + try: + if not future.done(): + future.set_exception(asyncio.CancelledError()) + finally: + self._request_queue.task_done() + raise + + def close(self): + self._closing = True + # Cancel worker if running + if self._worker_task is not None: + self._worker_task.cancel() + try: + self.scheduler_socket.close() + finally: + try: + self.context.term() + except Exception: + pass + + +# Singleton instance +scheduler_client = SchedulerClient() + + +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 + request_batch = await socket.recv_pyobj() + 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 scheduler_client.forward(request_batch) + + # 3. Send the Scheduler's reply back to the offline client + await socket.send_pyobj(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 + await socket.send_pyobj({"status": "error", "message": str(e)}) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/server_args.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/server_args.py new file mode 100644 index 00000000..4abc5cbf --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/server_args.py @@ -0,0 +1,1053 @@ +# 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 inspect +import json +import random +import sys +import tempfile +from contextlib import contextmanager +from dataclasses import field +from enum import Enum +from typing import Any, Optional + +from sglang.multimodal_gen.configs.pipeline_configs import FluxPipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig, STA_Mode +from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( + QwenImageEditPipelineConfig, + QwenImagePipelineConfig, +) +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 ( + configure_logger, + init_logger, +) +from sglang.multimodal_gen.utils import FlexibleArgumentParser, StoreBoolean + +logger = init_logger(__name__) + +ZMQ_TCP_PORT_DELTA = 233 + + +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. + - Fallback to str(...) for unknown types. + """ + # Handle simple types quickly + if obj is None or isinstance(obj, (str, int, float, bool)): + return obj + + # Enum -> value for readability + if isinstance(obj, Enum): + return obj.value + + # torch.Tensor handling (lazy import) + is_tensor, torch_mod = _is_torch_tensor(obj) + if is_tensor: + try: + ten = obj.detach().cpu() + if key_hint == "scaling_factor": + # Provide a compact, single-line summary for scaling_factor + stats = { + "shape": list(ten.shape), + "dtype": str(ten.dtype), + } + # Stats might fail for some dtypes; guard individually + 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} + # Generic tensor summary + return {"tensor": True, "shape": list(ten.shape), "dtype": str(ten.dtype)} + except Exception: + return "" + + # Dataclasses -> dict + 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: # drop noisy mappings + continue + try: + value = getattr(obj, name) + except Exception: + continue + result[name] = _sanitize_for_logging(value, key_hint=name) + return result + + # Dicts -> sanitize keys/values; drop 'param_names_mapping' + 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 + + # Sequences/Sets -> list + if isinstance(obj, (list, tuple, set)): + return [_sanitize_for_logging(x) for x in obj] + + # Functions / Callables -> qualified name + 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 + + # Fallback: string representation + try: + return str(obj) + except Exception: + return "" + + +class ExecutionMode(str, Enum): + """ + Enumeration for different pipeline modes. + + Inherits from str to allow string comparison for backward compatibility. + """ + + INFERENCE = "inference" + + @classmethod + def from_string(cls, value: str) -> "ExecutionMode": + """Convert string to ExecutionMode enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid mode: {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 [mode.value for mode in cls] + + +class WorkloadType(str, Enum): + """ + Enumeration for different workload types. + + Inherits from str to allow string comparison for backward compatibility. + """ + + I2V = "i2v" # Image to Video + T2V = "t2v" # Text to Video + T2I = "t2i" # Text to Image + I2I = "i2i" # Image to Image + + @classmethod + def from_string(cls, value: str) -> "WorkloadType": + """Convert string to WorkloadType enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid workload type: {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 [workload.value for workload in cls] + + +# args for sgl_diffusion framework +@dataclasses.dataclass +class ServerArgs: + # Model and path configuration (for convenience) + model_path: str + + # Attention + attention_backend: str = None + + # Running mode + mode: ExecutionMode = ExecutionMode.INFERENCE + + # Workload type + workload_type: WorkloadType = WorkloadType.T2V + + # Cache strategy + cache_strategy: str = "none" + + # Distributed executor backend + distributed_executor_backend: str = "mp" + nccl_port: Optional[int] = None + + # HuggingFace specific parameters + trust_remote_code: bool = False + revision: str | None = None + + # Parallelism + num_gpus: int = 1 + tp_size: int = -1 + sp_degree: int = -1 + # 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: int = -1 + dist_timeout: int | None = None # timeout for torch.distributed + + pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False) + + # 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 + # 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 + + output_type: str = "pil" + + # CPU offload parameters + dit_cpu_offload: bool = True + use_fsdp_inference: bool = False + text_encoder_cpu_offload: bool = True + image_encoder_cpu_offload: bool = True + vae_cpu_offload: bool = True + pin_cpu_memory: bool = True + + # 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 + + # Compilation + enable_torch_compile: bool = False + + disable_autocast: bool = False + + # VSA parameters + VSA_sparsity: float = 0.0 # inference/validation sparsity + + # V-MoBA parameters + moba_config_path: str | None = None + moba_config: dict[str, Any] = field(default_factory=dict) + + # Master port for distributed inference + # TODO: do not hard code + master_port: int | None = None + + # http server endpoint config, would be ignored in local mode + host: str | None = None + port: int | None = None + + scheduler_port: int = 5555 + + # Stage verification + enable_stage_verification: bool = True + + # 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, + } + ) + override_transformer_cls_name: str | None = None + + # # 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 __post_init__(self): + # Add randomization to avoid race condition when multiple servers start simultaneously + if self.attention_backend in ["fa3", "fa4"]: + self.attention_backend = "fa" + + initial_scheduler_port = self.scheduler_port + random.randint(0, 100) + self.scheduler_port = self.settle_port(initial_scheduler_port) + # TODO: remove hard code + initial_master_port = (self.master_port or 30005) + random.randint(0, 100) + self.master_port = self.settle_port(initial_master_port, 37) + if self.moba_config_path: + try: + with open(self.moba_config_path) as f: + self.moba_config = json.load(f) + logger.info("Loaded V-MoBA config from %s", self.moba_config_path) + except (FileNotFoundError, json.JSONDecodeError) as e: + logger.error( + "Failed to load V-MoBA config from %s: %s", self.moba_config_path, e + ) + raise + self.check_server_args() + + configure_logger(server_args=self) + + # 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-dir", + type=str, + help="Directory containing StepVideo model", + ) + + # attention + parser.add_argument( + "--attention-backend", + type=str, + default=None, + choices=[e.name.lower() for e in AttentionBackendEnum] + ["fa3", "fa4"], + help="The attention backend to use. If not specified, the backend is automatically selected based on hardware and installed packages.", + ) + + # Running mode + parser.add_argument( + "--mode", + type=str, + choices=ExecutionMode.choices(), + default=ServerArgs.mode.value, + help="The mode to run SGLang-diffusion", + ) + + # Workload type + parser.add_argument( + "--workload-type", + type=str, + choices=WorkloadType.choices(), + default=ServerArgs.workload_type.value, + help="The workload type", + ) + + # distributed_executor_backend + parser.add_argument( + "--distributed-executor-backend", + type=str, + choices=["mp"], + default=ServerArgs.distributed_executor_backend, + help="The distributed executor backend to use", + ) + + # 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=ServerArgs.tp_size, + help="The tensor parallelism size.", + ) + parser.add_argument( + "--sp-degree", + type=int, + default=ServerArgs.sp_degree, + help="The sequence parallelism size.", + ) + 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=ServerArgs.hsdp_shard_dim, + help="The data parallelism shards.", + ) + parser.add_argument( + "--dist-timeout", + type=int, + default=ServerArgs.dist_timeout, + help="Set timeout for torch.distributed initialization.", + ) + + # Output type + parser.add_argument( + "--output-type", + type=str, + default=ServerArgs.output_type, + choices=["pil"], + help="Output type for the generated video", + ) + + # 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", + ) + + # STA (Sliding Tile Attention) parameters + parser.add_argument( + "--STA-mode", + type=str, + default=ServerArgs.STA_mode.value, + choices=[mode.value for mode in STA_Mode], + help="STA mode contains STA_inference, STA_searching, STA_tuning, STA_tuning_cfg, None", + ) + parser.add_argument( + "--skip-time-steps", + type=int, + default=ServerArgs.skip_time_steps, + help="Number of time steps to warmup (full attention) for STA", + ) + 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)", + ) + + 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( + "--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", + ) + + # VSA parameters + parser.add_argument( + "--VSA-sparsity", + type=float, + default=ServerArgs.VSA_sparsity, + help="Validation sparsity for VSA", + ) + + # 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.", + ) + + # Stage verification + parser.add_argument( + "--enable-stage-verification", + action=StoreBoolean, + default=ServerArgs.enable_stage_verification, + help="Enable input/output verification for pipeline stages", + ) + parser.add_argument( + "--override-transformer-cls-name", + type=str, + default=ServerArgs.override_transformer_cls_name, + help="Override transformer cls name", + ) + # 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", + ) + # 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.", + ) + return parser + + def url(self): + if is_valid_ipv6_address(self.host): + return f"http://[{self.host}]:{self.port}" + else: + return f"http://{self.host}:{self.port}" + + def scheduler_endpoint(self): + """ + Internal endpoint for scheduler + + """ + scheduler_host = self.host or "localhost" + 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. + + Args: + port: Initial port to check + port_inc: Port increment for each attempt + max_attempts: Maximum number of attempts to find an available port + + Returns: + An available port number + + Raises: + RuntimeError: If no available port is found after max_attempts + """ + 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})" + ) + + def post_init_serve(self): + """ + Post init when in serve mode + """ + if self.host is None: + self.host = "localhost" + if self.port is None: + self.port = 3000 + self.port = self.settle_port(self.port) + + @classmethod + def from_cli_args( + cls, args: argparse.Namespace, unknown_args: list[str] | None = None + ) -> "ServerArgs": + if unknown_args is None: + unknown_args = [] + 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 override config file args + provided_args = {**config_args, **provided_args} + + # Handle special cases + # if "tp_size" in provided_args: + # provided_args["tp"] = provided_args.pop("tp_size") + + return cls.from_dict(provided_args) + + @classmethod + def from_dict(cls, kwargs: dict[str, Any]) -> "ServerArgs": + """Create a ServerArgs object from a dictionary.""" + attrs = [attr.name for attr in dataclasses.fields(cls)] + server_args_kwargs: dict[str, Any] = {} + + 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 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 mode string to enum if necessary + if "mode" in kwargs and isinstance(kwargs["mode"], str): + kwargs["mode"] = ExecutionMode.from_string(kwargs["mode"]) + + # Convert workload_type string to enum if necessary + if "workload_type" in kwargs and isinstance(kwargs["workload_type"], str): + kwargs["workload_type"] = WorkloadType.from_string(kwargs["workload_type"]) + + 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 check_server_sp_args(self): + if self.sp_degree == -1: + # assume we leave all remaining gpus to sp + 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"{self.num_gpus=} % {num_gpus_per_group} != 0") + self.sp_degree = self.num_gpus // num_gpus_per_group + + 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.info( + f"Ulysses degree not set, " f"using default value {self.ulysses_degree}" + ) + + if self.ring_degree is None: + self.ring_degree = 1 + logger.info( + f"Ring degree not set, " f"using default value {self.ring_degree}" + ) + + if self.ring_degree > 1: + if self.attention_backend != None and self.attention_backend != "fa": + raise ValueError( + "Ring Attention is only supported for flash attention backend for now" + ) + else: + self.attention_backend = "fa" + logger.info( + "Ring Attention is currently only supported for flash attention, attention_backend has been automatically set to flash attention" + ) + + if self.sp_degree == -1: + self.sp_degree = self.ring_degree * self.ulysses_degree + logger.info( + f"sequence_parallel_degree is not provided, using ring_degree * ulysses_degree = {self.sp_degree}" + ) + + if self.sp_degree != self.ring_degree * self.ulysses_degree: + raise ValueError( + f"sequence_parallel_degree is not equal to ring_degree * ulysses_degree, {self.sp_degree} != {self.ring_degree} * {self.ulysses_degree}" + ) + + def check_server_dp_args(self): + assert self.num_gpus % self.dp_size == 0, f"{self.num_gpus=}, {self.dp_size=}" + assert self.dp_size >= 1, "--dp-size must be natural number" + # NOTE: disable temporarily + # self.dp_degree = self.num_gpus // self.dp_size + logger.info(f"Setting dp_degree to: {self.dp_degree}") + if self.dp_size > 1: + raise ValueError("DP is not yet supported") + + def check_server_args(self) -> None: + """Validate inference arguments for consistency""" + if current_platform.is_mps(): + self.use_fsdp_inference = False + + # autocast + is_flux = ( + isinstance(self.pipeline_config, FluxPipelineConfig) + or isinstance(self.pipeline_config, QwenImagePipelineConfig) + or isinstance(self.pipeline_config, QwenImageEditPipelineConfig) + ) + if is_flux: + self.disable_autocast = True + + # Validate mode consistency + assert isinstance( + self.mode, ExecutionMode + ), f"Mode must be an ExecutionMode enum, got {type(self.mode)}" + assert ( + self.mode in ExecutionMode.choices() + ), f"Invalid execution mode: {self.mode}" + + # Validate workload type + assert isinstance( + self.workload_type, WorkloadType + ), f"Workload type must be a WorkloadType enum, got {type(self.workload_type)}" + assert ( + self.workload_type in WorkloadType.choices() + ), f"Invalid workload type: {self.workload_type}" + + if self.tp_size == -1: + self.tp_size = 1 + + if self.hsdp_shard_dim == -1: + self.hsdp_shard_dim = self.num_gpus + + assert ( + self.sp_degree <= self.num_gpus and self.num_gpus % self.sp_degree == 0 + ), "num_gpus must >= and be divisible by sp_size" + assert ( + self.hsdp_replicate_dim <= self.num_gpus + and self.num_gpus % self.hsdp_replicate_dim == 0 + ), "num_gpus must >= and be divisible by hsdp_replicate_dim" + assert ( + self.hsdp_shard_dim <= self.num_gpus + and self.num_gpus % self.hsdp_shard_dim == 0 + ), "num_gpus must >= and be divisible by hsdp_shard_dim" + + if self.num_gpus < max(self.tp_size, self.sp_degree): + self.num_gpus = max(self.tp_size, self.sp_degree) + + if self.pipeline_config is None: + raise ValueError("pipeline_config is not set in ServerArgs") + + self.pipeline_config.check_pipeline_config() + + # parallelism + self.check_server_dp_args() + # allocate all remaining gpus for sp-size + self.check_server_sp_args() + + if self.enable_cfg_parallel: + if self.num_gpus == 1: + raise ValueError( + "CFG Parallelism is enabled via `--enable-cfg-parallel`, while -num-gpus==1" + ) + + +@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, + ) + + +# TODO: not sure what _current_server_args is for, using a _global_server_args instead +_current_server_args = None +_global_server_args = None + + +def prepare_server_args(argv: list[str]) -> ServerArgs: + """ + Prepare the inference arguments from the command line arguments. + + Args: + argv: The command line arguments. Typically, it should be `sys.argv[1:]` + to ensure compatibility with `parse_args` when no arguments are passed. + + Returns: + The inference arguments. + """ + parser = FlexibleArgumentParser() + ServerArgs.add_cli_args(parser) + raw_args = parser.parse_args(argv) + server_args = ServerArgs.from_cli_args(raw_args) + global _current_server_args + _current_server_args = server_args + return server_args + + +@contextmanager +def set_current_server_args(server_args: ServerArgs): + """ + Temporarily set the current sgl_diffusion config. + Used during model initialization. + We save the current sgl_diffusion config in a global variable, + so that all modules can access it, e.g. custom ops + can access the sgl_diffusion config to determine how to dispatch. + """ + global _current_server_args + old_server_args = _current_server_args + try: + _current_server_args = server_args + yield + finally: + _current_server_args = old_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_current_server_args() -> ServerArgs: + if _current_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("Current sgl_diffusion args is not set.") + return _current_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 + + +def parse_int_list(value: str) -> list[int]: + if not value: + return [] + return [int(x.strip()) for x in value.split(",")] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py new file mode 100644 index 00000000..3d32576f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py @@ -0,0 +1,92 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +from typing import Any + +import zmq + +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 SyncSchedulerClient: + """ + A synchronous, singleton client for communicating with the Scheduler service. + Designed for use in synchronous environments like the DiffGenerator or standalone scripts. + """ + + _instance = None + + def __new__(cls, *args, **kwargs): + if not cls._instance: + cls._instance = super(SyncSchedulerClient, cls).__new__(cls) + return cls._instance + + def initialize(self, server_args: ServerArgs): + if hasattr(self, "context") and not self.context.closed: + logger.warning( + "SyncSchedulerClient is already initialized. Re-initializing." + ) + self.close() + + self.server_args = server_args + self.context = zmq.Context() # Standard synchronous context + self.scheduler_socket = self.context.socket(zmq.REQ) + + # Set socket options for the main communication socket + self.scheduler_socket.setsockopt(zmq.LINGER, 0) + self.scheduler_socket.setsockopt( + zmq.RCVTIMEO, 6000000 + ) # 10 minute timeout for generation + + scheduler_endpoint = self.server_args.scheduler_endpoint() + self.scheduler_socket.connect(scheduler_endpoint) + logger.debug( + f"SyncSchedulerClient 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. + This avoids interfering with the state of the main REQ/REP socket. + """ + if not hasattr(self, "context") 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 hasattr(self, "scheduler_socket"): + self.scheduler_socket.close() + if hasattr(self, "context"): + self.context.term() + + +# Singleton instance for easy access +sync_scheduler_client = SyncSchedulerClient() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/common.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/common.py new file mode 100644 index 00000000..6907756e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/common.py @@ -0,0 +1,340 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +import importlib +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 == zmq.DEALER: + set_send_opt() + set_recv_opt() + elif socket_type == zmq.REQ: + set_send_opt() + set_recv_opt() + elif socket_type == zmq.REP: + 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_hip() -> bool: + return torch.version.hip is not None + + +@lru_cache(maxsize=1) +def is_cuda(): + return torch.cuda.is_available() and torch.version.cuda + + +@lru_cache(maxsize=1) +def is_cuda_alike(): + return is_cuda() or is_hip() + + +@lru_cache(maxsize=1) +def is_blackwell(): + if not is_cuda(): + return False + return torch.cuda.get_device_capability()[0] == 10 + + +@lru_cache(maxsize=1) +def is_hpu() -> bool: + return hasattr(torch, "hpu") and torch.hpu.is_available() + + +@lru_cache(maxsize=1) +def is_xpu() -> bool: + return hasattr(torch, "xpu") and torch.xpu.is_available() + + +@lru_cache(maxsize=1) +def is_npu() -> bool: + return hasattr(torch, "npu") and torch.npu.is_available() + + +@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() + ) + + +@lru_cache(maxsize=1) +def is_cpu() -> bool: + return os.getenv("SGLANG_USE_CPU_ENGINE", "0") == "1" and is_host_cpu_x86() + + +# 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 ''}" + + +def is_flashinfer_available(): + """ + Check whether flashinfer is available. + As of Oct. 6, 2024, it is only available on NVIDIA GPUs. + """ + # if not get_bool_env_var("SGLANG_IS_FLASHINFER_AVAILABLE", default="true"): + # return False + return importlib.util.find_spec("flashinfer") is not None and is_cuda() + + +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/distributed.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/distributed.py new file mode 100644 index 00000000..c89a31dc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/distributed.py @@ -0,0 +1,231 @@ +# 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 + + +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( + "cuda" if torch.cuda.is_available() and 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) + ).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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py new file mode 100644 index 00000000..a3c83b6f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -0,0 +1,406 @@ +# 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 hashlib +import json +import os +import tempfile +from functools import reduce +from pathlib import Path +from typing import Any, Optional, cast + +import filelock +from diffusers.loaders.lora_base import ( + _best_guess_weight_name, # watch out for potetential removal from diffusers +) +from huggingface_hub import snapshot_download +from transformers import AutoConfig, PretrainedConfig +from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES + +from sglang.multimodal_gen.runtime.utils.logging_utils import ( + init_logger, + suppress_other_loggers, +) + +logger = init_logger(__name__) +_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: + is_gguf = check_gguf_file(component_model_path) + if is_gguf: + 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) + + # Special architecture mapping check for GGUF models + if is_gguf: + if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: + raise RuntimeError(f"Can't get gguf config for {config.model_type}.") + model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type] + config.update({"architectures": [model_type]}) + + return config + + +def get_config( + model: str, + trust_remote_code: bool, + revision: Optional[str] = None, + model_override_args: Optional[dict] = None, + **kwargs, +): + try: + config = AutoConfig.from_pretrained( + model, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) + except ValueError as e: + raise e + + return config + + +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( + model_path: str, +) -> dict[str, Any]: + """Gets a configuration of a submodule for the given diffusers model. + + Args: + model_path: the path of the submodule + + Returns: + The loaded configuration. + """ + + # Check if the model path exists + if os.path.exists(model_path): + # tokenizer + config_names = ["generation_config.json"] + # By default, we load config.json, but scheduler_config.json for scheduler + if "scheduler" in model_path: + config_names.append("scheduler_config.json") + else: + config_names.append("config.json") + + config_file_paths = [ + os.path.join(model_path, config_name) for config_name in config_names + ] + + combined_config = reduce( + lambda acc, path: acc | load_dict(path), config_file_paths, {} + ) + + return combined_config + else: + raise RuntimeError(f"Diffusers config file not found at {model_path}") + + +# 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 = set( + [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 get_lock(model_name_or_path: str): + lock_dir = tempfile.gettempdir() + 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 + + +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 + """ + + local_path = maybe_download_model(model_name_or_path, local_dir, download) + # return directly if local_path is a file + if os.path.isfile(local_path): + return local_path + + 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." + ) + + # Check for transformer and vae directories + 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." + ) + + # 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"]) + 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 import hf_hub_download + 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.info( + "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 +) -> 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 + """ + + # If the path exists locally, return it + if os.path.exists(model_name_or_path): + logger.info("Model already exists locally") + return model_name_or_path + + # Otherwise, assume it's a HF Hub model ID and try to download it + try: + logger.info( + "Downloading model snapshot from HF Hub for %s...", model_name_or_path + ) + with get_lock(model_name_or_path).acquire( + poll_interval=2 + ), suppress_other_loggers(not_suppress_on_main_rank=True): + local_path = snapshot_download( + repo_id=model_name_or_path, + ignore_patterns=["*.onnx", "*.msgpack"], + local_dir=local_dir, + ) + logger.info("Downloaded model to %s", local_path) + return str(local_path) + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/logging_utils.py new file mode 100644 index 00000000..86d2e7b8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -0,0 +1,482 @@ +# 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 datetime +import logging +import os +import sys +import time +import warnings +from contextlib import contextmanager +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_CONFIGURE_LOGGING = envs.SGLANG_DIFFUSION_CONFIGURE_LOGGING +SGLANG_DIFFUSION_LOGGING_CONFIG_PATH = envs.SGLANG_DIFFUSION_LOGGING_CONFIG_PATH +SGLANG_DIFFUSION_LOGGING_LEVEL = envs.SGLANG_DIFFUSION_LOGGING_LEVEL +SGLANG_DIFFUSION_LOGGING_PREFIX = envs.SGLANG_DIFFUSION_LOGGING_PREFIX + +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 NewLineFormatter(logging.Formatter): + """Adds logging prefix to newlines to align multi-line messages.""" + + def __init__(self, fmt, datefmt=None, style="%"): + logging.Formatter.__init__(self, fmt, datefmt, style) + + def format(self, record): + msg = logging.Formatter.format(self, record) + if record.message != "": + parts = msg.split(record.message) + msg = msg.replace("\n", "\r\n" + parts[0]) + return msg + + +class ColoredFormatter(NewLineFormatter): + """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 level name.""" + original_levelname = record.levelname + color = self.LEVEL_COLORS.get(record.levelno) + if color: + record.levelname = f"{color}{original_levelname}{RESET}" + + formatted_message = super().format(record) + + if color: + record.levelname = original_levelname + + 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( + 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) + ) + + if should_log: + # stacklevel=3 to show the original caller's location, + # as this function is called by the patched methods. + 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) + + # 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( + 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, True), logger), + ) + + return cast(_SGLDiffusionLogger, logger) + + +logger = init_logger(__name__) + + +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" + logging.basicConfig( + level=getattr(logging, server_args.log_level.upper()), + format=log_format, + datefmt=datefmt, + force=True, + ) + + set_uvicorn_logging_configs() + + +def suppress_loggers(loggers_to_suppress: list[str]): + original_levels = {} + + for logger_name in loggers_to_suppress: + logger = logging.getLogger(logger_name) + original_levels[logger_name] = logger.level + logger.setLevel(logging.WARNING) + + return original_levels + + +@contextmanager +def suppress_other_loggers(not_suppress_on_main_rank: bool = False): + """ + A context manager to temporarily suppress specified loggers. + + Args: + not_suppress_on_main_rank (bool): If True, loggers will not be + suppressed on the main process (rank 0). + """ + # This is a global setting that we want to apply to all ranks + warnings.filterwarnings( + "ignore", category=UserWarning, message="The given NumPy array is not writable" + ) + + should_suppress = True + if not_suppress_on_main_rank: + if get_is_main_process() == 0: + should_suppress = False + + loggers_to_suppress = ["urllib3", "imageio", "imageio_ffmpeg", "PIL", "PIL_Image"] + original_levels = suppress_loggers(loggers_to_suppress) + + try: + yield + finally: + if should_suppress: + for logger_name, level in original_levels.items(): + logging.getLogger(logger_name).setLevel(level) + + +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, + prompt[:100], + ) + else: + max_len = 100 + suffix = "..." if len(prompt) > max_len else "" + logger.info(f"Processing prompt: {prompt[:100]}{suffix}") + + 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("Pixel data generated successfully in %.2f 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( + "Completed batch processing. Generated %d outputs in %.2f seconds.", + num_outputs, + total_time, + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/perf_logger.py new file mode 100644 index 00000000..16f94255 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/runtime/utils/perf_logger.py @@ -0,0 +1,263 @@ +# 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 + +from dateutil.tz import UTC + +import sglang +import sglang.multimodal_gen.envs as envs +from sglang.multimodal_gen.runtime.utils.logging_utils import ( + _SGLDiffusionLogger, + get_is_main_process, +) + + +@dataclasses.dataclass +class RequestTimings: + """A lightweight data class to store performance timings for a single request.""" + + 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 + + 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 to_dict(self) -> Dict[str, Any]: + """Serializes the timing data to a dictionary.""" + return { + "request_id": self.request_id, + "stages": self.stages, + "steps": self.steps, + "total_duration_ms": self.total_duration_ms, + } + + +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" + + +@dataclasses.dataclass +class RequestPerfRecord: + request_id: str + + timestamp: str + commit_hash: str + tag: str + + stages: list[dict] + steps: list[float] + total_duration_ms: float + + def __init__( + self, + request_id, + commit_hash, + tag, + stages, + steps, + total_duration_ms, + 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 + + +class StageProfiler: + """ + A unified context manager, records timing information (usually of a single Stage or a step) into a provided RequestTimings object (usually from a Req). + """ + + def __init__( + self, + stage_name: str, + logger: _SGLDiffusionLogger, + timings: Optional["RequestTimings"], + simple_log: bool = False, + ): + self.stage_name = stage_name + self.timings = timings + self.logger = logger + self.simple_log = simple_log + self.start_time = 0.0 + + # Check env var at runtime to ensure we pick up changes (e.g. from CLI args) + self.metrics_enabled = envs.SGLANG_DIFFUSION_STAGE_LOGGING + + def __enter__(self): + if self.simple_log: + self.logger.info(f"[{self.stage_name}] started...") + + if (self.metrics_enabled and self.timings) or self.simple_log: + self.start_time = time.perf_counter() + + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if not ((self.metrics_enabled and self.timings) or self.simple_log): + return False + + 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.simple_log: + self.logger.info( + f"[{self.stage_name}] finished in {execution_time_s:.4f} seconds" + ) + + if self.metrics_enabled and self.timings: + if "denoising_step_" in self.stage_name: + index = int(self.stage_name[len("denoising_step_") :]) + self.timings.record_steps(index, execution_time_s) + else: + self.timings.record_stage(self.stage_name, execution_time_s) + + 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 ""RequestTimings"" stores the performance metrics of a single request + """ + + @classmethod + def dump_benchmark_report( + cls, + file_path: str, + timings: "RequestTimings", + 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 timings.stages.items() + ] + + report = { + "timestamp": datetime.now(UTC).isoformat(), + "request_id": timings.request_id, + "commit_hash": get_git_commit_hash(), + "tag": tag, + "total_duration_ms": timings.total_duration_ms, + "steps": formatted_steps, + "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) + print(f"[Performance] Metrics dumped to: {abs_path}") + except IOError as e: + print(f"[Performance] Failed to dump metrics to {abs_path}: {e}") + logging.getLogger(__name__).error(f"Dump failed: {e}") + + @classmethod + def log_request_summary( + cls, + timings: "RequestTimings", + tag: str = "total_inference_time", + ): + """logs the stage metrics and total duration for a completed request + to the performance_log file. + """ + formatted_stages = [ + {"name": name, "execution_time_ms": duration_ms} + for name, duration_ms in timings.stages.items() + ] + + record = RequestPerfRecord( + timings.request_id, + commit_hash=get_git_commit_hash(), + tag="pipeline_stage_metrics", + stages=formatted_stages, + steps=timings.steps, + total_duration_ms=timings.total_duration_ms, + ) + + 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/__init__.py new file mode 100644 index 00000000..af2eb7d1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/__init__.py @@ -0,0 +1 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_common.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_common.py new file mode 100644 index 00000000..49dfa653 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_common.py @@ -0,0 +1,107 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +""" + Common generate cli test, one test for image and video each +""" + +import os +import unittest +from pathlib import Path + +from PIL import Image + +from sglang.multimodal_gen.test.test_utils import ( + TestCLIBase, + check_image_size, + is_mp4, + run_command, +) + + +class TestGenerate(TestCLIBase): + model_path = "black-forest-labs/FLUX.1-dev" + launch_file_name = "launch_flux.json" + output_name = "FLUX.1-dev, single gpu" + ext = "jpg" + + def test_generate_with_config(self): + test_dir = Path(__file__).parent + config_path = ( + (test_dir / ".." / "test_files" / self.launch_file_name) + .resolve() + .as_posix() + ) + command = [ + "sgl_diffusion", + "generate", + f"--config={config_path}", + ] + duration = run_command(command) + + self.assertIsNotNone(duration, f"Run command failed: {command}") + + # verify + self.verify_image(self.output_name) + + def test_generate_multiple_outputs(self): + command = [ + "sglang", + "generate", + "--prompt", + "A curious raccoon", + "--output-path=outputs", + f"--model-path={self.model_path}", + "--save-output", + f"--output-file-name={self.output_name}", + "--num-outputs-per-prompt=2", + "--width=720", + "--height=720", + ] + duration = run_command(command) + self.assertIsNotNone(duration, f"Run command failed: {command}") + + self.verify_image(f"{self.output_name}_0.{self.ext}") + self.verify_image(f"{self.output_name}_1.{self.ext}") + + def verify_image(self, output_name): + path = os.path.join("outputs", output_name) + with Image.open(path) as image: + check_image_size(self, image, 720, 720) + + def verify_video(self, output_name): + path = os.path.join("outputs", output_name) + with open(path, "rb") as f: + header = f.read(12) + assert is_mp4(header) + + +class TestWanGenerate(TestGenerate): + model_path = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" + launch_file_name = "launch_wan.json" + output_name = "Wan2.1-T2V-1.3B-Diffusers, single gpu" + ext = "mp4" + + def test_generate_multiple_outputs(self): + command = [ + "sglang", + "generate", + "--prompt", + "A curious raccoon", + "--output-path=outputs", + f"--model-path={self.model_path}", + "--save-output", + f"--output-file-name={self.output_name}", + "--num-outputs-per-prompt=2", + "--width=720", + "--height=720", + ] + duration = run_command(command) + self.assertIsNotNone(duration, f"Run command failed: {command}") + + self.verify_video(f"{self.output_name}_0.{self.ext}") + # FIXME: second video is a meaningless output + self.verify_video(f"{self.output_name}_1.{self.ext}") + + +if __name__ == "__main__": + unittest.main() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py new file mode 100644 index 00000000..400e8c71 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py @@ -0,0 +1,84 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +import unittest +from pathlib import Path + +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 TestGenerateBase + +logger = init_logger(__name__) + + +class TestFlux_T2V(TestGenerateBase): + model_path = "black-forest-labs/FLUX.1-dev" + extra_args = [] + data_type: DataType = DataType.IMAGE + thresholds = { + "test_single_gpu": 6.5 * 1.05, + "test_usp": 8.3 * 1.05, + } + + def test_cfg_parallel(self): + pass + + def test_mixed(self): + pass + + +class TestQwenImage(TestGenerateBase): + model_path = "Qwen/Qwen-Image" + extra_args = [] + data_type: DataType = DataType.IMAGE + thresholds = { + "test_single_gpu": 10.4 * 1.05, + "test_usp": 20.2 * 1.05, + } + + def test_cfg_parallel(self): + pass + + def test_mixed(self): + pass + + +class TestQwenImageEdit(TestGenerateBase): + model_path = "Qwen/Qwen-Image-Edit" + extra_args = [] + data_type: DataType = DataType.IMAGE + thresholds = { + "test_single_gpu": 33.4 * 1.05, + "test_usp": 26.9 * 1.05, + } + + prompt: str | None = ( + "Change the rabbit's color to purple, with a flash light background." + ) + + def setUp(self): + test_dir = Path(__file__).parent + img_path = (test_dir / ".." / "test_files" / "rabbit.jpg").resolve().as_posix() + self.base_command = [ + "sglang", + "generate", + "--text-encoder-cpu-offload", + "--pin-cpu-memory", + f"--prompt", + f"{self.prompt}", + "--save-output", + "--log-level=debug", + f"--width={self.width}", + f"--height={self.height}", + f"--output-path={self.output_path}", + ] + [f"--image-path={img_path}"] + + def test_cfg_parallel(self): + pass + + def test_mixed(self): + pass + + +if __name__ == "__main__": + del TestGenerateBase + unittest.main() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py new file mode 100644 index 00000000..1739bc8c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py @@ -0,0 +1,81 @@ +# 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.test_utils import TestGenerateBase + +logger = init_logger(__name__) + + +class TestFastWan2_1_T2V(TestGenerateBase): + model_path = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers" + extra_args = ["--attention-backend=video_sparse_attn"] + data_type: DataType = DataType.VIDEO + thresholds = { + "test_single_gpu": 13.0, + "test_cfg_parallel": 15.0, + "test_usp": 15.0, + "test_mixed": 15.0 * 1.05, + } + + # disabled for vsa + def test_usp(self): + pass + + +class TestFastWan2_2_T2V(TestGenerateBase): + model_path = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers" + extra_args = [] + data_type: DataType = DataType.VIDEO + thresholds = { + "test_single_gpu": 25.0, + "test_cfg_parallel": 30.0, + "test_usp": 30.0, + "test_mixed": 30.0, + } + + +class TestWan2_1_T2V(TestGenerateBase): + model_path = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" + extra_args = [] + data_type: DataType = DataType.VIDEO + thresholds = { + "test_single_gpu": 76.0 * 1.05, + "test_cfg_parallel": 46.5 * 1.05, + "test_usp": 39.8 * 1.05, + "test_mixed": 37.3 * 1.05, + } + + def test_mixed(self): + pass + + def test_cfg_parallel(self): + pass + + +class TestWan2_2_T2V(TestGenerateBase): + model_path = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" + extra_args = [] + data_type: DataType = DataType.VIDEO + thresholds = { + "test_single_gpu": 904.3 * 1.05, + "test_cfg_parallel": 446, + "test_usp": 316 * 1.05, + "test_mixed": 159, + } + + def test_single_gpu(self): + pass + + def test_mixed(self): + pass + + def test_cfg_parallel(self): + pass + + +if __name__ == "__main__": + del TestGenerateBase + unittest.main() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py new file mode 100644 index 00000000..3dc3c7c8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py @@ -0,0 +1,71 @@ +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.test_utils import TestGenerateBase + +logger = init_logger(__name__) + + +class TestGenerateTI2VBase(TestGenerateBase): + data_type: DataType = DataType.VIDEO + + @classmethod + def setUpClass(cls): + cls.base_command = [ + "sglang", + "generate", + "--prompt", + "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside.", + "--image-path", + "https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true", + "--save-output", + "--log-level=debug", + f"--output-path={cls.output_path}", + ] + cls.extra_args + + def test_single_gpu(self): + pass + + def test_cfg_parallel(self): + pass + + def test_mixed(self): + pass + + +class TestWan2_1_I2V_14B_480P(TestGenerateTI2VBase): + model_path = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" + thresholds = { + "test_usp": 557.9 * 1.05, + } + + +class TestWan2_1_I2V_14B_720P(TestGenerateTI2VBase): + model_path = "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers" + thresholds = { + "test_usp": 558.4 * 1.05, + } + + +class TestWan2_2_TI2V_5B(TestGenerateTI2VBase): + model_path = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" + # FIXME: doesn't work with vsa at the moment + # extra_args = ["--attention-backend=video_sparse_attn"] + thresholds = { + "test_usp": 82.3 * 1.05, + } + + +# OOM +# class TestWan2_2_I2V_A14B(TestGenerateTI2VBase): +# model_path = "Wan-AI/Wan2.2-I2V-A14B-Diffusers" +# # FIXME: doesn't work with vsa at the moment +# thresholds = { +# "test_usp": 66.3 * 1.05, +# } + + +if __name__ == "__main__": + del TestGenerateTI2VBase, TestGenerateBase + unittest.main() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_serve.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_serve.py new file mode 100644 index 00000000..156cf215 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/cli/test_serve.py @@ -0,0 +1,301 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +import asyncio +import base64 +import subprocess +import tempfile +import time +import unittest +import uuid +from contextlib import contextmanager +from pathlib import Path +from urllib.request import urlopen + +from openai import OpenAI + +from sglang.multimodal_gen.runtime.utils.common import kill_process_tree +from sglang.multimodal_gen.test.test_utils import is_mp4, is_png, wait_for_port + + +@contextmanager +def downloaded_temp_file(url: str, prefix: str = "i2v_input_", suffix: str = ".jpg"): + tmp_path = Path(tempfile.gettempdir()) / f"{prefix}{uuid.uuid4().hex}{suffix}" + with urlopen(url) as resp: + tmp_path.write_bytes(resp.read()) + try: + yield tmp_path + finally: + try: + tmp_path.unlink(missing_ok=True) + except Exception: + pass + + +def wait_for_video_completion(client, video_id, timeout=300, check_interval=3): + start = time.time() + video = client.videos.retrieve(video_id) + + while video.status not in ("completed", "failed"): + time.sleep(check_interval) + video = client.videos.retrieve(video_id) + assert time.time() - start < timeout, "video generate timeout" + + return video + + +class TestVideoHttpServer(unittest.TestCase): + model_name = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" + timeout = 500 + extra_args = [] + + def _create_wait_and_download( + self, client: OpenAI, prompt: str, size: str + ) -> bytes: + + video = client.videos.create(prompt=prompt, size=size) + video_id = video.id + self.assertEqual(video.status, "queued") + + video = wait_for_video_completion(client, video_id, timeout=self.timeout) + self.assertEqual(video.status, "completed", "video generate failed") + + response = client.videos.download_content( + video_id=video_id, + ) + content = response.read() + return content + + @classmethod + def setUpClass(cls): + cls.base_command = [ + "sglang", + "serve", + "--model-path", + f"{cls.model_name}", + "--port", + "30010", + ] + + process = subprocess.Popen( + cls.base_command + cls.extra_args, + # stdout=subprocess.PIPE, + # stderr=subprocess.PIPE, + text=True, + bufsize=1, + ) + cls.pid = process.pid + wait_for_port(host="127.0.0.1", port=30010) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.pid) + + def test_http_server_basic(self): + client = OpenAI( + api_key="sk-proj-1234567890", base_url="http://localhost:30010/v1" + ) + content = self._create_wait_and_download( + client, "A plane is taking off.", "832x480" + ) + self.assertTrue(is_mp4(content)) + + def test_concurrent_requests(self): + client = OpenAI( + api_key="sk-proj-1234567890", base_url="http://localhost:30010/v1" + ) + + num_requests = 2 + + async def generate_and_check_video(prompt, size): + content = await asyncio.to_thread( + self._create_wait_and_download, client, prompt, size + ) + self.assertTrue(is_mp4(content)) + + async def send_concurrent_requests(): + tasks = [ + generate_and_check_video( + "A ship is beside the port.", + "832x480", + ) + for _ in range(num_requests) + ] + await asyncio.gather(*tasks) + + asyncio.run(send_concurrent_requests()) + + +class TestImage2VideoHttpServer(unittest.TestCase): + model_name = "Wan-AI/Wan2.2-I2V-A14B-Diffusers" + timeout = 1200 + extra_args = [] + + def _create_wait_and_download( + self, client: OpenAI, prompt: str, size: str + ) -> bytes: + + image_url = "https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true" + with downloaded_temp_file( + image_url, prefix="i2v_input_", suffix=".jpg" + ) as tmp_path: + video = client.videos.create( + prompt=prompt, + input_reference=tmp_path, + size=size, + seconds=10, + extra_body={"fps": 16, "num_frames": 125}, + ) + # TODO: Some combinations of num_frames and fps may cause errors and need further investigation. + video_id = video.id + self.assertEqual(video.status, "queued") + + video = wait_for_video_completion(client, video_id, timeout=self.timeout) + self.assertEqual(video.status, "completed", "video generate failed") + + response = client.videos.download_content( + video_id=video_id, + ) + content = response.read() + return content + + @classmethod + def setUpClass(cls): + cls.base_command = [ + "sglang", + "serve", + "--model-path", + f"{cls.model_name}", + "--num-gpus", + "4", + "--ulysses-degree", + "4", + "--port", + "30010", + ] + + process = subprocess.Popen( + cls.base_command + cls.extra_args, + text=True, + bufsize=1, + ) + cls.pid = process.pid + wait_for_port(host="127.0.0.1", port=30010) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.pid) + + def test_http_server_basic(self): + client = OpenAI( + api_key="sk-proj-1234567890", base_url="http://localhost:30010/v1" + ) + content = self._create_wait_and_download( + client, "A cat surfing on the sea.", "832x480" + ) + self.assertTrue(is_mp4(content)) + + def test_concurrent_requests(self): + client = OpenAI( + api_key="sk-proj-1234567890", base_url="http://localhost:30010/v1" + ) + + num_requests = 2 + + async def generate_and_check_video(prompt, size): + content = await asyncio.to_thread( + self._create_wait_and_download, client, prompt, size + ) + self.assertTrue(is_mp4(content)) + + async def send_concurrent_requests(): + tasks = [ + generate_and_check_video( + "A cat surfing on the sea.", + "832x480", + ) + for _ in range(num_requests) + ] + await asyncio.gather(*tasks) + + asyncio.run(send_concurrent_requests()) + + +class TestImageHttpServer(unittest.TestCase): + + def _create_wait_and_download( + self, client: OpenAI, prompt: str, size: str + ) -> bytes: + img = client.images.generate( + model="gpt-image-1", + prompt=prompt, + n=1, + size=size, + response_format="b64_json", + output_format="png", + ) + image_bytes = base64.b64decode(img.data[0].b64_json) + return image_bytes + + @classmethod + def setUpClass(cls): + cls.base_command = [ + "sglang", + "serve", + "--model-path", + "Qwen/Qwen-Image", + "--port", + "30020", + ] + + process = subprocess.Popen( + cls.base_command, + # stdout=subprocess.PIPE, + # stderr=subprocess.PIPE, + text=True, + bufsize=1, + ) + cls.pid = process.pid + wait_for_port(host="127.0.0.1", port=30020) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.pid) + + def test_http_server_basic(self): + client = OpenAI( + api_key="sk-proj-1234567890", base_url="http://localhost:30020/v1" + ) + content = self._create_wait_and_download( + client, "A calico cat playing a piano on stage", "832x480" + ) + self.assertTrue(is_png(content)) + + def test_concurrent_requests(self): + client = OpenAI( + api_key="sk-proj-1234567890", base_url="http://localhost:30020/v1" + ) + + num_requests = 2 + + async def generate_and_check_image(prompt, size): + content = await asyncio.to_thread( + self._create_wait_and_download, client, prompt, size + ) + self.assertTrue(is_png(content)) + + async def send_concurrent_requests(): + tasks = [ + generate_and_check_image( + "A dog playing a piano on stage", + "832x480", + ) + for _ in range(num_requests) + ] + await asyncio.gather(*tasks) + + asyncio.run(send_concurrent_requests()) + + +if __name__ == "__main__": + # del TestPerform·anceBase + unittest.main() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/run_suite.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/run_suite.py new file mode 100644 index 00000000..206d06cc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/run_suite.py @@ -0,0 +1,173 @@ +""" +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 2 +""" + +import argparse +import os +import subprocess +import sys +from pathlib import Path + +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +SUITES = { + "1-gpu": [ + "test_server_a.py", + "test_server_b.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 + ], +} + + +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", + ) + return parser.parse_args() + + +def run_pytest(files): + if not files: + print("No files to run.") + return 0 + + base_cmd = [sys.executable, "-m", "pytest", "-s", "-v", "--log-cli-level=INFO"] + + max_retries = 2 + # 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") + cmd.extend(files) + + if i > 0: + logger.info( + f"Performance assertion failed. Retrying ({i}/{max_retries}) with --last-failed..." + ) + + logger.info(f"Running command: {' '.join(cmd)}") + + process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + + output_lines = [] + while True: + line = process.stdout.readline() + if not line and process.poll() is not None: + break + if line: + sys.stdout.write(line) + output_lines.append(line) + + returncode = process.poll() + + if returncode == 0: + return 0 + + # check if the failure is due to an assertion in test_server_utils.py + full_output = "".join(output_lines) + is_perf_assertion = ( + "multimodal_gen/test/server/test_server_utils.py" in full_output + and "AssertionError" in full_output + ) + + if not is_perf_assertion: + return returncode + + return returncode + + +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] + + suite_files_abs = [] + for f_rel in suite_files_rel: + f_abs = target_dir / f_rel + if not f_abs.exists(): + print(f"Warning: Test file {f_rel} not found in {target_dir}. 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(0) + + # 3. partitioning + my_files = [ + f + for i, f in enumerate(suite_files_abs) + if i % args.total_partitions == args.partition_id + ] + + print( + f"Suite: {args.suite} | Partition: {args.partition_id}/{args.total_partitions}" + ) + print(f"Selected {len(my_files)} files:") + for f in my_files: + print(f" - {os.path.basename(f)}") + + if not my_files: + print("No files assigned to this partition. Exiting success.") + sys.exit(0) + + # 4. execute + exit_code = run_pytest(my_files) + sys.exit(exit_code) + + +if __name__ == "__main__": + main() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/conftest.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/conftest.py new file mode 100644 index 00000000..96b49591 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/conftest.py @@ -0,0 +1,53 @@ +_GLOBAL_PERF_RESULTS = [] + + +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. + """ + if not _GLOBAL_PERF_RESULTS: + return + + 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(_GLOBAL_PERF_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(_GLOBAL_PERF_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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/perf_baselines.json b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/perf_baselines.json new file mode 100644 index 00000000..0757aa51 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -0,0 +1,1018 @@ +{ + "metadata": { + "model": "Diffusion Server", + "hardware": "CI H100 80GB pool", + "description": "Reference numbers captured from the CI diffusion server baseline run" + }, + "tolerances": { + "e2e": 0.1, + "denoise_stage": 0.05, + "non_denoise_stage": 0.4, + "denoise_step": 0.2, + "denoise_agg": 0.1 + }, + "improvement_reporting": { + "threshold": 0.2 + }, + "sampling": { + "step_fractions": [ + 0.0, + 0.2, + 0.4, + 0.6, + 0.8, + 1.0 + ], + "warmup_requests": { + "text": 1, + "image_edit": 0 + } + }, + "scenarios": { + "qwen_image_t2i": { + "notes": "Single-image generation using the default prompt", + "expected_e2e_ms": 74500.0, + "expected_avg_denoise_ms": 422.42, + "expected_median_denoise_ms": 410.62, + "stages_ms": { + "InputValidationStage": 0.1, + "TextEncodingStage": 834.2, + "ConditioningStage": 0.1, + "TimestepPreparationStage": 10.6, + "LatentPreparationStage": 11.8, + "DenoisingStage": 21202.6, + "DecodingStage": 751.1 + }, + "denoise_step_ms": { + "0": 1077.77, + "1": 345.13, + "2": 413.8, + "3": 405.49, + "4": 408.14, + "5": 409.06, + "6": 408.85, + "7": 410.53, + "8": 407.51, + "9": 409.44, + "10": 408.65, + "11": 410.14, + "12": 411.74, + "13": 409.59, + "14": 409.17, + "15": 410.78, + "16": 410.66, + "17": 410.58, + "18": 411.27, + "19": 410.51, + "20": 409.03, + "21": 410.16, + "22": 409.42, + "23": 411.03, + "24": 410.18, + "25": 409.72, + "26": 410.26, + "27": 410.21, + "28": 410.71, + "29": 470.76, + "30": 411.06, + "31": 410.1, + "32": 410.55, + "33": 410.77, + "34": 410.74, + "35": 411.75, + "36": 410.78, + "37": 411.56, + "38": 410.85, + "39": 411.08, + "40": 411.12, + "41": 411.1, + "42": 411.09, + "43": 410.87, + "44": 411.37, + "45": 411.68, + "46": 411.0, + "47": 410.09, + "48": 412.72, + "49": 410.42 + } + }, + "qwen_image_t2i_2_gpus": { + "stages_ms": { + "InputValidationStage": 0.04, + "TextEncodingStage": 693.2, + "ConditioningStage": 0.02, + "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": { + "InputValidationStage": 0.03, + "TextEncodingStage": 81.49, + "ConditioningStage": 0.01, + "TimestepPreparationStage": 2.43, + "LatentPreparationStage": 6.29, + "DenoisingStage": 8381.3, + "DecodingStage": 653.03 + }, + "denoise_step_ms": { + "0": 165.27, + "1": 58.88, + "2": 166.85, + "3": 166.51, + "4": 166.77, + "5": 167.55, + "6": 172.4, + "7": 167.77, + "8": 167.51, + "9": 167.22, + "10": 168.19, + "11": 167.74, + "12": 168.48, + "13": 168.08, + "14": 168.16, + "15": 167.15, + "16": 167.05, + "17": 169.27, + "18": 167.96, + "19": 167.74, + "20": 168.21, + "21": 167.07, + "22": 167.35, + "23": 167.06, + "24": 169.28, + "25": 169.41, + "26": 168.92, + "27": 167.59, + "28": 167.57, + "29": 170.42, + "30": 166.24, + "31": 168.33, + "32": 168.56, + "33": 168.62, + "34": 167.28, + "35": 167.12, + "36": 168.21, + "37": 168.78, + "38": 168.89, + "39": 167.74, + "40": 168.57, + "41": 167.89, + "42": 168.03, + "43": 167.61, + "44": 167.75, + "45": 168.03, + "46": 168.81, + "47": 168.29, + "48": 168.64, + "49": 168.78 + }, + "expected_e2e_ms": 9275.51, + "expected_avg_denoise_ms": 165.83, + "expected_median_denoise_ms": 169.33 + }, + "flux_2_image_t2i": { + "stages_ms": { + "InputValidationStage": 0.05, + "TextEncodingStage": 530.93, + "ImageVAEEncodingStage": 0.0, + "ConditioningStage": 0.02, + "LatentPreparationStage": 12.71, + "TimestepPreparationStage": 2.91, + "DenoisingStage": 26403.1, + "DecodingStage": 286.85 + }, + "denoise_step_ms": { + "0": 511.3, + "1": 132.57, + "2": 541.19, + "3": 518.93, + "4": 541.2, + "5": 520.28, + "6": 532.47, + "7": 525.68, + "8": 538.25, + "9": 525.84, + "10": 526.13, + "11": 525.67, + "12": 524.63, + "13": 530.57, + "14": 530.46, + "15": 529.94, + "16": 532.47, + "17": 527.88, + "18": 527.7, + "19": 525.08, + "20": 525.72, + "21": 529.3, + "22": 522.59, + "23": 529.75, + "24": 523.46, + "25": 528.72, + "26": 526.92, + "27": 528.62, + "28": 522.77, + "29": 528.35, + "30": 528.05, + "31": 528.89, + "32": 525.34, + "33": 530.36, + "34": 529.19, + "35": 526.92, + "36": 528.16, + "37": 525.03, + "38": 527.33, + "39": 527.96, + "40": 527.81, + "41": 524.79, + "42": 528.46, + "43": 532.49, + "44": 526.95, + "45": 533.14, + "46": 529.32, + "47": 528.51, + "48": 532.14, + "49": 529.29 + }, + "expected_e2e_ms": 27648.69, + "expected_avg_denoise_ms": 520.09, + "expected_median_denoise_ms": 528.0 + }, + "flux_2_ti2i": { + "stages_ms": { + "InputValidationStage": 99.82, + "TextEncodingStage": 519.88, + "ImageVAEEncodingStage": 254.56, + "ConditioningStage": 0.01, + "LatentPreparationStage": 12.4, + "TimestepPreparationStage": 2.71, + "DenoisingStage": 54705.41, + "DecodingStage": 311.13 + }, + "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_image_t2i_2_gpus": { + "stages_ms": { + "InputValidationStage": 0.03, + "TextEncodingStage": 74.47, + "ConditioningStage": 0.01, + "TimestepPreparationStage": 2.23, + "LatentPreparationStage": 6.17, + "DenoisingStage": 8400.49, + "DecodingStage": 381.56 + }, + "denoise_step_ms": { + "0": 166.27, + "1": 59.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": 104.21, + "ConditioningStage": 0.01, + "TimestepPreparationStage": 1.33, + "LatentPreparationStage": 1.13, + "DenoisingStage": 850.85, + "DecodingStage": 289.32 + }, + "denoise_step_ms": { + "0": 101.56, + "1": 28.26, + "2": 101.74, + "3": 101.68, + "4": 102.19, + "5": 102.05, + "6": 102.03, + "7": 102.28, + "8": 105.54 + }, + "expected_e2e_ms": 1383.47, + "expected_avg_denoise_ms": 94.15, + "expected_median_denoise_ms": 102.03 + }, + "qwen_image_edit_ti2i": { + "notes": "single uploaded reference image, Qwen/Qwen-Image-Edit", + "expected_e2e_ms": 138500.0, + "expected_avg_denoise_ms": 720.0, + "expected_median_denoise_ms": 718.0, + "stages_ms": { + "InputValidationStage": 23, + "ImageEncodingStage": 1485.0, + "ImageVAEEncodingStage": 350.0, + "ConditioningStage": 0.13, + "TimestepPreparationStage": 13.78, + "LatentPreparationStage": 15.0, + "DenoisingStage": 36000.0, + "DecodingStage": 850.0 + }, + "denoise_step_ms": { + "0": 720.0, + "1": 720.0, + "2": 720.0, + "3": 720.0, + "4": 720.0, + "5": 720.0, + "6": 720.0, + "7": 720.0, + "8": 720.0, + "9": 720.0, + "10": 720.0, + "11": 720.0, + "12": 720.0, + "13": 720.0, + "14": 720.0, + "15": 720.0, + "16": 720.0, + "17": 720.0, + "18": 720.0, + "19": 720.0, + "20": 720.0, + "21": 720.0, + "22": 720.0, + "23": 720.0, + "24": 720.0, + "25": 720.0, + "26": 720.0, + "27": 720.0, + "28": 720.0, + "29": 720.0, + "30": 720.0, + "31": 720.0, + "32": 720.0, + "33": 720.0, + "34": 720.0, + "35": 720.0, + "36": 720.0, + "37": 720.0, + "38": 720.0, + "39": 720.0, + "40": 720.0, + "41": 720.0, + "42": 720.0, + "43": 720.0, + "44": 720.0, + "45": 720.0, + "46": 720.0, + "47": 720.0, + "48": 720.0, + "49": 720.0 + } + }, + "wan2_1_t2v_1.3b": { + "stages_ms": { + "InputValidationStage": 0.06, + "TextEncodingStage": 3595.12, + "ConditioningStage": 0.02, + "TimestepPreparationStage": 2.39, + "LatentPreparationStage": 15.27, + "DenoisingStage": 91099.4, + "DecodingStage": 4330.65, + "per_frame_generation": null + }, + "denoise_step_ms": { + "0": 2918.67, + "1": 1784.23, + "2": 1797.72, + "3": 1798.8, + "4": 1798.19, + "5": 1799.27, + "6": 1798.54, + "7": 1798.67, + "8": 1798.76, + "9": 1798.34, + "10": 1799.22, + "11": 1798.61, + "12": 1799.4, + "13": 1799.04, + "14": 1797.41, + "15": 1799.05, + "16": 1798.32, + "17": 1799.12, + "18": 1799.56, + "19": 1797.01, + "20": 1798.28, + "21": 1799.06, + "22": 1800.05, + "23": 1797.76, + "24": 1798.16, + "25": 1798.62, + "26": 1798.64, + "27": 1799.44, + "28": 1798.79, + "29": 1798.13, + "30": 1797.47, + "31": 1799.4, + "32": 1798.77, + "33": 1799.47, + "34": 1798.49, + "35": 1796.51, + "36": 1799.68, + "37": 1799.24, + "38": 1798.49, + "39": 1799.66, + "40": 1797.04, + "41": 1799.58, + "42": 1797.35, + "43": 1798.07, + "44": 1798.6, + "45": 1798.95, + "46": 1799.51, + "47": 1798.25, + "48": 1799.04, + "49": 1798.34 + }, + "expected_e2e_ms": 99083.75, + "expected_avg_denoise_ms": 1820.02, + "expected_median_denoise_ms": 1798.65 + }, + "wan2_2_ti2v_5b": { + "stages_ms": { + "InputValidationStage": 96.27, + "TextEncodingStage": 2238.81, + "ConditioningStage": 0.02, + "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 + }, + "fastwan2_2_ti2v_5b": { + "stages_ms": { + "InputValidationStage": 88.86, + "TextEncodingStage": 2327.87, + "ConditioningStage": 0.01, + "TimestepPreparationStage": 58.66, + "LatentPreparationStage": 28.55, + "DmdDenoisingStage": 4438.3, + "DecodingStage": 14177.77, + "per_frame_generation": null + }, + "denoise_step_ms": { + "0": 2022.21, + "1": 1263.17, + "2": 1149.59 + }, + "expected_e2e_ms": 21133.36, + "expected_avg_denoise_ms": 1478.32, + "expected_median_denoise_ms": 1263.17 + }, + "fast_hunyuan_video": { + "stages_ms": { + "InputValidationStage": 0.09, + "TextEncodingStage": 845.64, + "ConditioningStage": 0.04, + "TimestepPreparationStage": 125.22, + "LatentPreparationStage": 29.34, + "DenoisingStage": 3860.64, + "DecodingStage": 2580.55 + }, + "denoise_step_ms": { + "0": 2063.08, + "1": 164.02, + "2": 406.99, + "3": 407.95, + "4": 407.51, + "5": 404.2 + }, + "expected_e2e_ms": 7487.87, + "expected_avg_denoise_ms": 642.29, + "expected_median_denoise_ms": 407.25 + }, + "wan2_2_i2v_a14b_2gpu": { + "stages_ms": { + "InputValidationStage": 59.33, + "TextEncodingStage": 6062.41, + "ConditioningStage": 0.02, + "TimestepPreparationStage": 2.2, + "LatentPreparationStage": 8.93, + "ImageVAEEncodingStage": 2075.47, + "DenoisingStage": 382628.41, + "DecodingStage": 2820.89 + }, + "denoise_step_ms": { + "0": 31228.27, + "1": 7723.86, + "2": 7769.69, + "3": 7795.93, + "4": 7815.58, + "5": 7829.48, + "6": 7827.34, + "7": 7825.35, + "8": 7828.05, + "9": 7809.53, + "10": 7801.29, + "11": 7790.96, + "12": 7785.88, + "13": 7785.5, + "14": 7780.32, + "15": 55411.1, + "16": 7722.27, + "17": 7761.31, + "18": 7789.46, + "19": 7800.6, + "20": 7814.91, + "21": 7799.62, + "22": 7801.25, + "23": 7798.27, + "24": 7797.67, + "25": 7795.97, + "26": 7781.74, + "27": 7784.16, + "28": 7796.64, + "29": 7789.75, + "30": 7792.13, + "31": 7790.99, + "32": 7778.1, + "33": 7777.78, + "34": 7780.56, + "35": 7778.22, + "36": 7770.88, + "37": 7771.56, + "38": 7767.82, + "39": 7769.23 + }, + "expected_e2e_ms": 393606.77, + "expected_avg_denoise_ms": 9565.48, + "expected_median_denoise_ms": 7790.98 + }, + "wan2_1_i2v_14b_480P_2gpu": { + "stages_ms": { + "InputValidationStage": 38.23, + "TextEncodingStage": 3550.36, + "ImageEncodingStage": 3462.55, + "ConditioningStage": 0.01, + "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, + "ConditioningStage": 0.01, + "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": 2507.83, + "ConditioningStage": 0.02, + "TimestepPreparationStage": 3.22, + "LatentPreparationStage": 2.99, + "DenoisingStage": 103136.69, + "DecodingStage": 1431.71 + }, + "denoise_step_ms": { + "0": 24471.86, + "1": 757.31, + "2": 760.07, + "3": 758.74, + "4": 762.4, + "5": 755.83, + "6": 760.06, + "7": 756.38, + "8": 755.38, + "9": 754.25, + "10": 754.51, + "11": 753.46, + "12": 753.67, + "13": 753.08, + "14": 754.83, + "15": 753.04, + "16": 754.28, + "17": 754.45, + "18": 758.19, + "19": 756.23, + "20": 755.14, + "21": 755.92, + "22": 759.52, + "23": 762.09, + "24": 756.8, + "25": 758.86, + "26": 48787.27, + "27": 758.5, + "28": 757.57, + "29": 757.16, + "30": 758.43, + "31": 763.31, + "32": 753.69, + "33": 754.91, + "34": 752.03, + "35": 763.65, + "36": 760.96, + "37": 754.31, + "38": 753.64, + "39": 756.95 + }, + "expected_e2e_ms": 106895.63, + "expected_avg_denoise_ms": 2550.47, + "expected_median_denoise_ms": 756.59 + }, + "wan2_1_t2v_14b_2gpu": { + "stages_ms": { + "InputValidationStage": 0.05, + "TextEncodingStage": 2310.34, + "ConditioningStage": 0.02, + "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 + } + } +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_2_gpu_a.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_2_gpu_a.py new file mode 100644 index 00000000..3668f63e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_2_gpu_a.py @@ -0,0 +1,25 @@ +""" +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_2_gpu_b.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_2_gpu_b.py new file mode 100644 index 00000000..2c9b5cdc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_2_gpu_b.py @@ -0,0 +1,25 @@ +""" +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_a.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_a.py new file mode 100644 index 00000000..fdf072ec --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_a.py @@ -0,0 +1,31 @@ +""" +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_b.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_b.py new file mode 100644 index 00000000..1a0432db --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_b.py @@ -0,0 +1,31 @@ +""" +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_common.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_common.py new file mode 100644 index 00000000..18b56f5d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -0,0 +1,384 @@ +""" +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 os +from pathlib import Path +from typing import Any, Callable + +import openai +import pytest +from openai import OpenAI + +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.conftest import _GLOBAL_PERF_RESULTS +from sglang.multimodal_gen.test.server.test_server_utils import ( + VALIDATOR_REGISTRY, + PerformanceValidator, + ServerContext, + ServerManager, + WarmupRunner, + download_image_from_url, + 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 ( + get_dynamic_server_port, + is_image_url, + read_perf_logs, + 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.""" + default_port = get_dynamic_server_port() + port = int(os.environ.get("SGLANG_TEST_SERVER_PORT", default_port)) + server_args = case.server_args + sampling_params = case.sampling_params + extra_args = os.environ.get("SGLANG_TEST_SERVE_ARGS", "") + extra_args += ( + f" --num-gpus {server_args.num_gpus} --ulysses-degree {server_args.num_gpus}" + ) + + # 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, + ) + ctx = manager.start() + + try: + # Reconstruct output size for OpenAI API + output_size = sampling_params.output_size + warmup = WarmupRunner( + port=ctx.port, + model=server_args.model_path, + prompt=sampling_params.prompt or "A colorful raccoon icon", + output_size=output_size, + ) + warmup.run_text_warmups(server_args.warmup_text) + + if ( + case.server_args.warmup_edit > 0 + and case.sampling_params.prompt + and sampling_params.image_path + ): + # Handle URL or local path + 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) + + warmup.run_edit_warmups( + count=server_args.warmup_edit, + edit_prompt=sampling_params.prompt, + image_path=image_path, + ) + 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]] = [] + + @classmethod + def setup_class(cls): + cls._perf_results = [] + cls._improved_baselines = [] + + @classmethod + def teardown_class(cls): + for result in cls._perf_results: + result["class_name"] = cls.__name__ + _GLOBAL_PERF_RESULTS.append(result) + + 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) + + 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], str], + ) -> RequestPerfRecord: + """Run generation and collect performance records.""" + log_path = ctx.perf_log_path + prev_len = len(read_perf_logs(log_path)) + log_wait_timeout = 30 + + client = self._client(ctx) + rid = generate_fn(case_id, client) + + req_perf_record, _ = wait_for_req_perf_record( + rid, + prev_len, + log_path, + timeout=log_wait_timeout, + ) + + return req_perf_record + + 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 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) + + 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) + + 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 test_diffusion_perf( + self, + case: DiffusionTestCase, + diffusion_server: ServerContext, + ): + """Single parametrized test that runs for all cases. + + Pytest will execute this test once per case in ONE_GPU_CASES, + with test IDs like: + - test_diffusion_perf[qwen_image_text] + - test_diffusion_perf[qwen_image_edit] + - etc. + """ + generate_fn = get_generate_fn( + model_path=case.server_args.model_path, + modality=case.server_args.modality, + sampling_params=case.sampling_params, + ) + perf_record = self.run_and_collect( + diffusion_server, + case.id, + generate_fn, + ) + self._validate_and_record(case, perf_record) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_utils.py new file mode 100644 index 00000000..4e3e788a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -0,0 +1,696 @@ +""" +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, OpenAI + +from sglang.multimodal_gen.benchmarks.compare_perf import calculate_upper_bound +from sglang.multimodal_gen.runtime.utils.common import kill_process_tree +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.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 ( + is_image_url, + prepare_perf_log, + validate_image, + validate_openai_video, +) + +logger = init_logger(__name__) + + +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 + + +@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 + + +class ServerManager: + """Manages diffusion server lifecycle.""" + + def __init__( + self, + model: str, + port: int, + wait_deadline: float = 1200.0, + extra_args: str = "", + ): + self.model = model + self.port = port + self.wait_deadline = wait_deadline + self.extra_args = extra_args + + def start(self) -> ServerContext: + """Start the diffusion server and wait for readiness.""" + 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() + + # 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) + 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." + + 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) + logger.info("[server-test] Waiting for server... elapsed=%ss", elapsed) + time.sleep(5) + + 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 WarmupRunner: + """Handles warmup requests for a server.""" + + def __init__( + self, + port: int, + model: str, + prompt: str, + output_size: str, + ): + self.client = OpenAI( + api_key="sglang-anything", + base_url=f"http://localhost:{port}/v1", + ) + self.model = model + self.prompt = prompt + self.output_size = output_size + + def run_text_warmups(self, count: int) -> None: + """Run text-to-image warmup requests.""" + if count <= 0: + return + + logger.info("[server-test] Running %s text warm-up(s)", count) + for _ in range(count): + result = self.client.images.generate( + model=self.model, + prompt=self.prompt, + n=1, + size=self.output_size, + response_format="b64_json", + ) + validate_image(result.data[0].b64_json) + + def run_edit_warmups( + self, + count: int, + edit_prompt: str, + image_path: Path, + ) -> None: + """Run image-edit warmup requests.""" + if count <= 0: + return + + if not image_path.exists(): + logger.warning( + "[server-test] Skipping edit warmup: image missing at %s", image_path + ) + return + + logger.info("[server-test] Running %s edit warm-up(s)", count) + for _ in range(count): + with image_path.open("rb") as fh: + result = self.client.images.edit( + model=self.model, + image=fh, + prompt=edit_prompt, + n=1, + size=self.output_size, + response_format="b64_json", + ) + validate_image(result.data[0].b64_json) + + +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. + """ + 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, + ) + + +# Registry of validators by name +VALIDATOR_REGISTRY = { + "default": PerformanceValidator, + "video": VideoPerformanceValidator, +} + + +def get_generate_fn( + model_path: str, + modality: str, + sampling_params: DiffusionSamplingParams, +) -> Callable[[str, Client], str]: + """Return appropriate generation function for the case.""" + + 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, + ) -> str: + """ + Create a video job via /v1/videos, poll until completion, + then download the binary content and validate it. + """ + + 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 + + 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" + timeout = 3600.0 if is_baseline_generation_mode else 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"{id}: video job {video_id} timed out during baseline generation. " + "Attempting to collect performance data anyway." + ) + return video_id + + pytest.fail(f"{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) + + tmp_path = f"{video_id}.mp4" + with open(tmp_path, "wb") as f: + f.write(content) + 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 + + video_seconds = sampling_params.seconds or 4 + + def generate_image(case_id, client) -> str: + """T2I: Text to Image generation.""" + if not sampling_params.prompt: + pytest.skip(f"{id}: no text prompt configured") + + response = client.images.with_raw_response.generate( + model=model_path, + prompt=sampling_params.prompt, + n=1, + size=sampling_params.output_size, + response_format="b64_json", + ) + result = response.parse() + validate_image(result.data[0].b64_json) + + img_data = base64.b64decode(result.data[0].b64_json) + tmp_path = f"{result.created}.png" + with open(tmp_path, "wb") as f: + f.write(img_data) + upload_file_to_slack( + case_id=case_id, + model=model_path, + prompt=sampling_params.prompt, + file_path=tmp_path, + ) + os.remove(tmp_path) + + return str(result.created) + + def generate_image_edit(case_id, client) -> str: + """TI2I: Text + Image ? Image edit.""" + if not sampling_params.prompt or not sampling_params.image_path: + pytest.skip(f"{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"{id}: file missing: {image_path}") + + with image_path.open("rb") as fh: + response = client.images.with_raw_response.edit( + model=model_path, + image=fh, + prompt=sampling_params.prompt, + n=1, + size=sampling_params.output_size, + response_format="b64_json", + ) + rid = response.headers.get("x-request-id", "") + + result = response.parse() + validate_image(result.data[0].b64_json) + + img_data = base64.b64decode(result.data[0].b64_json) + tmp_path = f"{rid}.png" + with open(tmp_path, "wb") as f: + f.write(img_data) + 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 + + def generate_video(case_id, client) -> str: + """T2V: Text ? Video.""" + if not sampling_params.prompt: + pytest.skip(f"{id}: no text prompt configured") + + return _create_and_download_video( + client, + case_id, + model=model_path, + prompt=sampling_params.prompt, + size=sampling_params.output_size, + seconds=video_seconds, + ) + + def generate_image_to_video(case_id, client) -> str: + """I2V: Image ? Video (optional prompt).""" + if not sampling_params.image_path: + pytest.skip(f"{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"{id}: file missing: {image_path}") + + with image_path.open("rb") as fh: + return _create_and_download_video( + client, + case_id, + model=model_path, + prompt=sampling_params.prompt, + size=sampling_params.output_size, + seconds=video_seconds, + input_reference=fh, + ) + + def generate_text_image_to_video(case_id, client) -> str: + """TI2V: Text + Image ? Video.""" + if not sampling_params.prompt or not sampling_params.image_path: + pytest.skip(f"{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"{id}: file missing: {image_path}") + + with image_path.open("rb") as fh: + return _create_and_download_video( + client, + case_id, + model=model_path, + prompt=sampling_params.prompt, + size=sampling_params.output_size, + seconds=video_seconds, + input_reference=fh, + ) + + if modality == "video": + if sampling_params.image_path and sampling_params.prompt: + 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: + fn = generate_image_edit + else: + fn = generate_image + + return fn diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/testcase_configs.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/testcase_configs.py new file mode 100644 index 00000000..75022e0b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -0,0 +1,484 @@ +""" +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/test_server_a.py -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 +from pathlib import Path +from typing import Sequence + +from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord + + +@dataclass +class ToleranceConfig: + """Tolerance ratios for performance validation.""" + + e2e: float + denoise_stage: float + non_denoise_stage: float + denoise_step: float + denoise_agg: float + + +@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] + warmup_defaults: dict[str, int] + 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) + + tol_data = data["tolerances"] + tolerances = ToleranceConfig( + 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"]) + ), + ) + + 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"]), + warmup_defaults=data["sampling"].get("warmup_requests", {}), + tolerances=tolerances, + improvement_threshold=data.get("improvement_reporting", {}).get( + "threshold", 0.2 + ), + ) + + +@dataclass(frozen=True) +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" + + warmup_text: int = 1 # number of text-to-image/video warmups + warmup_edit: int = 0 # number of image/video-edit warmups + custom_validator: str | None = None # optional custom validator name + # resources + num_gpus: int = 1 + + +@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 + + +@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 + + +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, + ) + + +# 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="Qwen/Qwen-Image", + modality="image", + warmup_text=1, + warmup_edit=0, + ), + DiffusionSamplingParams( + prompt="A futuristic cityscape at sunset with flying cars", + output_size="1024x1024", + ), + ), + DiffusionTestCase( + "flux_image_t2i", + DiffusionServerArgs( + model_path="black-forest-labs/FLUX.1-dev", + modality="image", + warmup_text=1, + warmup_edit=0, + ), + DiffusionSamplingParams( + prompt="A futuristic cityscape at sunset with flying cars", + output_size="1024x1024", + ), + ), + DiffusionTestCase( + "flux_2_image_t2i", + DiffusionServerArgs( + model_path="black-forest-labs/FLUX.2-dev", + modality="image", + warmup_text=1, + warmup_edit=0, + ), + DiffusionSamplingParams( + prompt="A futuristic cityscape at sunset with flying cars", + output_size="1024x1024", + ), + ), + DiffusionTestCase( + "zimage_image_t2i", + DiffusionServerArgs( + model_path="Tongyi-MAI/Z-Image-Turbo", + modality="image", + warmup_text=1, + warmup_edit=0, + ), + DiffusionSamplingParams( + prompt="Doraemon is eating dorayaki.", + output_size="1024x1024", + ), + ), + # === Text and Image to Image (TI2I) === + # TODO: Timeout with Torch2.9. Add back when it can pass CI + # DiffusionTestCase( + # id="qwen_image_edit_ti2i", + # model_path="Qwen/Qwen-Image-Edit", + # modality="image", + # prompt=None, # not used for editing + # output_size="1024x1536", + # warmup_text=0, + # warmup_edit=1, + # edit_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", + # ), +] + +ONE_GPU_CASES_B: list[DiffusionTestCase] = [ + # === Text to Video (T2V) === + DiffusionTestCase( + "wan2_1_t2v_1.3b", + DiffusionServerArgs( + model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + modality="video", + warmup_text=0, + warmup_edit=0, + custom_validator="video", + ), + DiffusionSamplingParams( + prompt="A curious raccoon", + output_size="848x480", + ), + ), + # NOTE(mick): flaky + # DiffusionTestCase( + # id="hunyuan_video", + # model_path="hunyuanvideo-community/HunyuanVideo", + # modality="video", + # prompt="A curious raccoon", + # output_size="720x480", + # warmup_text=0, + # warmup_edit=0, + # custom_validator="video", + # ), + DiffusionTestCase( + "flux_2_ti2i", + DiffusionServerArgs( + model_path="black-forest-labs/FLUX.2-dev", + modality="image", + warmup_text=0, + warmup_edit=1, + ), + 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", + ), + ), + DiffusionTestCase( + "fast_hunyuan_video", + DiffusionServerArgs( + model_path="FastVideo/FastHunyuan-diffusers", + modality="video", + warmup_text=0, + warmup_edit=0, + custom_validator="video", + ), + DiffusionSamplingParams( + prompt="A curious raccoon", + output_size="720x480", + ), + ), + # === Text and Image to Video (TI2V) === + DiffusionTestCase( + "wan2_2_ti2v_5b", + DiffusionServerArgs( + model_path="Wan-AI/Wan2.2-TI2V-5B-Diffusers", + modality="video", + warmup_text=0, + warmup_edit=0, + custom_validator="video", + ), + DiffusionSamplingParams( + output_size="832x1104", + prompt="Add dynamic motion to the scene", + image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg", + ), + ), + DiffusionTestCase( + "fastwan2_2_ti2v_5b", + DiffusionServerArgs( + model_path="FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", + modality="video", + warmup_text=0, + warmup_edit=0, + custom_validator="video", + ), + DiffusionSamplingParams( + output_size="832x1104", + prompt="Add dynamic motion to the scene", + image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg", + ), + ), +] + +TWO_GPU_CASES_A = [ + # TODO: Timeout with Torch2.9. Add back when it can pass CI + # DiffusionTestCase( + # id="wan2_2_i2v_a14b_2gpu", + # model_path="Wan-AI/Wan2.2-I2V-A14B-Diffusers", + # modality="video", + # prompt="generate", + # warmup_text=0, + # warmup_edit=0, + # output_size="832x1104", + # edit_prompt="generate", + # image_path="https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true", + # custom_validator="video", + # num_gpus=2, + # num_frames=1, + # ), + DiffusionTestCase( + "wan2_2_t2v_a14b_2gpu", + DiffusionServerArgs( + model_path="Wan-AI/Wan2.2-T2V-A14B-Diffusers", + modality="video", + warmup_text=0, + warmup_edit=0, + custom_validator="video", + num_gpus=2, + ), + DiffusionSamplingParams( + prompt="A curious raccoon", + output_size="720x480", + ), + ), + DiffusionTestCase( + "wan2_1_t2v_14b_2gpu", + DiffusionServerArgs( + model_path="Wan-AI/Wan2.1-T2V-14B-Diffusers", + warmup_text=0, + warmup_edit=0, + modality="video", + num_gpus=2, + custom_validator="video", + ), + DiffusionSamplingParams( + prompt="A curious raccoon", + output_size="720x480", + ), + ), +] + +TWO_GPU_CASES_B = [ + DiffusionTestCase( + "wan2_1_i2v_14b_480P_2gpu", + DiffusionServerArgs( + model_path="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", + warmup_text=0, + warmup_edit=0, + modality="video", + custom_validator="video", + num_gpus=2, + ), + DiffusionSamplingParams( + output_size="832x1104", + prompt="Add dynamic motion to the scene", + image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg", + ), + ), + DiffusionTestCase( + "wan2_1_i2v_14b_720P_2gpu", + DiffusionServerArgs( + model_path="Wan-AI/Wan2.1-I2V-14B-720P-Diffusers", + modality="video", + warmup_text=0, + warmup_edit=0, + custom_validator="video", + num_gpus=2, + ), + DiffusionSamplingParams( + prompt="Add dynamic motion to the scene", + image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg", + output_size="832x1104", + ), + ), + DiffusionTestCase( + "qwen_image_t2i_2_gpus", + DiffusionServerArgs( + model_path="Qwen/Qwen-Image", + modality="image", + warmup_text=1, + warmup_edit=0, + num_gpus=2, + ), + DiffusionSamplingParams( + prompt="A futuristic cityscape at sunset with flying cars", + output_size="1024x1024", + ), + ), + DiffusionTestCase( + "flux_image_t2i_2_gpus", + DiffusionServerArgs( + model_path="black-forest-labs/FLUX.1-dev", + modality="image", + warmup_text=1, + warmup_edit=0, + ), + DiffusionSamplingParams( + prompt="A futuristic cityscape at sunset with flying cars", + output_size="1024x1024", + ), + ), +] + +# Load global configuration +BASELINE_CONFIG = BaselineConfig.load(Path(__file__).with_name("perf_baselines.json")) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/slack_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/slack_utils.py new file mode 100644 index 00000000..7237c891 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/slack_utils.py @@ -0,0 +1,193 @@ +""" + 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 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: str = None, +) -> bool: + temp_path = None + 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 + + if origin_file_path and origin_file_path.startswith(("http", "https")): + suffix = os.path.splitext(urlparse(origin_file_path).path)[1] or ".tmp" + with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tf: + with urlopen(origin_file_path) as response: + tf.write(response.read()) + temp_path = tf.name + origin_file_path = temp_path + + uploads = [{"file": file_path, "title": "Generated Image"}] + if origin_file_path and os.path.exists(origin_file_path): + uploads.insert(0, {"file": origin_file_path, "title": "Original 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: + if temp_path and os.path.exists(temp_path): + os.remove(temp_path) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_files/launch_flux.json b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_files/launch_flux.json new file mode 100644 index 00000000..6a9d8382 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_files/launch_flux.json @@ -0,0 +1,11 @@ +{ + "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" +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_files/launch_wan.json b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_files/launch_wan.json new file mode 100644 index 00000000..eeb9ddf9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_files/launch_wan.json @@ -0,0 +1,11 @@ +{ + "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" +} diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_offline_api.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_offline_api.py new file mode 100644 index 00000000..9c45e271 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_offline_api.py @@ -0,0 +1,75 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +""" + Testing the performance of generate command of sgl_diffusion' CLI +""" + +import unittest + +import torch + +from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class TestGeneratorAPIBase(unittest.TestCase): + # server args + server_kwargs = {} + + # sampling + output_path: str = "test_outputs" + + results = [] + + @classmethod + def setUpClass(cls): + cls.results = [] + + def verify_single_generation_result(self, result): + self.assertIsNotNone(result, "Generation failed") + self.assertTrue( + "samples" in result and isinstance(result["samples"], torch.Tensor), + f"Incorrect Generation result", + ) + + def _run_test(self, name, server_kwargs, test_key: str): + generator = DiffGenerator.from_pretrained(**server_kwargs) + result = generator.generate(prompt="A curious raccoon") + self.verify_single_generation_result(result) + + def test_single_gpu(self): + self._run_test( + name=self.server_kwargs["model_path"], + server_kwargs=self.server_kwargs | dict(num_gpus=1), + test_key="test_single_gpu", + ) + + def test_cfg_parallel(self): + self._run_test( + name=self.server_kwargs["model_path"], + server_kwargs=self.server_kwargs + | dict(num_gpus=2, enable_cfg_parallel=True), + test_key="test_cfg_parallel", + ) + + def test_multiple_prompts(self): + generator = DiffGenerator.from_pretrained( + **self.server_kwargs | dict(num_gpus=2, enable_cfg_parallel=True) + ) + prompts = ["A curious raccoon", "A curious cat"] + results = generator.generate(prompt=prompts) + + self.assertEqual(len(results), len(prompts), "Some generation tasks fail") + for result in results: + self.verify_single_generation_result(result) + + +class TestWan2_1_T2V(TestGeneratorAPIBase): + server_kwargs = {"model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"} + + +if __name__ == "__main__": + del TestGeneratorAPIBase + unittest.main() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_utils.py new file mode 100644 index 00000000..7af786c4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/test/test_utils.py @@ -0,0 +1,420 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +import base64 +import dataclasses +import json +import os +import shlex +import socket +import subprocess +import sys +import time +import unittest +from pathlib import Path +from typing import Optional + +from PIL import Image + +from sglang.multimodal_gen.configs.sample.sampling_params import DataType +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__) + + +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 run_command(command) -> Optional[float]: + """Runs a command and returns the execution time and status.""" + print(f"Running command: {shlex.join(command)}") + + duration = None + 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) + if "Pixel data generated" in line: + words = line.split(" ") + duration = float(words[-2]) + + if process.returncode == 0: + return duration + else: + print(f"Command failed with exit code {process.returncode}") + return None + + +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 is_mp4(data): + idx = data.find(b"ftyp") + return 0 <= idx <= 32 + + +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 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, + prev_len: int, + log_path: Path, + timeout: float = 30.0, +) -> tuple[RequestPerfRecord | None, int]: + """ + 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) + if len(records) >= prev_len + 1: + # FIXME: unable to get rid from openai apis, this is a hack. we should compare rid + # potential error when there are multiple servers + return records[-1], len(records) + + time.sleep(0.5) + + if os.environ.get("SGLANG_GEN_BASELINE", "0") == "1": + records = read_perf_logs(log_path) + return None, len(records) + + 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_mp4 = ( + video_bytes[:4] == b"\x00\x00\x00\x18" or video_bytes[:4] == b"\x00\x00\x00\x1c" + ) + is_webm = video_bytes[:4] == b"\x1a\x45\xdf\xa3" + assert is_mp4 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_mp4 = ( + video_bytes.startswith(b"\x00\x00\x00\x18") + or video_bytes.startswith(b"\x00\x00\x00\x1c") + or video_bytes[4:8] == b"ftyp" + ) + is_webm = video_bytes.startswith(b"\x1a\x45\xdf\xa3") + assert is_mp4 or is_webm, "Video must be MP4 or WebM" + + +@dataclasses.dataclass +class TestResult: + name: str + key: str + duration: Optional[float] + succeed: bool + + @property + def duration_str(self): + return f"{self.duration:.4f}" if self.duration else "NA" + + +class TestCLIBase(unittest.TestCase): + model_path: str = None + extra_args = [] + data_type: DataType = None + # tested on h100 + thresholds = {} + + width: int = 720 + height: int = 720 + output_path: str = "test_outputs" + + base_command = [ + "sglang", + "generate", + "--text-encoder-cpu-offload", + "--pin-cpu-memory", + "--prompt", + "A curious raccoon", + "--save-output", + "--log-level=debug", + f"--width={width}", + f"--height={height}", + f"--output-path={output_path}", + ] + + results = [] + + @classmethod + def setUpClass(cls): + cls.results = [] + + def _run_command(self, name: str, model_path: str, test_key: str = "", args=[]): + command = ( + self.base_command + + [f"--model-path={model_path}"] + + shlex.split(args or "") + + ["--output-file-name", f"{name}"] + + self.extra_args + ) + duration = run_command(command) + status = "Success" if duration else "Failed" + succeed = duration is not None + + duration = float(duration) if succeed else None + self.results.append(TestResult(name, test_key, duration, succeed)) + + return name, duration, status + + +class TestGenerateBase(TestCLIBase): + model_path: str = None + extra_args = [] + data_type: DataType = None + # tested on h100 + thresholds = {} + + width: int = 720 + height: int = 720 + output_path: str = "test_outputs" + image_path: str | None = None + prompt: str | None = "A curious raccoon" + + base_command = [ + "sglang", + "generate", + # "--text-encoder-cpu-offload", + # "--pin-cpu-memory", + f"--prompt", + f"{prompt}", + "--save-output", + "--log-level=debug", + f"--width={width}", + f"--height={height}", + f"--output-path={output_path}", + ] + + results: list[TestResult] = [] + + @classmethod + def setUpClass(cls): + cls.results = [] + + @classmethod + def tearDownClass(cls): + # Print markdown table + print("\n## Test Results\n") + print("| Test Case | Duration | Status |") + print("|--------------------------------|----------|---------|") + test_keys = ["test_single_gpu", "test_cfg_parallel", "test_usp", "test_mixed"] + test_key_to_order = { + test_key: order for order, test_key in enumerate(test_keys) + } + + ordered_results: list[TestResult] = [None] * len(test_keys) + for result in cls.results: + order = test_key_to_order[result.key] + ordered_results[order] = result + + for result in ordered_results: + if not result: + continue + status = ( + "Succeed" + if ( + result.succeed + and float(result.duration) <= float(cls.thresholds[result.key]) + ) + else "Failed" + ) + print(f"| {result.name:<30} | {result.duration_str:<8} | {status:<7} |") + print() + durations = [result.duration_str for result in cls.results] + print(" | ".join([""] + durations + [""])) + + def _run_test(self, name: str, args, model_path: str, test_key: str): + time_threshold = self.thresholds[test_key] + name, duration, status = self._run_command( + name, args=args, model_path=model_path, test_key=test_key + ) + self.verify(status, name, duration, time_threshold) + + def verify(self, status, name, duration, time_threshold): + print("-" * 80) + print("\n" * 3) + + # test task status + self.assertEqual(status, "Success", f"{name} command failed") + self.assertIsNotNone(duration, f"Could not parse duration for {name}") + self.assertLessEqual( + duration, + time_threshold, + f"{name} failed with {duration:.4f}s > {time_threshold}s", + ) + + # 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) + logger.info(f"{name} passed in {duration:.4f}s (threshold: {time_threshold}s)") + + 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", + ) + + def test_cfg_parallel(self): + """cfg parallel""" + self._run_test( + name=f"{self.model_name()}_cfg_parallel", + args="--num-gpus 2 --enable-cfg-parallel", + model_path=self.model_path, + test_key="test_cfg_parallel", + ) + + def test_usp(self): + """usp""" + self._run_test( + name=f"{self.model_name()}_usp", + args="--num-gpus 4 --ulysses-degree=2 --ring-degree=2", + model_path=self.model_path, + test_key="test_usp", + ) + + def test_mixed(self): + """mixed""" + self._run_test( + name=f"{self.model_name()}_mixed", + args="--num-gpus 4 --ulysses-degree=2 --ring-degree=1 --enable-cfg-parallel", + model_path=self.model_path, + test_key="test_mixed", + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/third_party/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/third_party/__init__.py new file mode 100644 index 00000000..af2eb7d1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/third_party/__init__.py @@ -0,0 +1 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/third_party/pynvml.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/third_party/pynvml.py new file mode 100644 index 00000000..546dc8b8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/third_party/pynvml.py @@ -0,0 +1,7227 @@ +# 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 == 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 == 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 == None: + # lock to ensure only one caller loads the library + libLoadLock.acquire() + + try: + # ensure the library still isn't loaded + if nvmlLib == 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 == 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/utils.py new file mode 100644 index 00000000..6a8f5958 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/multimodal_gen/utils.py @@ -0,0 +1,813 @@ +# 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 socket +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 imageio +import numpy as np +import torch +import torchvision +import yaml +from einops import rearrange +from remote_pdb import RemotePdb +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") + +# 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` or `librccl.so.1` 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" + else: + raise ValueError("NCCL only supports CUDA and ROCm 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}") + + +# For non-torch.distributed debugging +def remote_breakpoint() -> None: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("localhost", 0)) # Let the OS pick an ephemeral port. + port = s.getsockname()[1] + RemotePdb(host="localhost", port=port).set_trace() + + +@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. + + Returns: + torch.dtype: The compute dtype to use, defaults to get_default_dtype() if no policy set + """ + 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 save_decoded_latents_as_video( + decoded_latents: list[torch.Tensor], output_path: str, fps: int +): + # Process outputs + videos = rearrange(decoded_latents, "b c t h w -> t b c h w") + frames = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + + os.makedirs(os.path.dirname(output_path), exist_ok=True) + imageio.mimsave(output_path, frames, fps=fps, format="mp4") + + +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/profiler.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/profiler.py new file mode 100644 index 00000000..0ffa07f2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/profiler.py @@ -0,0 +1,155 @@ +""" +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, +) -> 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, + } + + 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 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, + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/_custom_ops.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/_custom_ops.py new file mode 100644 index 00000000..3353aa2e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/_custom_ops.py @@ -0,0 +1,165 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/_custom_ops.py +import logging +from typing import List, Optional, Tuple + +import torch + +from sglang.srt.utils import is_hip, is_hpu, is_npu + +logger = logging.getLogger(__name__) + + +if not is_hpu(): + try: + import sgl_kernel + except ImportError as e: + logger.warning("Failed to import from custom_ar with %r", e) + + +if not is_hip() and not is_npu(): + custom_op = sgl_kernel.allreduce + + # custom allreduce + def init_custom_ar( + ipc_tensors: List[torch.Tensor], + rank_data: torch.Tensor, + rank: int, + full_nvlink: bool, + ) -> int: + return custom_op.init_custom_ar(ipc_tensors, rank_data, rank, full_nvlink) + + def all_reduce( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, + ) -> None: + custom_op.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) + + def dispose(fa: int) -> None: + custom_op.dispose(fa) + + def meta_size() -> int: + return custom_op.meta_size() + + def register_buffer(fa: int, ipc_tensors: List[int]) -> None: + return custom_op.register_buffer(fa, ipc_tensors) + + def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]: + return custom_op.get_graph_buffer_ipc_meta(fa) + + def register_graph_buffers( + fa: int, handles: List[List[int]], offsets: List[List[int]] + ) -> None: + custom_op.register_graph_buffers(fa, handles, offsets) + +else: + # ROCM custom allreduce + + def init_custom_ar( + meta: torch.Tensor, + rank_data: torch.Tensor, + handles: List[str], + offsets: List[int], + rank: int, + full_nvlink: bool, + ) -> int: + return sgl_kernel.allreduce.init_custom_ar( + meta, rank_data, handles, offsets, rank, full_nvlink + ) + + def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None: + sgl_kernel.allreduce.all_reduce_reg(fa, inp, out) + + def all_reduce_unreg( + fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor + ) -> None: + sgl_kernel.allreduce.all_reduce_unreg(fa, inp, reg_buffer, out) + + def dispose(fa: int) -> None: + sgl_kernel.allreduce.dispose(fa) + + def meta_size() -> int: + return sgl_kernel.allreduce.meta_size() + + def register_buffer( + fa: int, t: torch.Tensor, handles: List[str], offsets: List[int] + ) -> None: + return sgl_kernel.allreduce.register_buffer(fa, t, handles, offsets) + + def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]: + return sgl_kernel.allreduce.get_graph_buffer_ipc_meta(fa) + + def register_graph_buffers( + fa: int, handles: List[str], offsets: List[List[int]] + ) -> None: + sgl_kernel.allreduce.register_graph_buffers(fa, handles, offsets) + + def allocate_meta_buffer(size: int) -> torch.Tensor: + return sgl_kernel.allreduce.allocate_meta_buffer(size) + + def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor: + return sgl_kernel.allreduce.get_meta_buffer_ipc_handle(inp) + + # ROCM custom quick allreduce + + def init_custom_qr( + rank: int, world_size: int, qr_max_size: Optional[int] = None + ) -> int: + return sgl_kernel.allreduce.init_custom_qr(world_size, rank, qr_max_size) + + def qr_get_handle(fa: int) -> torch.Tensor: + return sgl_kernel.allreduce.qr_get_handle(fa) + + def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None: + sgl_kernel.allreduce.qr_open_handles(fa, handles) + + def qr_all_reduce( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + quant_level: int, + cast_bf2half: bool, + ) -> None: + sgl_kernel.allreduce.qr_all_reduce(fa, inp, out, quant_level, cast_bf2half) + + def qr_destroy(fa: int) -> None: + sgl_kernel.allreduce.qr_destroy(fa) + + def qr_max_size() -> int: + return sgl_kernel.allreduce.qr_max_size() + + +def mscclpp_generate_unique_id() -> bytes: + return sgl_kernel.allreduce.mscclpp_generate_unique_id() + + +def mscclpp_init_context( + unique_id: bytes, + rank: int, + world_size: int, + scratch: torch.Tensor, + put_buffer: torch.Tensor, + nranks_per_node: int, + rank_to_node: List[int], + rank_to_ib: List[int], + context_selection: int, +) -> int: + return sgl_kernel.allreduce.mscclpp_init_context( + unique_id, + rank, + world_size, + scratch, + put_buffer, + nranks_per_node, + rank_to_node, + rank_to_ib, + context_selection, + ) + + +def mscclpp_allreduce( + context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int +) -> None: + return sgl_kernel.allreduce.mscclpp_allreduce(context, inp, out, nthreads, nblocks) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_invariant_ops/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_invariant_ops/__init__.py new file mode 100644 index 00000000..3caa60f0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_invariant_ops/__init__.py @@ -0,0 +1,29 @@ +# 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", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py new file mode 100644 index 00000000..67bb1c5c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py @@ -0,0 +1,991 @@ +# 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: + return torch.einsum("ik,kj->ij", a, b) + + return _matmul_persistent_triton(a=a, b=b, bias=bias) + + +@triton.jit +def _log_softmax_kernel( + input_ptr, + output_ptr, + input_row_stride, + output_row_stride, + n_cols, + 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, + output_row_stride, + n_cols, + 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/operations.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/operations.py new file mode 100644 index 00000000..9d824587 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/operations.py @@ -0,0 +1,211 @@ +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.input_ids.shape[0] + 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] + + if self._global_dp_buffer_len is not None: + 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 clearning. This may indicate you do not release memory early enough but leave it to 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, + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/operations_strategy.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/operations_strategy.py new file mode 100644 index 00000000..152e4874 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/operations_strategy.py @@ -0,0 +1,211 @@ +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 + + +@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 + ] + ) + 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 = 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 = 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(), + ], + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/single_batch_overlap.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/single_batch_overlap.py new file mode 100644 index 00000000..7b324c37 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/single_batch_overlap.py @@ -0,0 +1,116 @@ +# 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.layers.moe import get_moe_runner_backend +from sglang.srt.layers.moe.utils import is_sbo_enabled +from sglang.srt.utils import get_int_env_var + + +class SboFlags: + # TODO may have: "enable_dispatch_shared_one_stream_overlap", "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() + ) + + @classmethod + def enable_combine_shared_two_stream_overlap(cls): + return is_sbo_enabled() + + @classmethod + def fuse_shared_experts_inside_sbo(cls): + # TODO after antgroup's PR, should be `... or cls.enable_dispatch_shared_one_stream_overlap()` + return cls.enable_combine_shared_two_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: int + signal: Optional[torch.Tensor] = None + threshold: 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 + communicate_num_sms = get_int_env_var("SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS", 32) + 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 + combine_signal = torch.zeros( + num_local_experts, dtype=torch.uint32, 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/two_batch_overlap.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/two_batch_overlap.py new file mode 100644 index 00000000..becb54e4 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -0,0 +1,1027 @@ +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.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, +) +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): + local_can_run_tbo_aggregated = min(partial_global_info[:, 0].tolist()) + forward_modes = partial_global_info[:, 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 = [ + x for x in forward_modes if x != ForwardMode.IDLE.value + ] + + if not forward_modes_excluding_idle: + return ForwardMode.IDLE, False + + forward_mode_agree = TboDPAttentionPreparer._is_all_same( + forward_modes_excluding_idle + ) + global_forward_mode = ( + ForwardMode(forward_modes_excluding_idle[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] + + 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", + ]: + 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", + "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, + 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, + ) + ) + + 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, + ) + 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) + + +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) + return dict( + hidden_states=hidden_states[token_slice], + residual=None if residual is None else residual[token_slice], + positions=positions[token_slice], + forward_batch=output_forward_batch, + tbo_subbatch_index=tbo_subbatch_index, + ) + + +def _model_forward_tbo_merge_outputs(output_a, output_b): + 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 + return torch.concat([value_a, value_b], dim=0) + + 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) + ] + + 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 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/checkpoint_engine/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/checkpoint_engine/__init__.py new file mode 100644 index 00000000..d8a77f90 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/checkpoint_engine/__init__.py @@ -0,0 +1,9 @@ +""" +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"] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py new file mode 100644 index 00000000..dd8805e6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py @@ -0,0 +1,142 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/checkpoint_engine/update.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/checkpoint_engine/update.py new file mode 100644 index 00000000..93c8b4b6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/checkpoint_engine/update.py @@ -0,0 +1,317 @@ +""" +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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/backend.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/backend.py new file mode 100644 index 00000000..9892e472 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/backend.py @@ -0,0 +1,455 @@ +# 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.pass_manager import PostGradPassManager +from sglang.srt.utils.common import rank0_log + +logger = logging.getLogger(__name__) + + +SPLIT_OPS = [ + "sglang.unified_attention_with_output", + "sglang.gdn_with_output", +] + + +def add_split_ops(ops): + SPLIT_OPS.extend(ops) + + +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}") + + +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] = CUDAPiecewiseBackend( + 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, + ): + rank0_log(f"Initializing SGLangBackend") + 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: + rank0_log(f"SGLangBackend __call__") + 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, + 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) + + rank0_log(f"Computation graph saved to {graph_path}") + + self._called = True + return self.split_gm diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compilation_config.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compilation_config.py new file mode 100644 index 00000000..5ddafe8a --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compilation_config.py @@ -0,0 +1,29 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/compilation_config.py + +from typing import List + + +# 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 + + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compilation_counter.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compilation_counter.py new file mode 100644 index 00000000..e973f8f2 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compilation_counter.py @@ -0,0 +1,47 @@ +# 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() diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compile.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compile.py new file mode 100644 index 00000000..b9ff7f6b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compile.py @@ -0,0 +1,215 @@ +import contextvars +import inspect +import logging +import os +import sys +import types +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Callable, Optional, Union + +import torch + +from sglang.srt.compilation.compilation_config import CompilationConfig +from sglang.srt.utils.common import rank0_log + +logger = logging.getLogger(__name__) + +_COMPILE_ENABLED = contextvars.ContextVar("_COMPILE_ENABLED", default=False) + + +@contextmanager +def set_compiled(enabled: bool = True): + token = _COMPILE_ENABLED.set(enabled) + try: + yield + finally: + _COMPILE_ENABLED.reset(token) + + +@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.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.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, +): + rank0_log(f"install_torch_compiled") + 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 = _COMPILE_ENABLED.get() + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compiler_interface.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compiler_interface.py new file mode 100644 index 00000000..8310f75c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/compiler_interface.py @@ -0,0 +1,503 @@ +# 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 + + +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.__version__.startswith("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.__version__ >= "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.__version__.startswith("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.__version__ >= "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") diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/cuda_piecewise_backend.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/cuda_piecewise_backend.py new file mode 100644 index 00000000..2e45d34d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/cuda_piecewise_backend.py @@ -0,0 +1,215 @@ +# 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, Union +from unittest.mock import patch + +import torch +import torch.fx as fx +from sgl_kernel import weak_ref_tensor + +from sglang.srt.compilation.compilation_config import CompilationConfig +from sglang.srt.compilation.compilation_counter import compilation_counter + +logger = logging.getLogger(__name__) + + +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") + + +@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) + 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() + + # Skip CUDA graphs if this entry doesn't use them OR + # if we're supposed to skip them globally + # skip_cuda_graphs = get_forward_context().skip_cuda_graphs + # if not entry.use_cudagraph or skip_cuda_graphs: + # 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. + with torch.cuda.graph(cudagraph, 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 = 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/fix_functionalization.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/fix_functionalization.py new file mode 100644 index 00000000..8673e357 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/fix_functionalization.py @@ -0,0 +1,134 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/fx_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/fx_utils.py new file mode 100644 index 00000000..b2e863e6 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/fx_utils.py @@ -0,0 +1,83 @@ +# 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)) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/inductor_pass.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/inductor_pass.py new file mode 100644 index 00000000..acbde65b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/inductor_pass.py @@ -0,0 +1,140 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/pass_manager.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/pass_manager.py new file mode 100644 index 00000000..9173976f --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/pass_manager.py @@ -0,0 +1,66 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/piecewise_context_manager.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/piecewise_context_manager.py new file mode 100644 index 00000000..a0d3c7dd --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/compilation/piecewise_context_manager.py @@ -0,0 +1,61 @@ +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, List, Optional + +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + +_in_piecewise_cuda_graph = False + + +def is_in_piecewise_cuda_graph(): + return _in_piecewise_cuda_graph + + +@contextmanager +def enable_piecewise_cuda_graph(): + global _in_piecewise_cuda_graph + _in_piecewise_cuda_graph = True + + yield + + _in_piecewise_cuda_graph = False + + +@dataclass +class ForwardContext: + def __init__(self): + self.forward_batch = None + self.attention_layer = 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 + + +_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 +): + 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) + try: + yield + finally: + _forward_context = None diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/__init__.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/__init__.py new file mode 100644 index 00000000..b35cc1dc --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/__init__.py @@ -0,0 +1,47 @@ +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.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_linear import KimiLinearConfig +from sglang.srt.configs.kimi_vl import KimiVLConfig +from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig +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_next import Qwen3NextConfig +from sglang.srt.configs.step3_vl import ( + Step3TextConfig, + Step3VisionEncoderConfig, + Step3VLConfig, +) + +__all__ = [ + "ExaoneConfig", + "ChatGLMConfig", + "DbrxConfig", + "DeepseekVL2Config", + "LongcatFlashConfig", + "MultiModalityConfig", + "KimiVLConfig", + "MoonViTConfig", + "Step3VLConfig", + "Step3TextConfig", + "Step3VisionEncoderConfig", + "Olmo3Config", + "KimiLinearConfig", + "Qwen3NextConfig", + "DotsVLMConfig", + "DotsOCRConfig", + "FalconH1Config", + "NemotronHConfig", + "NemotronH_Nano_VL_V2_Config", + "JetNemotronConfig", + "JetVLMConfig", +] diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/chatglm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/chatglm.py new file mode 100644 index 00000000..9370c218 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/chatglm.py @@ -0,0 +1,78 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/dbrx.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/dbrx.py new file mode 100644 index 00000000..75ccbde9 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/dbrx.py @@ -0,0 +1,279 @@ +# 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, + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/deepseek_ocr.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/deepseek_ocr.py new file mode 100644 index 00000000..b1f2488d --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/deepseek_ocr.py @@ -0,0 +1,794 @@ +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, + **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 + + 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) + + 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, + ) + + 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 + ) + + 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 + projector_config: MlpProjectorConfig + + 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/deepseekvl2.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/deepseekvl2.py new file mode 100644 index 00000000..9621f058 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/deepseekvl2.py @@ -0,0 +1,687 @@ +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 + projector_config: DeepseekVL2MlpProjectorConfig + language_config: DeepseekV2Config + + 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/device_config.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/device_config.py new file mode 100644 index 00000000..20b9af9b --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/device_config.py @@ -0,0 +1,19 @@ +import logging +from typing import Optional + +import torch + +logger = logging.getLogger(__name__) + + +class DeviceConfig: + device: Optional[torch.device] + gpu_id: Optional[int] + + def __init__(self, device: str = "cuda", gpu_id: int = -1) -> None: + if device in ["cuda", "xpu", "hpu", "cpu", "npu"]: + 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/dots_ocr.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/dots_ocr.py new file mode 100644 index 00000000..8b0693b8 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/dots_ocr.py @@ -0,0 +1,64 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/dots_vlm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/dots_vlm.py new file mode 100644 index 00000000..dc921582 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/dots_vlm.py @@ -0,0 +1,134 @@ +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/exaone.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/exaone.py new file mode 100644 index 00000000..7b0a2d29 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/exaone.py @@ -0,0 +1,195 @@ +# 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 + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/falcon_h1.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/falcon_h1.py new file mode 100644 index 00000000..1f524b89 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/falcon_h1.py @@ -0,0 +1,310 @@ +# 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 + +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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/internvl.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/internvl.py new file mode 100644 index 00000000..3ba9c61c --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/internvl.py @@ -0,0 +1,706 @@ +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.clean_up_tokenization(out_string) + 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 +) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/janus_pro.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/janus_pro.py new file mode 100644 index 00000000..d574953e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/janus_pro.py @@ -0,0 +1,634 @@ +# 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 + aligner_config: AlignerConfig + + gen_vision_config: GenVisionConfig + gen_aligner_config: GenAlignerConfig + gen_head_config: GenHeadConfig + + language_config: LlamaConfig + + 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 + min_size: int + image_mean: Union[Tuple[float, float, float], List[float]] + image_std: Union[Tuple[float, float, float], List[float]] + rescale_factor: float + do_normalize: bool + + 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/jet_nemotron.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/jet_nemotron.py new file mode 100644 index 00000000..c05a2ec1 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/jet_nemotron.py @@ -0,0 +1,74 @@ +from dataclasses import dataclass +from typing import Any + +from transformers.configuration_utils import PretrainedConfig + +from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape + + +@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]] + hidden_act: str + hidden_size: int + initializer_range: float + intermediate_size: int + layer_types: list[str] + max_position_embeddings: int + num_attention_heads: int + num_key_value_heads: int + rms_norm_eps: float + rope_scaling: None + rope_theta: float + + @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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/jet_vlm.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/jet_vlm.py new file mode 100644 index 00000000..9b8cba6e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/jet_vlm.py @@ -0,0 +1,53 @@ +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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/kimi_linear.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/kimi_linear.py new file mode 100644 index 00000000..e7360904 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/kimi_linear.py @@ -0,0 +1,161 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/kimi_vl.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/kimi_vl.py new file mode 100644 index 00000000..3c7d20f5 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/kimi_vl.py @@ -0,0 +1,38 @@ +# 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) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/kimi_vl_moonvit.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/kimi_vl_moonvit.py new file mode 100644 index 00000000..166809eb --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/kimi_vl_moonvit.py @@ -0,0 +1,32 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/load_config.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/load_config.py new file mode 100644 index 00000000..042eb322 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/load_config.py @@ -0,0 +1,121 @@ +# 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 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" + JAX = "jax" + REMOTE = "remote" + REMOTE_INSTANCE = "remote_instance" + RDMA = "rdma" + LOCAL_CACHED = "local_cached" + + +@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. + 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 + + # 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 + + 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}" + ) diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/longcat_flash.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/longcat_flash.py new file mode 100644 index 00000000..e6a2dfb0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/longcat_flash.py @@ -0,0 +1,104 @@ +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, + **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" diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/mamba_utils.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/mamba_utils.py new file mode 100644 index 00000000..d2ff3762 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/mamba_utils.py @@ -0,0 +1,183 @@ +# 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, etc.""" + +import os +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 + + +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 + + +CONV_DTYPE = torch.bfloat16 + + +def mamba2_state_dtype() -> Mamba2StateDType: + dtype_map = { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + } + ssm_dtype = dtype_map[os.environ["SGLANG_MAMBA_SSM_DTYPE"]] + return Mamba2StateDType(conv=CONV_DTYPE, temporal=ssm_dtype) + + +@dataclass(kw_only=True, frozen=True) +class BaseLinearStateParams(ABC): + dtype: Mamba2StateDType = field(default_factory=mamba2_state_dtype) + 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 = conv_state_k_shape[1], conv_state_k_shape[0] + + return KimiLinearStateShape( + conv=[conv_state_shape, conv_state_k_shape, conv_state_k_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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/model_config.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/model_config.py new file mode 100644 index 00000000..26dfbe5e --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/model_config.py @@ -0,0 +1,1043 @@ +# 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 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, 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: PretrainedConfig) -> bool: + return ( + config.architectures is not None + and config.architectures[0] + in [ + "DeepseekV3ForCausalLM", + "DeepseekV32ForCausalLM", + "DeepseekV3ForCausalLMNextN", + ] + and getattr(config, "index_topk", None) 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, + hybrid_kvcache_ratio: Optional[ + float + ] = None, # TODO: remove this, it is not a model config + model_impl: Union[str, ModelImpl] = ModelImpl.AUTO, + sampling_defaults: str = "openai", + quantize_and_serve: 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 + + # 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.is_hybrid_swa = is_hybrid_model( + self.hf_config.architectures, + hybrid_kvcache_ratio=hybrid_kvcache_ratio, + context_length=context_length, + attention_chunk_size=self.attention_chunk_size, + ) + if self.is_hybrid_swa is not None: + self.swa_attention_layer_ids, self.full_attention_layer_ids = ( + get_hybrid_layer_ids( + self.hf_config.architectures, self.hf_text_config.num_hidden_layers + ) + ) + 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.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() + + # Verify quantization + self._verify_quantization() + + # 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) + + # 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, + **kwargs, + ): + 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=server_args.quantization, + hybrid_kvcache_ratio=server_args.hybrid_kvcache_ratio, + model_impl=server_args.model_impl, + sampling_defaults=server_args.sampling_defaults, + quantize_and_serve=server_args.quantize_and_serve, + override_config_file=server_args.decrypted_config_file, + **kwargs, + ) + + def _config_draft_model(self): + is_draft_model = self.is_draft_model + + if ( + is_draft_model + and self.hf_config.architectures[0] == "DeepseekV3ForCausalLM" + ): + self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN" + + if is_draft_model and self.hf_config.architectures[0] == "Glm4MoeForCausalLM": + self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN" + + 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] in [ + "BailingMoeV2ForCausalLM", + "BailingMoeForCausalLM", + ]: + 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 + + 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, + ) + + # 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 "LongcatFlashForCausalLM" in self.hf_config.architectures + or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures + or "DotsVLMForCausalLM" in self.hf_config.architectures + ): + self.head_dim = 256 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_config.kv_lora_rank + self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim + self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim + self.v_head_dim = self.hf_config.v_head_dim + self.index_head_dim = ( + get_nsa_index_head_dim(self.hf_config) + if is_deepseek_nsa(self.hf_config) + else None + ) + + # 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: + mscale_all_dim = self.hf_config.rope_scaling.get( + "mscale_all_dim", False + ) + scaling_factor = self.hf_config.rope_scaling["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + self.scaling = self.scaling * mscale * mscale + + 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 + 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 + 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) + + 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 + 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) + + # 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 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: + import huggingface_hub + + try: + from huggingface_hub import HfApi, hf_hub_download + + hf_api = HfApi() + # Retry HF API call up to 3 times + 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 file_exists: + # Download and parse the quantization config for remote models + quant_config_file = hf_hub_download( + repo_id=self.model_path, + filename="hf_quant_config.json", + revision=self.revision, + ) + 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.OfflineModeIsEnabled: + logger.warning( + "Offline mode is enabled, skipping hf_quant_config.json check" + ) + except Exception as e: + logger.warning( + f"Failed to check hf_quant_config.json: {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 _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": + return {"quant_method": "w4afp8"} + elif quant_algo and ("FP4" in quant_algo or "NVFP4" in quant_algo): + return {"quant_method": "modelopt_fp4"} + elif quant_algo and "FP8" in quant_algo: + return {"quant_method": "modelopt_fp8"} + else: + return None + + def _is_already_quantized(self) -> bool: + """Check if the model is already quantized based on config files.""" + # 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": + # 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 _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_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", + ] + optimized_quantization_methods = [ + "fp8", + "marlin", + "modelopt_fp8", + "modelopt_fp4", + "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", + ] + compatible_quantization_methods = { + "modelopt_fp8": ["modelopt"], + "modelopt_fp4": ["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 model config, if available. + quant_cfg = self._parse_quant_hf_config() + + 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: + if ( + self.quantization not in compatible_quantization_methods + or quant_method + not in compatible_quantization_methods[self.quantization] + ): + 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})." + ) + + 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: + 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 _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. + config_dtype = getattr(config, "dtype", None) + 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 config.model_type.startswith("gemma"): + if config.model_type == "gemma": + gemma_version = "" + else: + gemma_version = config.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 "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 + ): + return False + else: + return not is_embedding + + +multimodal_model_archs = [ + "CLIPModel", + "DeepseekVL2ForCausalLM", + "Gemma3ForConditionalGeneration", + "Gemma3nForConditionalGeneration", + "Glm4vForConditionalGeneration", + "Glm4vMoeForConditionalGeneration", + "Grok1VForCausalLM", + "Grok1AForCausalLM", + "LlavaLlamaForCausalLM", + "Llama4ForConditionalGeneration", + "LlavaMistralForCausalLM", + "LlavaQwenForCausalLM", + "LlavaForConditionalGeneration", + "LlavaVidForCausalLM", + "MiniCPMO", + "MiniCPMV", + "Mistral3ForConditionalGeneration", + "MultiModalityCausalLM", + "MllamaForConditionalGeneration", + "NemotronH_Nano_VL_V2", + "Qwen2AudioForConditionalGeneration", + "Qwen2VLForConditionalGeneration", + "Qwen2_5_VLForConditionalGeneration", + "Qwen3VLForConditionalGeneration", + "Qwen3VLMoeForConditionalGeneration", + "Qwen3OmniMoeForConditionalGeneration", + "KimiVLForConditionalGeneration", + "InternVLChatModel", + "InternS1ForConditionalGeneration", + "Phi4MMForCausalLM", + "Step3VLForConditionalGeneration", + "POINTSV15ChatModel", + "DotsVLMForCausalLM", + "DotsOCRForCausalLM", + "Sarashina2VisionForCausalLM", + "NVILAForConditionalGeneration", + "NVILALiteForConditionalGeneration", + "DeepseekOCRForCausalLM", + "JetVLMForConditionalGeneration", +] + +if envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.value: + multimodal_model_archs.append(envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.value) + + +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]): + return False + + +def is_encoder_decoder_model(model_architectures: List[str]): + return "MllamaForConditionalGeneration" 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 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 is_hybrid_model( + model_architectures: List[str], + hybrid_kvcache_ratio: Optional[float], + context_length: Optional[int], + attention_chunk_size: Optional[int], +): + if hybrid_kvcache_ratio is None: + return None + elif ( + hybrid_kvcache_ratio > 0 + and model_architectures[0] == "Llama4ForConditionalGeneration" + and context_length > attention_chunk_size + ): + return hybrid_kvcache_ratio + else: + return None + + +def get_hybrid_layer_ids(model_architectures: List[str], num_hidden_layers: int): + 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 + ] + else: + swa_attention_layer_ids = None + full_attention_layer_ids = None + return swa_attention_layer_ids, full_attention_layer_ids diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/modelopt_config.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/modelopt_config.py new file mode 100644 index 00000000..911b4ce0 --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/modelopt_config.py @@ -0,0 +1,30 @@ +# 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 diff --git a/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/nano_nemotron_vl.py b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/nano_nemotron_vl.py new file mode 100644 index 00000000..09ab29ab --- /dev/null +++ b/runners/nvidia_sglang_inference_opt_4fd1e59e/.vendor/sglang/python/sglang/srt/configs/nano_nemotron_vl.py @@ -0,0 +1,114 @@ +# 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 = "