diff --git a/.gitignore b/.gitignore index 92a7e21..43815d7 100644 --- a/.gitignore +++ b/.gitignore @@ -23,5 +23,8 @@ datasets/ # 原始请求级输出（方案 A：极简版，不存原始 jsonl） bench_results/**/raw_outputs/ +# 实验级原始结果目录（report.md / results.json 可单独保留，默认忽略整个目录） +experiments/*/results/ + # 无关项目 loomeval_yy/ diff --git a/BENCHMARK_WORKFLOW.md b/BENCHMARK_WORKFLOW.md index c9bb135..b0bd601 100644 --- a/BENCHMARK_WORKFLOW.md +++ b/BENCHMARK_WORKFLOW.md @@ -4,20 +4,29 @@ ``` /data/user1/yy/ -├── scripts/ # all benchmark/orchestrator/utility scripts -│ ├── benchmark_dspark_0707/ # DSpark benchmark suite -│ ├── benchmark_dsv4_backend_comparison.sh -│ ├── start_dsv4_dspark_8card.sh -│ ├── start_sglang_dsv4_8card.sh +├── platforms/ # chip/accelerator platform configs +│ ├── kunlun_p800.env +│ ├── nvidia_h200.env +│ └── patches/kunlun_p800/ # runtime patches required by some images +├── scripts/ # shared benchmark/orchestrator/utility scripts +│ ├── common/ # reusable components (lib.sh, platform.sh, ...) +│ ├── analysis/ # cross-experiment comparison tools +│ ├── benchmark_dspark_0707/ # legacy DSpark benchmark suite │ └── ... -├── bench_results/ # all benchmark outputs and reports -│ ├── dsv4_backend_comparison_20260707/ -│ │ ├── raw_outputs/ # JSONL raw outputs -│ │ ├── logs/ # per-run logs -│ │ └── README.md # output manifest + provenance +├── experiments/ # experiment-centric directories (preferred) +│ └── dsv4_p800_sglang/ +│ ├── README.md +│ ├── config.env # experiment-level configuration +│ ├── start_server.sh +│ ├── run_bench.sh +│ ├── parse_results.py +│ └── results/ +│ └── 20260708-XXXXXX/ +│ ├── report.md +│ ├── results.json +│ └── logs/ +├── bench_results/ # legacy benchmark outputs (read-only history) │ ├── dspark_grid_20260707-132641/ -│ ├── dspark_st_comparison_20260707-150649/ -│ ├── eagle_grid/ │ └── ... ├── logs/ # server logs (stdout/stderr from start scripts) ├── datasets/ # benchmark datasets @@ -26,16 +35,16 @@ ## Rules -1. **Scripts live in `scripts/` only.** - - Group related scripts into subdirectories, e.g. `scripts/benchmark_dspark_0707/`. - - Each script group should have its own `README.md` listing scripts, purpose, and outputs. +1. **Experiments are the primary organization unit.** + - Each experiment lives under `experiments//` and contains its scripts, configuration, and results. + - Shared orchestration code lives in `scripts/common/`; do not copy server start / health check logic into every experiment. + - Legacy experiments may remain under `scripts//` with outputs in `bench_results/`, but new work should use `experiments/`. -2. **Benchmark outputs live in `bench_results/` only.** - - Never leave `.jsonl`, `.json`, `.log`, or `.md` reports in the project root. - - Each benchmark run gets its own directory: `bench_results/_/`. - - Raw outputs go in `raw_outputs/`. - - Logs go in `logs/`. - - Reports (e.g. `report.md`, `comparison_report.md`) go in the run root. +2. **Benchmark outputs live with their experiment.** + - For `experiments//`, results go in `experiments//results//`. + - Each run directory must contain `report.md` (human-readable) and `results.json` (structured data). + - Raw outputs go in `raw_outputs/`; logs go in `logs/`. + - Legacy `bench_results/_/` directories remain valid for archived runs. 3. **Each `bench_results//` directory must contain two final artifacts.** - A Markdown report for human reading (e.g. `report.md`, `comparison_report.md`). @@ -50,7 +59,9 @@ - Keep the schema stable so downstream Python scripts can parse all experiments uniformly. - See [Final JSON Schema](#final-json-schema) below for the recommended structure. -5. **Scripts should default `RESULT_ROOT` to `bench_results/_${RUN_ID}`.** +5. **Scripts should default `RESULT_ROOT` to the experiment's results directory.** + - For `experiments//run_bench.sh`, default to `experiments//results/${RUN_ID}/`. + - For legacy scripts, default to `bench_results/_${RUN_ID}/`. - Allow override via `RESULT_ROOT` env var. - Use `RUN_ID=$(date '+%Y%m%d-%H%M%S')` unless specified. @@ -63,9 +74,33 @@ - Final reports and JSON outputs must include both the accelerator/chip family and the inference engine/backend used for the run. - Do not rely on the experiment name alone to identify the platform or engine. +8. **Use platform configuration files for chip-specific constants.** + - Put per-platform settings in `platforms/.env` (e.g. `platforms/kunlun_p800.env`, `platforms/nvidia_h200.env`). + - Scripts load the platform file via `scripts/common/platform.sh`; the active platform is selected by the `PLATFORM` env var or auto-detected. + - Keep experiment scripts free of hardcoded device IDs, image names, or model root paths. + ## Naming Conventions -### Result directories +### Experiment result directories + +For experiment-centric layout: + +``` +experiments//results// +``` + +Examples: + +- `experiments/dsv4_p800_sglang/results/20260708-120000/` +- `experiments/dspark_grid/results/20260707-132641/` + +The `results.json` metadata already records `chip`/`accelerator` and `engine`, so the directory path does not need to encode them. If a single experiment must distinguish across platforms in its directory tree, use: + +``` +experiments//results/__/ +``` + +### Legacy result directories ``` bench_results/_/ @@ -76,13 +111,6 @@ Examples: - `bench_results/dspark_grid_20260707-132641/` - `bench_results/dsv4_backend_comparison_20260707/` -When the same experiment is repeated across chips or engines, include them in the directory name or organize by subdirectories so results are not confused: - -``` -bench_results/___/ -bench_results///// -``` - ### Raw output files Include the accelerator and inference engine in raw output filenames so files from different platforms cannot overwrite each other. @@ -190,7 +218,13 @@ The JSON file inside each `bench_results//` directory should follow a stabl ## Quick Start -### Run DSpark grid benchmark +### Run P800 SGLang benchmark + +```bash +bash experiments/dsv4_p800_sglang/run_bench.sh +``` + +### Run legacy DSpark grid benchmark ```bash bash scripts/benchmark_dspark_0707/run_dspark_benchmark_grid.sh @@ -202,28 +236,33 @@ bash scripts/benchmark_dspark_0707/run_dspark_benchmark_grid.sh bash scripts/benchmark_dspark_0707/run_dspark_st_comparison.sh ``` -### Run SGLang vs vLLM backend comparison - -```bash -# Start SGLang on port 30000 and vLLM on port 8000, then: -bash scripts/benchmark_dsv4_backend_comparison.sh all -``` - ### Parse results ```bash +# Experiment-centric layout +python3 experiments/dsv4_p800_sglang/parse_results.py \ + experiments/dsv4_p800_sglang/results/ + +# Legacy layout /data/user1/yy/envs/sglang/bin/python scripts/benchmark_dspark_0707/parse_results.py \ /data/user1/yy/bench_results/dspark_grid_ ``` +### Cross-experiment comparison + +```bash +python3 scripts/analysis/compare_experiments.py +``` + ## Checklist Before Committing / Archiving -- [ ] No `.jsonl`, `.json`, `.log`, or `.md` files left in `/data/user1/yy/` root. -- [ ] All outputs moved to `bench_results/_/`. -- [ ] `bench_results//report.md` (or equivalent human-readable `.md`) exists. -- [ ] `bench_results//results.json` exists and follows the [Final JSON Schema](#final-json-schema). -- [ ] `bench_results//results.json` metadata records the `chip`/`accelerator` and `engine`/`backend` used. +- [ ] No `.jsonl`, `.json`, `.log`, or `.md` files left in the project root. +- [ ] For `experiments//`, outputs live in `experiments//results//`. +- [ ] For legacy runs, outputs live in `bench_results/_/`. +- [ ] `/report.md` (or equivalent human-readable `.md`) exists. +- [ ] `/results.json` exists and follows the [Final JSON Schema](#final-json-schema). +- [ ] `/results.json` metadata records the `chip`/`accelerator` and `engine`/`backend` used. - [ ] Raw output filenames include the chip/accelerator and engine when cross-platform runs may collide. -- [ ] `bench_results//README.md` exists and documents provenance (or the report itself covers provenance). -- [ ] Scripts moved to `scripts/` (or `scripts//`). +- [ ] `/README.md` exists and documents provenance (or the report itself covers provenance). +- [ ] Scripts either live under `experiments//` or in `scripts/` (or `scripts//`). - [ ] Script path references updated after moving. diff --git a/README.md b/README.md index 5d01ebc..e039add 100644 --- a/README.md +++ b/README.md @@ -6,8 +6,12 @@ | 目录/文件 | 说明 | |---|---| -| `scripts/` | 所有 benchmark 脚本、服务启动脚本、结果解析脚本 | -| `bench_results/` | 各次实验的最终产物：人可读 `.md` 报告 + 结构化 `results.json`（含原始数据与分位统计） | +| `platforms/` | 芯片/加速器平台配置（`*.env`），如 P800、H200 | +| `scripts/common/` | 跨实验复用的 orchestration 组件（server 启停、health check、元数据生成） | +| `scripts/analysis/` | 跨实验对比分析工具 | +| `scripts/` | 旧版/legacy benchmark 脚本 | +| `experiments/` | 以实验为单位的自包含目录（脚本 + 配置 + 结果） | +| `bench_results/` | 历史归档的实验产物（只读） | | `BENCHMARK_WORKFLOW.md` | benchmark 目录与命名规范 | | `scripts/SLO_STANDARDS.md` | 推理服务 SLO 标准（TTFT/TPOT） | @@ -15,6 +19,14 @@ ## 实验索引 +### 新结构（experiments/） + +| 实验 | 脚本 | 说明 | +|---|---|---| +| DSV4 P800 SGLang | `experiments/dsv4_p800_sglang/run_bench.sh` | Kunlun P800 + SGLang + DeepSeek-V4-Flash-INT8 | + +### 旧结构（scripts/ + bench_results/） + | 实验 | 脚本 | 最终报告 | 说明 | |---|---|---|---| | DSpark grid benchmark | `scripts/benchmark_dspark_0707/run_dspark_benchmark_grid.sh` | `bench_results/dspark_grid_20260707-132641/report.md` | P1/P2/P3 全量网格 | @@ -28,6 +40,12 @@ ## 快速复现 +### P800 SGLang + +```bash +bash experiments/dsv4_p800_sglang/run_bench.sh +``` + ### DSpark grid ```bash @@ -43,6 +61,11 @@ bash scripts/benchmark_dspark_0707/run_dspark_st_comparison.sh ### 解析已有结果 ```bash +# 新结构 +python3 experiments/dsv4_p800_sglang/parse_results.py \ + experiments/dsv4_p800_sglang/results/ + +# 旧结构 /data/user1/yy/envs/sglang/bin/python scripts/benchmark_dspark_0707/parse_results.py \ /data/user1/yy/bench_results/dspark_grid_20260707-132641 ``` @@ -58,6 +81,7 @@ bash scripts/benchmark_dspark_0707/run_dspark_st_comparison.sh ## 环境要求 -- Python env：`/data/user1/yy/envs/vllm-dspark`（服务端）、`/data/user1/yy/envs/sglang`（压测客户端） -- 模型：`/data/models/DeepSeek-V4-Flash`、`/data/models/DeepSeek-V4-Flash-DSpark` -- 硬件：8× H200（当前配置） +- **NVIDIA H200**：Python env `/data/user1/yy/envs/vllm-dspark`（服务端）、`/data/user1/yy/envs/sglang`（压测客户端）；模型 `/data/models/DeepSeek-V4-Flash`、`/data/models/DeepSeek-V4-Flash-DSpark`。 +- **Kunlun P800**：Docker 镜像 `iregistry.baidu-int.com/xpu/sglang-p800-pd-disagg-0510:20260511_4202`；模型 `/data1/models/DeepSeek-V4-Flash-INT8`；压测客户端在容器内运行。 + +平台相关常量见 `platforms/*.env`。 diff --git a/experiments/dsv4_p800_sglang/README.md b/experiments/dsv4_p800_sglang/README.md new file mode 100644 index 0000000..589413f --- /dev/null +++ b/experiments/dsv4_p800_sglang/README.md @@ -0,0 +1,28 @@ +# DSV4 P800 SGLang Benchmark + +Kunlun P800 XPU + SGLang + `DeepSeek-V4-Flash-INT8` benchmark experiment. + +## Quick Start + +```bash +# Run all configured scenarios (starts server, runs benchmark, stops server) +bash experiments/dsv4_p800_sglang/run_bench.sh + +# Reuse an already-running server +SKIP_MANAGE_SERVER=1 bash experiments/dsv4_p800_sglang/run_bench.sh +``` + +Results land in `experiments/dsv4_p800_sglang/results//`. + +## Files + +| File | Purpose | +|---|---| +| `config.env` | Experiment-level configuration (model, port, scenarios) | +| `start_server.sh` | Start the P800 SGLang Docker container | +| `run_bench.sh` | Orchestrator: server → benchmark → stop server | +| `parse_results.py` | Parse logs and generate `results.json` + `report.md` | + +## Platform + +This experiment targets `platforms/kunlun_p800.env`. diff --git a/experiments/dsv4_p800_sglang/config.env b/experiments/dsv4_p800_sglang/config.env new file mode 100644 index 0000000..0921fb3 --- /dev/null +++ b/experiments/dsv4_p800_sglang/config.env @@ -0,0 +1,35 @@ +# Experiment-level configuration for dsv4_p800_sglang +# All values can be overridden via environment variables. + +EXPERIMENT="${EXPERIMENT:-dsv4_p800_sglang}" +MODEL_NAME="${MODEL_NAME:-DeepSeek-V4-Flash-INT8}" +MODEL_PATH="${MODEL_PATH:-/data1/models/DeepSeek-V4-Flash-INT8}" +SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-deepseek-v4-flash-int8}" +PORT="${PORT:-30000}" +BACKEND="${BACKEND:-sglang}" +ENGINE="${ENGINE:-sglang-xpu}" +DATASET="${DATASET:-random}" +# Path to a local ShareGPT-style dataset inside the container. +# Required when HF Hub is offline. The repo ships a dummy dataset under +# platforms/patches/kunlun_p800/dummy_sharegpt.json; for managed mode it is +# mounted to /workspace/dummy_sharegpt.json. +DATASET_PATH="${DATASET_PATH:-/workspace/dummy_sharegpt.json}" +WARMUP="${WARMUP:-100}" +NUM_PROMPTS="${NUM_PROMPTS:-512}" + +# Server launch mode for P800 SGLang. +# - "fp8" : non-INT8 DeepSeek-V4-Flash checkpoint +# - "w8a8_int8": DeepSeek-V4-Flash-INT8 checkpoint +SERVER_MODE="${SERVER_MODE:-w8a8_int8}" + +# Scenario list: "concurrency input_len output_len" (one per line) +# Override example: SCENARIOS=("32 512 256") +if [[ -z "${SCENARIOS:-}" ]]; then + SCENARIOS=( + "32 512 256" + "128 512 256" + "256 512 256" + "512 512 256" + ) +fi + diff --git a/experiments/dsv4_p800_sglang/parse_results.py b/experiments/dsv4_p800_sglang/parse_results.py new file mode 100755 index 0000000..341f69e --- /dev/null +++ b/experiments/dsv4_p800_sglang/parse_results.py @@ -0,0 +1,313 @@ +#!/usr/bin/env python3 +"""Parse P800 SGLang benchmark outputs and produce results.json + report.md. + +Usage: + python3 parse_results.py +""" + +import json +import os +import re +import sys +from pathlib import Path +from collections import OrderedDict + +METRIC_PATTERNS = OrderedDict( + [ + ("successful_requests", [r"Successful requests:\s+(\d+)"]), + ("benchmark_duration_s", [r"Benchmark duration \(s\):\s+([\d.]+)"]), + ("request_throughput", [r"Request throughput \(req/s\):\s+([\d.]+)"]), + ("input_token_throughput", [r"Input token throughput \(tok/s\):\s+([\d.]+)"]), + ("output_token_throughput", [r"Output token throughput \(tok/s\):\s+([\d.]+)"]), + ("total_token_throughput", [r"Total token throughput \(tok/s\):\s+([\d.]+)"]), + ("ttft_mean", [r"Mean TTFT \(ms\):\s+([\d.]+)"]), + ("ttft_p50", [r"Median TTFT \(ms\):\s+([\d.]+)", r"P50 TTFT \(ms\):\s+([\d.]+)"]), + ("ttft_p90", [r"P90 TTFT \(ms\):\s+([\d.]+)"]), + ("ttft_p99", [r"P99 TTFT \(ms\):\s+([\d.]+)"]), + ("tpot_mean", [r"Mean TPOT \(ms\):\s+([\d.]+)"]), + ("tpot_p50", [r"Median TPOT \(ms\):\s+([\d.]+)", r"P50 TPOT \(ms\):\s+([\d.]+)"]), + ("tpot_p90", [r"P90 TPOT \(ms\):\s+([\d.]+)"]), + ("tpot_p99", [r"P99 TPOT \(ms\):\s+([\d.]+)"]), + ("e2e_mean", [r"Mean E2E \(ms\):\s+([\d.]+)"]), + ("e2e_p50", [r"Median E2E \(ms\):\s+([\d.]+)", r"P50 E2E \(ms\):\s+([\d.]+)"]), + ("e2e_p90", [r"P90 E2E \(ms\):\s+([\d.]+)"]), + ("e2e_p99", [r"P99 E2E \(ms\):\s+([\d.]+)"]), + ] +) + + +def find_float(text: str, patterns: list[str]) -> float | None: + for pat in patterns: + m = re.search(pat, text) + if m: + try: + return float(m.group(1)) + except ValueError: + return None + return None + + +def parse_summary_log(log_text: str) -> dict: + result = {} + for key, patterns in METRIC_PATTERNS.items(): + result[key] = find_float(log_text, patterns) + return result + + +def percentile(values: list[float], p: float) -> float: + if not values: + return 0.0 + sorted_values = sorted(values) + k = (len(sorted_values) - 1) * p / 100.0 + f = int(k) + c = min(f + 1, len(sorted_values) - 1) + if f == c: + return sorted_values[f] + return sorted_values[f] * (c - k) + sorted_values[c] * (k - f) + + +def parse_jsonl(jsonl_path: Path) -> tuple[list[dict], dict]: + """Parse sglang.bench_serving --output-file output. + + Newer sglang writes one aggregate JSON object. Older versions write one + JSON object per request. Returns (raw_requests, aggregates). + """ + text = jsonl_path.read_text(encoding="utf-8", errors="replace").strip() + if not text: + return [], {} + + # Try single aggregate JSON object first. + try: + data = json.loads(text) + if isinstance(data, dict) and "mean_e2e_latency_ms" in data: + aggregates = { + "success": data.get("total_output_tokens", 0) > 0 and 1 or 0, + "failed": 0, + "input_tokens": data.get("total_input_tokens", 0), + "output_tokens": data.get("total_output_tokens", 0), + "latencies": { + "e2e_ms": { + "mean": data.get("mean_e2e_latency_ms"), + "p50": data.get("median_e2e_latency_ms"), + "p90": data.get("p90_e2e_latency_ms"), + "p95": None, + "p99": data.get("p99_e2e_latency_ms"), + }, + "ttft_ms": { + "mean": data.get("mean_ttft_ms"), + "p50": data.get("median_ttft_ms"), + "p90": None, + "p95": None, + "p99": data.get("p99_ttft_ms"), + }, + "tpot_ms": { + "mean": data.get("mean_tpot_ms"), + "p50": data.get("median_tpot_ms"), + "p90": None, + "p95": None, + "p99": data.get("p99_tpot_ms"), + }, + "itl_ms": { + "mean": data.get("mean_itl_ms"), + "p50": data.get("median_itl_ms"), + "p90": None, + "p95": data.get("p95_itl_ms"), + "p99": data.get("p99_itl_ms"), + }, + }, + } + return [data], aggregates + except json.JSONDecodeError: + pass + + # Fall back to JSONL per-request parsing. + raw_requests = [] + ttfts = [] + tpots = [] + itls = [] + e2es = [] + input_tokens = [] + output_tokens = [] + success = 0 + failed = 0 + + for line in text.splitlines(): + line = line.strip() + if not line: + continue + try: + req = json.loads(line) + except json.JSONDecodeError: + continue + + raw_requests.append(req) + + ttft = req.get("ttft") or req.get("ttft_ms") or 0 + tpot = req.get("tpot") or req.get("tpot_ms") or 0 + itl = req.get("inter_token_latency") or req.get("itl") or req.get("itl_ms") or 0 + e2e = req.get("e2e_latency") or req.get("e2e") or req.get("e2e_ms") or 0 + in_tok = req.get("input_tokens") or req.get("prompt_tokens") or 0 + out_tok = req.get("output_tokens") or req.get("completion_tokens") or 0 + + if ttft: + ttfts.append(float(ttft)) + if tpot: + tpots.append(float(tpot)) + if itl: + itls.append(float(itl)) + if e2e: + e2es.append(float(e2e)) + if in_tok: + input_tokens.append(int(in_tok)) + if out_tok: + output_tokens.append(int(out_tok)) + + if req.get("success", True): + success += 1 + else: + failed += 1 + + def latency_bucket(values: list[float]) -> dict: + if not values: + return {"mean": None, "p50": None, "p90": None, "p95": None, "p99": None} + return { + "mean": round(sum(values) / len(values), 2), + "p50": round(percentile(values, 50), 2), + "p90": round(percentile(values, 90), 2), + "p95": round(percentile(values, 95), 2), + "p99": round(percentile(values, 99), 2), + } + + aggregates = { + "success": success, + "failed": failed, + "input_tokens": sum(input_tokens), + "output_tokens": sum(output_tokens), + "latencies": { + "e2e_ms": latency_bucket(e2es), + "ttft_ms": latency_bucket(ttfts), + "tpot_ms": latency_bucket(tpots), + "itl_ms": latency_bucket(itls), + }, + } + + return raw_requests, aggregates + + +def parse_scenario(result_root: Path, raw_file: Path) -> dict | None: + """Parse one raw output file into a scenario dict.""" + # Filename: {chip}_{engine}_{MMDD}_{concurrency}_{input_len}_{output_len}.jsonl + parts = raw_file.stem.split("_") + if len(parts) < 6: + return None + try: + concurrency = int(parts[-3]) + input_len = int(parts[-2]) + output_len = int(parts[-1]) + except ValueError: + return None + + log_file = result_root / "logs" / f"sglang_c{concurrency}_i{input_len}_o{output_len}.log" + summary = {} + if log_file.exists(): + summary = parse_summary_log(log_file.read_text(encoding="utf-8", errors="replace")) + + raw_requests, aggregates = parse_jsonl(raw_file) + + # For newer sglang aggregate JSON, success/failed are not present in the + # raw output file. Override them from the human-readable summary log when + # it is available. + if summary.get("successful_requests") is not None: + aggregates["success"] = int(summary["successful_requests"]) + # The summary log only reports successes; assume failures are zero + # unless the aggregate JSON already provided a non-zero failed count. + if not aggregates.get("failed"): + aggregates["failed"] = 0 + + scenario = { + "name": f"c{concurrency}_i{input_len}_o{output_len}", + "concurrency": concurrency, + "input_len": input_len, + "output_len": output_len, + "success": aggregates["success"], + "failed": aggregates["failed"], + "duration_s": summary.get("benchmark_duration_s"), + "request_throughput": summary.get("request_throughput"), + "input_token_throughput": summary.get("input_token_throughput"), + "output_token_throughput": summary.get("output_token_throughput"), + "total_token_throughput": summary.get("total_token_throughput"), + "accept_length": None, + "latencies": aggregates["latencies"], + "raw_requests": raw_requests[:100] if len(raw_requests) <= 100 else None, + } + + return scenario + + +def main() -> None: + result_root = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("results") + raw_dir = result_root / "raw_outputs" + json_path = result_root / "results.json" + report_path = result_root / "report.md" + + if not json_path.exists(): + raise SystemExit(f"metadata results.json not found: {json_path}") + + with open(json_path, "r", encoding="utf-8") as f: + data = json.load(f) + + scenarios = [] + if raw_dir.exists(): + for raw_file in sorted(raw_dir.glob("*.jsonl")): + scenario = parse_scenario(result_root, raw_file) + if scenario: + scenarios.append(scenario) + + data["scenarios"] = scenarios + + with open(json_path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) + + # Generate report.md + with open(report_path, "w", encoding="utf-8") as f: + meta = data["metadata"] + f.write(f"# Benchmark Report: {meta['experiment']}\n\n") + f.write("## Metadata\n\n") + f.write(f"- **Run ID**: {meta['run_id']}\n") + f.write(f"- **Timestamp**: {meta['timestamp']}\n") + f.write(f"- **Chip/Accelerator**: {meta.get('accelerator', '')} / {meta.get('chip', '')}\n") + f.write(f"- **Engine/Backend**: {meta.get('engine', '')} / {meta.get('backend', '')}\n") + f.write(f"- **Hardware**: {meta.get('hardware', '')}\n") + f.write(f"- **Model**: {meta['model']}\n") + f.write(f"- **Git Commit**: {meta.get('git_commit', 'unknown')}\n") + f.write("\n## Results\n\n") + f.write( + "| Scenario | Conc | In/Out | Success | Failed | Req/s | OutTok/s | " + "TTFT p50 | TTFT p99 | TPOT p50 | TPOT p99 | E2E p99 |\n" + ) + f.write( + "|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|\n" + ) + for s in scenarios: + lat = s["latencies"] + f.write( + f"| {s['name']} " + f"| {s['concurrency']} " + f"| {s['input_len']}/{s['output_len']} " + f"| {s['success']} " + f"| {s['failed']} " + f"| {s.get('request_throughput') or ''} " + f"| {s.get('output_token_throughput') or ''} " + f"| {lat['ttft_ms']['p50'] or ''} " + f"| {lat['ttft_ms']['p99'] or ''} " + f"| {lat['tpot_ms']['p50'] or ''} " + f"| {lat['tpot_ms']['p99'] or ''} " + f"| {lat['e2e_ms']['p99'] or ''} |\n" + ) + + print(f"Updated {json_path}") + print(f"Wrote {report_path}") + + +if __name__ == "__main__": + main() diff --git a/experiments/dsv4_p800_sglang/run_bench.sh b/experiments/dsv4_p800_sglang/run_bench.sh new file mode 100755 index 0000000..3b8d29e --- /dev/null +++ b/experiments/dsv4_p800_sglang/run_bench.sh @@ -0,0 +1,112 @@ +#!/usr/bin/env bash +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +EXPERIMENT_NAME="$(basename "$SCRIPT_DIR")" + +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../scripts/common/server_docker.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../scripts/common/bench_client_docker.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" + +RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}" +RESULT_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}" +RAW_DIR="${RESULT_ROOT}/raw_outputs" +LOG_DIR="${RESULT_ROOT}/logs" + +ensure_result_root "$RESULT_ROOT" +log_init "${LOG_DIR}/orchestrator.log" + +log "experiment=${EXPERIMENT_NAME}" +log "run_id=${RUN_ID}" +log "result_root=${RESULT_ROOT}" +log "platform=${PLATFORM}" +log "chip=${CHIP}" +log "accelerator=${ACCELERATOR}" +log "engine=${ENGINE}" +log "hardware=${HARDWARE}" + +# Write initial metadata for this run. +METADATA_JSON="${RESULT_ROOT}/results.json" +write_metadata_json \ + "$METADATA_JSON" \ + "$EXPERIMENT_NAME" \ + "$RUN_ID" \ + "$MODEL_PATH" \ + "$BACKEND" \ + "$ENGINE" \ + "$HARDWARE" \ + "$ACCELERATOR" \ + "$CHIP" \ + "experiments/${EXPERIMENT_NAME}/run_bench.sh" \ + "" \ + "Kunlun P800 SGLang benchmark for DeepSeek-V4-Flash-INT8" + +# Start server unless SKIP_MANAGE_SERVER is set. +if [[ -n "${SKIP_MANAGE_SERVER:-}" ]]; then + log "SKIP_MANAGE_SERVER is set, assuming server is already running on port ${PORT}" + if ! health_check 127.0.0.1 "${PORT}" 30; then + log "ERROR: no healthy server found at port ${PORT}" + exit 1 + fi +else + docker_server_start "${MODEL_PATH}" "${PORT}" "${LOG_DIR}/server.outer.log" "${SERVER_MODE}" +fi + +on_exit() { + local code=$? + if [[ -z "${SKIP_MANAGE_SERVER:-}" ]]; then + docker_server_stop + fi + log "orchestrator exiting with code=${code}" + exit "$code" +} +trap on_exit EXIT + +log "===== BENCHMARK START =====" + +for scenario in "${SCENARIOS[@]}"; do + # Allow SCENARIOS to be passed as a string like "(32 512 256)" from env. + scenario="${scenario//[()]/}" + read -r concurrency input_len output_len <<< "$scenario" + + tmp_output_file="/tmp/bench_outputs/${CHIP}_${ENGINE}_$(date '+%m%d')_${concurrency}_${input_len}_${output_len}.jsonl" + host_output_file="${RAW_DIR}/${CHIP}_${ENGINE}_$(date '+%m%d')_${concurrency}_${input_len}_${output_len}.jsonl" + detail_log="${LOG_DIR}/${BACKEND}_c${concurrency}_i${input_len}_o${output_len}.log" + + log "running scenario: concurrency=${concurrency} input=${input_len} output=${output_len}" + + run_random_case \ + "$BACKEND" \ + "$PORT" \ + "$MODEL_PATH" \ + "$tmp_output_file" \ + "$concurrency" \ + "$input_len" \ + "$output_len" \ + "$NUM_PROMPTS" \ + "${DATASET_PATH:-}" \ + > "$detail_log" 2>&1 || { + log "ERROR: scenario c=${concurrency} i=${input_len} o=${output_len} failed" + continue + } + + # Copy the JSONL from the container to the host result directory. + docker cp "${CONTAINER_NAME}:${tmp_output_file}" "$host_output_file" || { + log "ERROR: failed to copy ${tmp_output_file} from container" + continue + } + log "finished scenario: output=${host_output_file}" +done + +log "===== BENCHMARK DONE =====" +log "parsing results..." +"${SCRIPT_DIR}/parse_results.py" "$RESULT_ROOT" + +log "all results saved to ${RESULT_ROOT}" diff --git a/experiments/dsv4_p800_sglang/start_server.sh b/experiments/dsv4_p800_sglang/start_server.sh new file mode 100755 index 0000000..63467b6 --- /dev/null +++ b/experiments/dsv4_p800_sglang/start_server.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +EXPERIMENT_NAME="$(basename "$SCRIPT_DIR")" + +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../scripts/common/server_docker.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" + +log_init "${RESULT_ROOT:-/tmp/${EXPERIMENT_NAME}}/logs/start_server.log" +log "starting server for ${EXPERIMENT_NAME}" +log "model: ${MODEL_PATH}" +log "port: ${PORT}" + +docker_server_start "${MODEL_PATH}" "${PORT}" "${LOG_DIR}/server.outer.log" "${SERVER_MODE}" diff --git a/platforms/README.md b/platforms/README.md new file mode 100644 index 0000000..a03c9f2 --- /dev/null +++ b/platforms/README.md @@ -0,0 +1,36 @@ +# Platform Configurations + +Each `.env` file in this directory describes one accelerator platform. +They are meant to be sourced by benchmark scripts through +`scripts/common/platform.sh`, not executed directly. + +## Usage + +```bash +# Default platform for the current machine +bash experiments/dsv4_p800_sglang/run_bench.sh + +# Explicitly select a platform +PLATFORM=kunlun_p800 bash experiments/dsv4_p800_sglang/run_bench.sh +``` + +## Current platforms + +| File | Chip/Accelerator | Engine | Notes | +|---|---|---|---| +| `kunlun_p800.env` | Kunlun P800 XPU | `sglang-xpu` | Docker-based SGLang serving image | +| `nvidia_h200.env` | NVIDIA H200 | native vllm/sglang | Host virtual environments | + +## What belongs here + +- Chip/accelerator identity (`CHIP`, `ACCELERATOR`, `HARDWARE`). +- Device selection environment variables. +- Platform-wide paths that rarely change (model root, default port). +- Container image / interpreter paths for Docker-based platforms. + +## What does NOT belong here + +- Specific model names or experiment scenarios — those go in + `experiments//config.env`. +- Engine-specific launch flags — those go in the experiment's + `start_server.sh` or `run_bench.sh`. diff --git a/platforms/kunlun_p800.env b/platforms/kunlun_p800.env new file mode 100644 index 0000000..931edb7 --- /dev/null +++ b/platforms/kunlun_p800.env @@ -0,0 +1,27 @@ +# Platform configuration for Kunlun P800 XPU +# Source this file via scripts/common/platform.sh + +CHIP="kunlun_p800" +ACCELERATOR="Kunlun P800 XPU" +HARDWARE="8x Kunlun P800 XPU" +ENGINE="sglang-xpu" + +# Container runtime +DOCKER_IMAGE="iregistry.baidu-int.com/xpu/sglang-p800-pd-disagg-0510:20260511_4202" +CONTAINER_NAME="${CONTAINER_NAME:-sglang-dsv4-flash}" + +# Device selection +DEVICE_SELECT_ENV="XPU_VISIBLE_DEVICES=0,1,2,3,4,5,6,7" +CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7" +CUDA_DEVICE_ORDER="OAM_ID" + +# Default serving port and model root on the host +DEFAULT_PORT="30000" +MODEL_ROOT="/data1/models" + +# Python interpreter inside the container +CONTAINER_PYTHON="/root/miniconda/envs/python310_torch25_cuda/bin/python" + +# Host mount points for patches required by the P800 SGLang image +# These are baked into the container start command in server_docker.sh. +PATCH_ROOT="${PATCH_ROOT:-/data1/yy/sskj/platforms/patches/kunlun_p800}" diff --git a/platforms/nvidia_h200.env b/platforms/nvidia_h200.env new file mode 100644 index 0000000..a7a0e8e --- /dev/null +++ b/platforms/nvidia_h200.env @@ -0,0 +1,18 @@ +# Platform configuration for NVIDIA H200 +# Source this file via scripts/common/platform.sh + +CHIP="nvidia_h200" +ACCELERATOR="NVIDIA H200" +HARDWARE="8x NVIDIA H200 143GB" + +# Device selection +DEVICE_SELECT_ENV="CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7" +CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7" + +# Default serving port and model root on the host +DEFAULT_PORT="30004" +MODEL_ROOT="/data/models" + +# Virtual environments on the host (used by native H200 scripts) +VENV_VLLM_DSPARK="/data/user1/yy/envs/vllm-dspark" +VENV_SGLANG="/data/user1/yy/envs/sglang" diff --git a/platforms/patches/kunlun_p800/config_backup_small_fp8.json b/platforms/patches/kunlun_p800/config_backup_small_fp8.json new file mode 100644 index 0000000..0d69072 --- /dev/null +++ b/platforms/patches/kunlun_p800/config_backup_small_fp8.json @@ -0,0 +1,114 @@ +{ + "architectures": [ + "DeepseekXYZForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 0, + "eos_token_id": 1, + "hc_eps": 1e-06, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "head_dim": 512, + "hidden_act": "silu", + "hidden_size": 4096, + "index_head_dim": 128, + "index_n_heads": 64, + "index_topk": 512, + "initializer_range": 0.02, + "max_position_embeddings": 1048576, + "model_type": "deepseek_ref", + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "n_shared_experts": 1, + "norm_topk_prob": true, + "num_attention_heads": 64, + "num_experts_per_tok": 6, + "num_hidden_layers": 43, + "num_hash_layers": 3, + "num_key_value_heads": 1, + "num_nextn_predict_layers": 1, + "o_groups": 8, + "o_lora_rank": 1024, + "q_lora_rank": 1024, + "qk_rope_head_dim": 64, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "beta_fast": 32.0, + "beta_slow": 1.0, + "factor": 16.0, + "original_max_position_embeddings": 65536, + "type": "yarn" + }, + "rope_theta": 10000, + "routed_scaling_factor": 1.5, + "scoring_func": "sqrtsoftplus", + "sliding_window": 128, + "swiglu_limit": 10.0, + "tie_word_embeddings": false, + "n_group": 8, + "topk_group": 8, + "topk_method": "noaux_tc", + "torch_dtype": "bfloat16", + "transformers_version": "4.57.1", + "use_cache": true, + "vocab_size": 129280, + "compress_rope_theta": 160000, + "compress_ratios": [ + 0, + 0, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 0 + ], + "quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "scale_fmt": "ue8m0", + "weight_block_size": [ + 128, + 128 + ] + }, + "expert_dtype": "fp4" +} \ No newline at end of file diff --git a/platforms/patches/kunlun_p800/config_backup_small_w8a8_int8.json b/platforms/patches/kunlun_p800/config_backup_small_w8a8_int8.json new file mode 100644 index 0000000..590de43 --- /dev/null +++ b/platforms/patches/kunlun_p800/config_backup_small_w8a8_int8.json @@ -0,0 +1,58 @@ +{ + "architectures": [ + "DeepseekXYZForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 0, + "eos_token_id": 1, + "hc_eps": 1e-06, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "head_dim": 512, + "hidden_act": "silu", + "hidden_size": 4096, + "index_head_dim": 128, + "index_n_heads": 64, + "index_topk": 512, + "initializer_range": 0.02, + "max_position_embeddings": 1048576, + "model_type": "deepseek_ref", + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "n_shared_experts": 1, + "norm_topk_prob": true, + "num_attention_heads": 64, + "num_experts_per_tok": 6, + "num_hidden_layers": 43, + "num_hash_layers": 3, + "num_key_value_heads": 1, + "num_nextn_predict_layers": 1, + "o_groups": 8, + "o_lora_rank": 1024, + "q_lora_rank": 1024, + "qk_rope_head_dim": 64, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "beta_fast": 32, + "beta_slow": 1, + "factor": 16, + "original_max_position_embeddings": 65536, + "type": "yarn" + }, + "rope_theta": 10000, + "routed_scaling_factor": 1.5, + "scoring_func": "sqrtsoftplus", + "sliding_window": 128, + "swiglu_limit": 10.0, + "tie_word_embeddings": false, + "n_group": 8, + "topk_group": 8, + "topk_method": "noaux_tc", + "torch_dtype": "bfloat16", + "transformers_version": "4.57.1", + "use_cache": true, + "vocab_size": 129280, + "compress_rope_theta": 160000, + "compress_ratios": [0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0] +} diff --git a/platforms/patches/kunlun_p800/fp8_utils.py.patched b/platforms/patches/kunlun_p800/fp8_utils.py.patched new file mode 100755 index 0000000..28ed63f --- /dev/null +++ b/platforms/patches/kunlun_p800/fp8_utils.py.patched @@ -0,0 +1,1689 @@ +from __future__ import annotations + +import logging +from enum import Enum +from functools import lru_cache +from typing import TYPE_CHECKING, Callable, List, Optional, Tuple + +import torch + +from sglang.srt.layers import deep_gemm_wrapper +from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8 +from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil +from sglang.srt.utils.common import torch_release + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs + +from sglang.srt.layers.quantization.fp8_kernel import ( + fp8_dtype, + fp8_max, + is_fp8_fnuz, + mxfp8_block_scaled_matmul_triton, + per_token_group_quant_fp8, + scaled_fp8_quant, + sglang_per_token_quant_fp8, + static_quant_fp8, + triton_scaled_mm, + w8a8_block_fp8_matmul_deepgemm, + w8a8_block_fp8_matmul_triton, +) +from sglang.srt.utils import ( + ceil_align, + ceil_div, + get_bool_env_var, + get_cuda_version, + get_device_capability, + is_blackwell_supported, + is_cuda, + is_flashinfer_available, + is_gfx95_supported, + is_hip, + is_sm90_supported, + is_sm100_supported, + is_sm120_supported, + offloader, +) +from sglang.srt.utils.custom_op import register_custom_op + +logger = logging.getLogger(__name__) + +_is_hip = is_hip() +_is_cuda = is_cuda() +_is_fp8_fnuz = is_fp8_fnuz() +_is_sm100_supported = is_sm100_supported() +_is_sm120_supported = is_sm120_supported() +_is_gfx95_supported = is_gfx95_supported() + +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_use_aiter_gfx95 = _use_aiter and _is_gfx95_supported + + +def use_aiter_triton_gemm_w8a8_tuned_gfx950(n: int, k: int) -> bool: + return (n, k) in [ + (1024, 8192), + (16384, 1536), + (2112, 7168), + (3072, 1536), + (32768, 8192), + (4096, 7168), + (4608, 7168), + (512, 7168), + (7168, 2048), + (7168, 2304), + (7168, 16384), + (7168, 256), + (8192, 1024), + (8192, 32768), + ] + + +if _use_aiter: + import aiter + + # from aiter import gemm_a8w8_blockscale, gemm_a8w8_bpreshuffle, get_hip_quant + from aiter import gemm_a8w8_blockscale as gemm_a8w8_blockscale + from aiter import gemm_a8w8_bpreshuffle, get_hip_quant + from aiter.ops.triton.gemm_a8w8_blockscale import ( + gemm_a8w8_blockscale as triton_gemm_a8w8_blockscale, + ) + + aiter_per1x128_quant = get_hip_quant(aiter.QuantType.per_1x128) + + +if _is_cuda: + from sgl_kernel import fp8_blockwise_scaled_mm, fp8_scaled_mm + + from sglang.srt.utils.patch_torch import register_fake_if_exists + + @register_fake_if_exists("sgl_kernel::fp8_scaled_mm") + def _fp8_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): + # mat_a: [M, K], mat_b: [K, N] or [N, K] depending on callsite layout; output is [M, N]. + M = mat_a.shape[-2] + N = mat_b.shape[-1] + return mat_a.new_empty((M, N), dtype=out_dtype) + + @register_fake_if_exists("sgl_kernel::fp8_blockwise_scaled_mm") + def _fp8_blockwise_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype): + # mat_a: [M, K], mat_b: [K, N] or [N, K] depending on callsite layout; output is [M, N]. + M = mat_a.shape[-2] + N = mat_b.shape[-1] + return mat_a.new_empty((M, N), dtype=out_dtype) + + +use_vllm_cutlass_w8a8_fp8_kernel = get_bool_env_var("USE_VLLM_CUTLASS_W8A8_FP8_KERNEL") +use_triton_w8a8_fp8_kernel = get_bool_env_var("USE_TRITON_W8A8_FP8_KERNEL") + +# Input scaling factors are no longer optional in _scaled_mm starting +# from pytorch 2.5. Allocating a dummy tensor to pass as input_scale +TORCH_DEVICE_IDENTITY = None + + +def use_rowwise_torch_scaled_mm(): + if _is_hip: + # The condition to determine if it is on a platform that supports + # torch._scaled_mm rowwise feature. + # The condition is determined once as the operations + # are time consuming. + return get_device_capability() >= (9, 4) and torch_release >= (2, 7) + return False + + +USE_ROWWISE_TORCH_SCALED_MM = use_rowwise_torch_scaled_mm() + + +@lru_cache(maxsize=1) +def cutlass_fp8_supported(): + if not _is_cuda: + return False + major, minor = get_device_capability() + cuda_version = get_cuda_version() + if major >= 9: + return cuda_version >= (12, 0) + elif major == 8 and minor == 9: + return cuda_version >= (12, 4) + return False + + +def normalize_e4m3fn_to_e4m3fnuz( + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + assert weight.dtype == torch.float8_e4m3fn + # The bits pattern 10000000(-128) represents zero in e4m3fn + # but NaN in e4m3fnuz. So here we set it to 0. + # https://onnx.ai/onnx/technical/float8.html + weight_as_int8 = weight.view(torch.int8) + ROCM_FP8_NAN_AS_INT = -128 + weight_as_int8[weight_as_int8 == ROCM_FP8_NAN_AS_INT] = 0 + weight = weight_as_int8.view(torch.float8_e4m3fnuz) + + # For the same bits representation, e4m3fnuz value is half of + # the e4m3fn value, so we should double the scaling factor to + # get the same dequantized value. + # https://onnx.ai/onnx/technical/float8.html + weight_scale = weight_scale * 2.0 + if input_scale is not None: + input_scale = input_scale * 2.0 + return weight, weight_scale, input_scale + + +class Fp8GemmRunnerBackend(Enum): + """Enum for FP8 GEMM runner backend selection.""" + + AUTO = "auto" + FLASHINFER_TRTLLM = "flashinfer_trtllm" + FLASHINFER_CUTLASS = "flashinfer_cutlass" + FLASHINFER_DEEPGEMM = "flashinfer_deepgemm" + CUTLASS = "cutlass" + DEEP_GEMM = "deep_gemm" + TRITON = "triton" + AITER = "aiter" + + def is_auto(self) -> bool: + return self == Fp8GemmRunnerBackend.AUTO + + def is_flashinfer_trtllm(self) -> bool: + return self == Fp8GemmRunnerBackend.FLASHINFER_TRTLLM + + def is_flashinfer_cutlass(self) -> bool: + return self == Fp8GemmRunnerBackend.FLASHINFER_CUTLASS + + def is_flashinfer_deepgemm(self) -> bool: + return self == Fp8GemmRunnerBackend.FLASHINFER_DEEPGEMM + + def is_cutlass(self) -> bool: + return self == Fp8GemmRunnerBackend.CUTLASS + + def is_deep_gemm(self) -> bool: + return self == Fp8GemmRunnerBackend.DEEP_GEMM + + def is_triton(self) -> bool: + return self == Fp8GemmRunnerBackend.TRITON + + def is_aiter(self) -> bool: + return self == Fp8GemmRunnerBackend.AITER + + +FP8_GEMM_RUNNER_BACKEND: Fp8GemmRunnerBackend | None = None + + +def _check_cutlass_block_fp8_hardware_support() -> bool: + """Return True if CUTLASS block FP8 is supported (Hopper or newer with CUDA 12.0+).""" + return is_sm90_supported() or is_blackwell_supported() + + +if is_blackwell_supported() and is_flashinfer_available(): + from flashinfer import SfLayout + from flashinfer import mm_mxfp8 as _raw_flashinfer_mm_mxfp8 + from flashinfer import mxfp8_quantize as _raw_flashinfer_mxfp8_quantize + from flashinfer.gemm import gemm_fp8_nt_groupwise as _raw_gemm_fp8_nt_groupwise + + from sglang.srt.utils.custom_op import register_custom_op + + @lru_cache(maxsize=1) + def _get_flashinfer_groupwise_backend() -> str: + if get_fp8_gemm_runner_backend().is_flashinfer_cutlass(): + return "cutlass" + if get_fp8_gemm_runner_backend().is_flashinfer_trtllm(): + return "trtllm" + + major, minor = get_device_capability() + # SM120/121: CUTLASS only. + # SM100/103: TRTLLM only. + if major >= 12: + return "cutlass" + return "trtllm" + + # Wrap gemm_fp8_nt_groupwise as a custom op so torch.compile does not trace + # into flashinfer's JIT compilation code (pathlib/cubin_loader ops). + @register_custom_op( + op_name="flashinfer_gemm_fp8_nt_groupwise", + mutates_args=[], + fake_impl=lambda q_input, weight, x_scale, weight_scale, out_dtype: ( + q_input.new_empty((q_input.shape[0], weight.shape[0]), dtype=out_dtype) + ), + ) + def gemm_fp8_nt_groupwise( + q_input: torch.Tensor, + weight: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, + ) -> torch.Tensor: + backend = _get_flashinfer_groupwise_backend() + if backend == "cutlass": + # FlashInfer CUTLASS groupwise kernel requires contiguous scale tensors + x_scale = x_scale.contiguous() + weight_scale = weight_scale.contiguous() + return _raw_gemm_fp8_nt_groupwise( + q_input, + weight, + x_scale, + weight_scale, + out_dtype=out_dtype, + backend="cutlass", + scale_major_mode="MN", + ) + return _raw_gemm_fp8_nt_groupwise( + q_input, + weight, + x_scale, + weight_scale, + out_dtype=out_dtype, + backend=backend, + ) + + # Wrap MXFP8 ops as custom ops so torch.compile does not trace into + # flashinfer's JIT compilation path (filesystem checks/cubin loader). + def _fake_flashinfer_mxfp8_quantize( + input: torch.Tensor, + _is_sf_swizzled_layout: bool = True, + alignment: int = 32, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Fake mode only needs dtypes and output rank to propagate compile graph. + # The scale tensor shape is not consumed before the following fake mm op. + k_aligned = ((input.shape[1] + alignment - 1) // alignment) * alignment + q_input = input.new_empty( + (input.shape[0], k_aligned), dtype=torch.float8_e4m3fn + ) + scale = input.new_empty((1,), dtype=torch.uint8) + return q_input, scale + + @register_custom_op( + op_name="flashinfer_mxfp8_quantize", + mutates_args=[], + fake_impl=_fake_flashinfer_mxfp8_quantize, + ) + def flashinfer_mxfp8_quantize( + input: torch.Tensor, + is_sf_swizzled_layout: bool = True, + alignment: int = 32, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return _raw_flashinfer_mxfp8_quantize( + input, + is_sf_swizzled_layout=is_sf_swizzled_layout, + alignment=alignment, + sf_swizzle_layout=SfLayout.layout_128x4, + ) + + @register_custom_op( + op_name="flashinfer_mm_mxfp8", + mutates_args=[], + fake_impl=lambda q_input, weight_t, x_scale_u8, weight_scale_t, out_dtype, use_8x4_sf_layout=False, backend="auto": ( + q_input.new_empty((q_input.shape[0], weight_t.shape[1]), dtype=out_dtype) + ), + ) + def flashinfer_mm_mxfp8( + q_input: torch.Tensor, + weight_t: torch.Tensor, + x_scale_u8: torch.Tensor, + weight_scale_t: torch.Tensor, + out_dtype: torch.dtype, + use_8x4_sf_layout: bool = False, + backend: str = "auto", + ) -> torch.Tensor: + return _raw_flashinfer_mm_mxfp8( + q_input, + weight_t, + x_scale_u8, + weight_scale_t, + out_dtype=out_dtype, + use_8x4_sf_layout=use_8x4_sf_layout, + backend=backend, + ) + + +if is_sm90_supported() and is_flashinfer_available(): + # FlashInfer SM90 DeepGEMM with automatic swapAB optimization for small M + from flashinfer.gemm import fp8_blockscale_gemm_sm90 + + +def dispatch_w8a8_block_fp8_linear() -> Callable: + """ + Dispatch to the appropriate FP8 block linear implementation. + + This function selects the backend based on: + 1. The --fp8-gemm-backend server argument (preferred) + 2. Auto-detection based on hardware capabilities + """ + backend = get_fp8_gemm_runner_backend() + + # Handle explicit backend selection via --fp8-gemm-backend + if not backend.is_auto(): + return _dispatch_explicit_backend(backend) + + # Auto mode: Select based purely on hardware/backend availability + return _dispatch_auto_backend() + + +def dispatch_w8a8_mxfp8_linear() -> Callable: + """Dispatch MXFP8 linear kernel by --fp8-gemm-backend. + + For MXFP8, Triton remains the default path. We only route to FlashInfer + when backend is explicitly set to flashinfer_cutlass or flashinfer_trtllm. + """ + backend = get_fp8_gemm_runner_backend() + if backend.is_flashinfer_trtllm(): + return flashinfer_mxfp8_blockscaled_linear + elif backend.is_flashinfer_cutlass(): + return flashinfer_mxfp8_blockscaled_linear + return triton_mxfp8_blockscaled_linear + + +def _dispatch_explicit_backend(backend: Fp8GemmRunnerBackend) -> Callable: + """Dispatch based on explicitly selected backend.""" + if backend.is_flashinfer_trtllm(): + if not (is_sm100_supported() and is_flashinfer_available()): + raise RuntimeError( + "FlashInfer FP8 GEMM requested via --fp8-gemm-backend=flashinfer_trtllm, " + "but FlashInfer is not available or not supported on this hardware. " + "FlashInfer TRTLLM FP8 GEMM requires SM100/SM103 GPUs and FlashInfer." + ) + return flashinfer_gemm_w8a8_block_fp8_linear_with_fallback + + elif backend.is_flashinfer_cutlass(): + if not (is_blackwell_supported() and is_flashinfer_available()): + raise RuntimeError( + "FlashInfer FP8 GEMM requested via --fp8-gemm-backend=flashinfer_cutlass, " + "but FlashInfer is not available or not supported on this hardware. " + "FlashInfer CUTLASS FP8 GEMM requires Blackwell GPUs and FlashInfer." + ) + return flashinfer_gemm_w8a8_block_fp8_linear_with_fallback + + elif backend.is_flashinfer_deepgemm(): + if not (is_sm90_supported() and is_flashinfer_available()): + raise RuntimeError( + "FlashInfer DeepGEMM with swapAB requested via --fp8-gemm-backend=flashinfer_deepgemm, " + "but it's not available. This backend requires Hopper (SM90) GPUs and FlashInfer " + "to be installed." + ) + return flashinfer_deepgemm_w8a8_block_fp8_linear_with_fallback + + elif backend.is_cutlass(): + if not _check_cutlass_block_fp8_hardware_support(): + raise RuntimeError( + "CUTLASS block FP8 requested via --fp8-gemm-backend=cutlass, " + "but hardware does not support it. CUTLASS block FP8 requires " + "Hopper (SM90+) GPUs with CUDA 12.0+." + ) + return cutlass_w8a8_block_fp8_linear_with_fallback + + elif backend.is_aiter(): + if not _use_aiter: + raise RuntimeError( + "AITER backend requested via --fp8-gemm-backend=aiter, " + "but AITER is not available. AITER requires AMD GPUs with " + "SGLANG_USE_AITER=1 environment variable set." + ) + return aiter_w8a8_block_fp8_linear + + elif backend.is_deep_gemm(): + if not deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: + raise RuntimeError( + "DeepGEMM backend requested via --fp8-gemm-backend=deep_gemm, " + "but DeepGEMM is not available. This usually means the deep_gemm package " + "is not installed or has been disabled via SGLANG_ENABLE_JIT_DEEPGEMM=0." + ) + return deepgemm_w8a8_block_fp8_linear_with_fallback + + elif backend.is_triton(): + return triton_w8a8_block_fp8_linear + + else: + raise ValueError(f"Unknown FP8 GEMM backend: {backend}") + + +def _dispatch_auto_backend() -> Callable: + """Auto-select the best backend based on hardware capabilities.""" + # Priority order for auto selection: + # 1. DeepGEMM (if enabled and available) + # 2. FlashInfer TRTLLM (if Blackwell GPU and FlashInfer available) + # 3. CUTLASS (if Hopper+ GPU and CUDA 12.0+) + # 4. AITER (if AMD GPU with AITER enabled) + # 5. Triton (fallback) + + if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: + return deepgemm_w8a8_block_fp8_linear_with_fallback + elif is_blackwell_supported() and is_flashinfer_available(): + return flashinfer_gemm_w8a8_block_fp8_linear_with_fallback + elif _check_cutlass_block_fp8_hardware_support(): + return cutlass_w8a8_block_fp8_linear_with_fallback + elif _use_aiter: + return aiter_w8a8_block_fp8_linear + else: + return triton_w8a8_block_fp8_linear + + +def initialize_fp8_gemm_config(server_args: ServerArgs) -> None: + """Initialize FP8 GEMM configuration.""" + global FP8_GEMM_RUNNER_BACKEND + + backend = server_args.fp8_gemm_runner_backend + if backend == "auto" and is_sm120_supported(): + # TODO(brayden): Verify if CUTLASS can be set by default once SwapAB is supported + backend = "triton" + + FP8_GEMM_RUNNER_BACKEND = Fp8GemmRunnerBackend(backend) + + +def get_fp8_gemm_runner_backend() -> Fp8GemmRunnerBackend: + """Get the current FP8 GEMM runner backend.""" + global FP8_GEMM_RUNNER_BACKEND + if FP8_GEMM_RUNNER_BACKEND is None: + FP8_GEMM_RUNNER_BACKEND = Fp8GemmRunnerBackend.AUTO + return FP8_GEMM_RUNNER_BACKEND + + +def flashinfer_gemm_w8a8_block_fp8_linear_with_fallback( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + + input_2d = input.view(-1, input.shape[-1]) + backend = _get_flashinfer_groupwise_backend() + # TRTLLM backend requires K dimension >= 256. + if backend == "trtllm" and input_2d.shape[1] < 256: + return triton_w8a8_block_fp8_linear( + input, weight, block_size, weight_scale, input_scale, bias + ) + + output_shape = [*input.shape[:-1], weight.shape[0]] + + # TRTLLM uses the existing SGLang column-major scale layout. + # CUTLASS with scale_major_mode="MN" expects (k//block_k, m), so we normalize below. + q_input, x_scale = sglang_per_token_group_quant_fp8( + input_2d, block_size[1], column_major_scales=(backend == "trtllm") + ) + if backend == "cutlass": + block_n, block_k = block_size + m, k = input_2d.shape + n = weight.shape[0] + expected_x_scale_shape = (k // block_k, m) + expected_weight_scale_shape = (k // block_k, n // block_n) + if x_scale.shape == (m, k // block_k): + x_scale = x_scale.transpose(-1, -2).contiguous() + if weight_scale.shape == (n // block_n, k // block_k): + weight_scale = weight_scale.transpose(-1, -2).contiguous() + assert x_scale.shape == expected_x_scale_shape, ( + "FlashInfer CUTLASS groupwise FP8 expects A scale layout " + f"(k//block_k, m) for scale_major_mode='MN', got {tuple(x_scale.shape)}; " + f"expected {expected_x_scale_shape}. " + f"strides={x_scale.stride()} is_contiguous={x_scale.is_contiguous()} " + f"m={m} n={n} k={k} block_size={block_size}" + ) + assert weight_scale.shape == expected_weight_scale_shape, ( + "FlashInfer CUTLASS groupwise FP8 expects B scale layout " + f"(k//block_k, n//block_n) for scale_major_mode='MN', got {tuple(weight_scale.shape)}; " + f"expected {expected_weight_scale_shape}. " + f"strides={weight_scale.stride()} is_contiguous={weight_scale.is_contiguous()} " + f"m={m} n={n} k={k} block_size={block_size}" + ) + assert x_scale.dtype == torch.float32, ( + "FlashInfer CUTLASS groupwise FP8 expects x_scale dtype float32, " + f"got {x_scale.dtype}." + ) + assert weight_scale.dtype == torch.float32, ( + "FlashInfer CUTLASS groupwise FP8 expects weight_scale dtype float32, " + f"got {weight_scale.dtype}." + ) + # TRTLLM path continues using the original quantized scale layout. + output = gemm_fp8_nt_groupwise( + q_input, + weight, + x_scale, + weight_scale, + out_dtype=input_2d.dtype, + ) + + if bias is not None: + output += bias + + return output.to(dtype=input_2d.dtype).view(*output_shape) + + +def flashinfer_deepgemm_w8a8_block_fp8_linear_with_fallback( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """ + FlashInfer DeepGEMM backend for SM90 (Hopper) with swapAB optimization. + + Uses flashinfer.gemm.fp8_blockscale_gemm_sm90 which automatically selects + the swapAB kernel for small M dimensions (M < 32) for better performance + during decoding/low batch size scenarios. + + For SM90 (Hopper), this uses the DeepGEMM JIT with automatic swapAB selection. + """ + assert input_scale is None + + output_dtype = input.dtype + dtype_supported = output_dtype == torch.bfloat16 + + # fp8_blockscale_gemm_sm90 requires: N % 64 == 0, K % 128 == 0 + shape_supported = weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0 + + if not (shape_supported and dtype_supported): + if weight_scale.dtype == torch.int32: + weight_scale = _unpack_ue8m0_scale_for_triton( + weight_scale, weight.shape, block_size + ) + return triton_w8a8_block_fp8_linear( + input, weight, block_size, weight_scale, input_scale, bias + ) + + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + # - input: (M, K) BF16 or FP8 + # - weight: (N, K) FP8 with weight_scale + # - weight_scale: (N, K//128) for per-token or (N//128, K//128) for per-block + + output = fp8_blockscale_gemm_sm90( + input_2d, + weight, + input_scale=None, # BF16 input, internal quantization + weight_scale=weight_scale, + out_dtype=output_dtype, + ) + + if bias is not None: + output += bias + return output.view(*output_shape) + + +def cutlass_w8a8_block_fp8_linear_with_fallback( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + + # TODO: add more robust shape check here + shape_supported = weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0 + + if not shape_supported: + # fallback to triton + return triton_w8a8_block_fp8_linear( + input, weight, block_size, weight_scale, input_scale, bias + ) + + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = per_token_group_quant_fp8( + input_2d, block_size[1], column_major_scales=True + ) + output = fp8_blockwise_scaled_mm( + q_input, weight.T, x_scale, weight_scale.T, out_dtype=input_2d.dtype + ) + if bias is not None: + output += bias + return output.to(dtype=input_2d.dtype).view(*output_shape) + + +def deepgemm_w8a8_block_fp8_linear_with_fallback( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + + output_dtype = input.dtype + dtype_supported = output_dtype == torch.bfloat16 + + # TODO: https://github.com/sgl-project/sglang/pull/6890#issuecomment-2943395737 + shape_supported = weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0 + + assert not get_bool_env_var( + "SGLANG_HACK_DEEPGEMM_W8A8_FORCE_TRITON" + ), "removed flag" + + if not (shape_supported and dtype_supported): + # fall back to triton + # If weight_scale is in UE8M0 packed format (int32), convert back to float32 + # UE8M0 format has shape (N, K//block_k//4) with dtype int32 + # Triton expects shape (N//block_n, K//block_k) with dtype float32 + if weight_scale.dtype == torch.int32: + weight_scale = _unpack_ue8m0_scale_for_triton( + weight_scale, weight.shape, block_size + ) + + assert not get_bool_env_var("SGLANG_HACK_CUSTOM_W8A8_GEMM"), "removed flag" + + return triton_w8a8_block_fp8_linear( + input, weight, block_size, weight_scale, input_scale, bias + ) + + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = sglang_per_token_group_quant_fp8( + input_2d, + block_size[1], + column_major_scales=True, + scale_tma_aligned=True, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + + output = w8a8_block_fp8_matmul_deepgemm( + q_input, weight, x_scale, weight_scale, block_size, output_dtype=output_dtype + ) + if bias is not None: + output += bias + return output.to(dtype=output_dtype).view(*output_shape) + + +def _unpack_ue8m0_scale_for_triton( + sf_packed: torch.Tensor, + weight_shape: Tuple[int, int], + block_size: List[int], +) -> torch.Tensor: + """ + Unpack UE8M0 packed scale tensor back to float32 format for triton kernel. + + The UE8M0 format packs scales as: + - Shape: (N, K//block_k//4) with dtype int32 + - Each int32 contains 4 uint8 scale values + + Triton expects: + - Shape: (N//block_n, K//block_k) with dtype float32 + + Args: + sf_packed: Packed scale tensor with shape (N, packed_k_groups) and dtype int32 + weight_shape: (N, K) shape of the weight tensor + block_size: [block_n, block_k] quantization block size + + Returns: + Unpacked scale tensor with shape (n_groups, k_groups) and dtype float32 + """ + assert sf_packed.dtype == torch.int32 + assert len(sf_packed.shape) == 2 + + N, K = weight_shape + block_n, block_k = block_size + n_groups = ceil_div(N, block_n) + k_groups = ceil_div(K, block_k) + + mn_repeat, k_div_4 = sf_packed.shape + k_packed = k_div_4 * 4 + + # Unpack int32 -> 4x uint8 -> float32 + # Each uint8 represents an exponent in UE8M0 format + sf_u8 = sf_packed.contiguous().view(torch.uint8).view(mn_repeat, k_packed) + sf_fp32 = (sf_u8.to(torch.int32) << 23).view(torch.float32) + + # Handle row dimension - may have 128x replication or direct mapping + if mn_repeat == N: + # Rows are replicated 128 times, take every 128th row + # sf_fp32 shape: (N, k_packed) -> (n_groups, k_packed) + # Select representative rows at indices 0, 128, 256, ... + indices = torch.arange(0, N, block_n, device=sf_packed.device) + sf_fp32 = sf_fp32.index_select(0, indices) + elif mn_repeat == n_groups: + # Already in the correct n_groups format + pass + else: + raise ValueError( + f"Unexpected scale shape: sf_packed.shape={sf_packed.shape}, " + f"weight_shape={weight_shape}, block_size={block_size}" + ) + + # Crop k dimension to expected size (remove padding if any) + sf_fp32 = sf_fp32[:, :k_groups].contiguous() + + return sf_fp32 + + +def aiter_w8a8_block_fp8_linear( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + # assert input_scale is None + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + # if input_scale not None, input is quanted + if input_scale is not None: + q_input = input_2d + x_scale = input_scale + + else: + q_input, x_scale = aiter_per1x128_quant(input_2d, quant_dtype=aiter.dtypes.fp8) + + n, k = weight.shape + + if _use_aiter_gfx95: + use_triton = use_aiter_triton_gemm_w8a8_tuned_gfx950(n, k) + else: + use_triton = True + + if use_triton: + gemm_a8w8_blockscale_op = triton_gemm_a8w8_blockscale + else: + gemm_a8w8_blockscale_op = gemm_a8w8_blockscale + + output = gemm_a8w8_blockscale_op( + q_input, + weight, + x_scale, + weight_scale, + dtype=torch.bfloat16 if input_scale is not None else input.dtype, + ) + + if bias is not None: + output += bias + + return output.to( + dtype=torch.bfloat16 if input_scale is not None else input_2d.dtype + ).view(*output_shape) + + +def triton_w8a8_block_fp8_linear( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert input_scale is None + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = per_token_group_quant_fp8( + input_2d, block_size[1], column_major_scales=False + ) + output = w8a8_block_fp8_matmul_triton( + q_input, weight, x_scale, weight_scale, block_size, output_dtype=input_2d.dtype + ) + if bias is not None: + output += bias + return output.to(dtype=input_2d.dtype).view(*output_shape) + + +@lru_cache(maxsize=1) +def _get_triton_mxfp8_downcast(): + try: + from triton_kernels.numerics_details.mxfp import downcast_to_mxfp + except Exception as err: + raise RuntimeError( + "MXFP8 quantization requires triton_kernels with MXFP8 support." + ) from err + return downcast_to_mxfp + + +def mxfp8_group_quantize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2D contiguous tensor to MXFP8 with UE8M0 scales per group (32).""" + assert x.dim() == 2, f"Expected 2D input, got {x.dim()}D" + assert x.is_contiguous(), "MXFP8 quantization requires a contiguous 2D tensor." + _, k = x.shape + assert k % 32 == 0, f"{k=} must be divisible by 32" + downcast_to_mxfp = _get_triton_mxfp8_downcast() + q_input, scale_u8 = downcast_to_mxfp(x, torch.float8_e4m3fn, axis=1) + return q_input.contiguous(), scale_u8.contiguous() + + +def _pack_mxfp8_scales(scale_u8: torch.Tensor) -> torch.Tensor: + # Pack (M, K//32) UE8M0 scales into the layout expected by tl.dot_scaled. + assert scale_u8.dim() == 2, f"Expected 2D scale tensor, got {scale_u8.dim()}D" + scale_u8 = scale_u8.contiguous() + m, k_groups = scale_u8.shape + assert ( + k_groups % 4 == 0 + ), f"{k_groups=} must be divisible by 4 (K must be multiple of 128)" + + scale_m = ceil_div(m, 128) + if m % 128 != 0: + pad_rows = scale_m * 128 - m + pad = torch.full( + (pad_rows, k_groups), + 127, + dtype=scale_u8.dtype, + device=scale_u8.device, + ) + scale_u8 = torch.cat([scale_u8, pad], dim=0) + + scale_k = k_groups // 4 + scale_u8 = scale_u8.view(scale_m, 128, scale_k, 4) + scale_u8 = scale_u8.view(scale_m, 4, 32, scale_k, 4) + packed = scale_u8.permute(0, 3, 2, 1, 4).contiguous() + return packed.view(1, scale_m, scale_k, 2, 256) + + +@register_custom_op( + op_name="triton_mxfp8_block_scaled_matmul", + mutates_args=[], + fake_impl=lambda a, a_scale, b, b_scale, output_dtype, block_m=128, block_n=256, block_k=128, num_stages=None: ( # noqa: E501 + a.new_empty((a.shape[0], b.shape[0]), dtype=output_dtype) + ), +) +def triton_mxfp8_block_scaled_matmul( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + output_dtype: torch.dtype, + *, + block_m: int = 128, + block_n: int = 256, + block_k: int = 128, + num_stages: Optional[int] = None, +) -> torch.Tensor: + """Opaque custom op wrapper to prevent Dynamo tracing Triton grid math.""" + return mxfp8_block_scaled_matmul_triton( + a, + a_scale, + b, + b_scale, + output_dtype=output_dtype, + block_m=block_m, + block_n=block_n, + block_k=block_k, + num_stages=num_stages, + ) + + +def _raw_triton_mxfp8_blockscaled_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + output_dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + if not (_is_cuda and (_is_sm100_supported or _is_sm120_supported)): + raise RuntimeError("MXFP8 dense linear requires Blackwell GPUs (SM100/SM120).") + + input_2d = input.view(-1, input.shape[-1]).contiguous() + output_shape = [*input.shape[:-1], weight.shape[0]] + + block_m = 128 + block_n = 256 if weight.shape[0] % 256 == 0 else 128 + block_k = 128 + + m, k = input_2d.shape + n, k_w = weight.shape + assert k == k_w, f"{k=} does not match {k_w=}" + assert k % 128 == 0, f"{k=} must be divisible by 128 for MXFP8" + assert n % block_n == 0, f"{n=} must be divisible by {block_n}" + assert weight.dtype == torch.float8_e4m3fn, "MXFP8 weight must be FP8 E4M3." + assert weight_scale.dtype == torch.uint8, "MXFP8 weight_scale must be UE8M0 uint8." + + if input_scale is None: + q_input, x_scale_u8 = mxfp8_group_quantize(input_2d) + else: + q_input = input_2d + x_scale_u8 = input_scale + assert x_scale_u8.dtype == torch.uint8, "MXFP8 input_scale must be UE8M0 uint8." + assert x_scale_u8.shape == (m, k // 32) + + if output_dtype is None: + if input_2d.dtype in (torch.float16, torch.bfloat16, torch.float32): + output_dtype = input_2d.dtype + else: + output_dtype = torch.bfloat16 + + if m % block_m != 0: + pad_rows = ceil_div(m, block_m) * block_m - m + q_input = torch.cat( + [ + q_input, + torch.zeros((pad_rows, k), device=q_input.device, dtype=q_input.dtype), + ], + dim=0, + ) + pad_scale = torch.full( + (pad_rows, k // 32), + 127, + device=x_scale_u8.device, + dtype=x_scale_u8.dtype, + ) + x_scale_u8 = torch.cat([x_scale_u8, pad_scale], dim=0) + + a_scale_packed = _pack_mxfp8_scales(x_scale_u8) + b_scale_packed = _pack_mxfp8_scales(weight_scale) + + num_stages = 1 if _is_sm120_supported else (4 if _is_sm100_supported else 1) + output = triton_mxfp8_block_scaled_matmul( + q_input, + a_scale_packed, + weight.contiguous(), + b_scale_packed, + output_dtype=output_dtype, + block_m=block_m, + block_n=block_n, + block_k=block_k, + num_stages=num_stages, + ) + output = output[:m, :] + if bias is not None: + output += bias + return output.to(dtype=output_dtype).view(*output_shape) + + +@register_custom_op( + op_name="triton_mxfp8_blockscaled_linear", + mutates_args=[], + fake_impl=lambda input, weight, weight_scale, input_scale=None, bias=None, output_dtype=None: ( + input.new_empty( + (*input.shape[:-1], weight.shape[0]), + dtype=(output_dtype if output_dtype is not None else input.dtype), + ) + ), +) +def triton_mxfp8_blockscaled_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + output_dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + """Opaque custom-op wrapper to prevent Dynamo guards on MXFP8 padding branches.""" + return _raw_triton_mxfp8_blockscaled_linear( + input=input, + weight=weight, + weight_scale=weight_scale, + input_scale=input_scale, + bias=bias, + output_dtype=output_dtype, + ) + + +def flashinfer_mxfp8_blockscaled_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + output_dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + """MXFP8 dense linear via FlashInfer mm_mxfp8.""" + input_2d = input.view(-1, input.shape[-1]).contiguous() + output_shape = [*input.shape[:-1], weight.shape[0]] + + m, k = input_2d.shape + n, k_w = weight.shape + if k != k_w: + raise ValueError(f"Input K={k} does not match weight K={k_w}.") + if k % 32 != 0: + raise ValueError(f"K={k} must be divisible by 32 for MXFP8.") + if weight.dtype != torch.float8_e4m3fn: + raise TypeError("MXFP8 weight must be FP8 E4M3.") + + if input_scale is None: + q_input, x_scale_u8 = flashinfer_mxfp8_quantize( + input_2d, is_sf_swizzled_layout=True, alignment=32 + ) + else: + q_input = input_2d + x_scale_u8 = input_scale.contiguous() + + if output_dtype is None: + if input_2d.dtype in (torch.float16, torch.bfloat16, torch.float32): + output_dtype = input_2d.dtype + else: + output_dtype = torch.bfloat16 + + # Ensure transposed tensors are contiguous for FlashInfer's internal runner. + weight_t = weight.contiguous().t() + + if get_fp8_gemm_runner_backend().is_flashinfer_trtllm(): + + weight_scale_t = weight_scale.contiguous().view(-1) + output = flashinfer_mm_mxfp8( + q_input, + weight_t, + x_scale_u8, + weight_scale_t, + out_dtype=output_dtype, + use_8x4_sf_layout=False, + backend="trtllm", + ) + elif get_fp8_gemm_runner_backend().is_flashinfer_cutlass(): + weight_scale_t = ( + weight_scale.contiguous().t() + if weight_scale.ndim == 2 + else weight_scale.contiguous() + ) + output = flashinfer_mm_mxfp8( + q_input, + weight_t, + x_scale_u8, + weight_scale_t, + out_dtype=output_dtype, + use_8x4_sf_layout=False, + backend="cutlass", + ) + + if bias is not None: + output += bias + return output.to(dtype=output_dtype).view(*output_shape) + + +def dequant_mxfp4( + w_block: torch.Tensor, + w_scale: torch.Tensor, + out_dtype, +) -> torch.Tensor: + """ + :param w_block: (batch, n, k, 16), uint8, pack two mxfp4 into one byte + :param w_scale: (batch, n, k), uint8 + :return: (batch, n, k * 32), float32 + """ + + assert w_block.dtype == torch.uint8 + assert w_scale.dtype == torch.uint8 + + batch, n, k, pack_dim = w_block.shape + batch_, n_, k_ = w_scale.shape + assert pack_dim == 16 + assert batch == batch_ + assert n == n_ + assert k == k_ + + out_raw = MXFP4QuantizeUtil.dequantize( + quantized_data=w_block, scale=w_scale, dtype=out_dtype, block_sizes=[32] + ) + return out_raw.reshape(batch, n, k * 32) + + +def input_to_float8( + x: torch.Tensor, dtype: torch.dtype = fp8_dtype +) -> Tuple[torch.Tensor, torch.Tensor]: + """This function quantizes input values to float8 values with tensor-wise quantization.""" + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).float().clamp(min=1e-12) + + if _is_fp8_fnuz: + dtype = fp8_dtype + fp_max = fp8_max + else: + finfo = torch.finfo(dtype) + fp_max = finfo.max + + scale = fp_max / amax + x_scl_sat = (x.float() * scale).clamp(min=-fp_max, max=fp_max) + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() + + +def block_quant_to_tensor_quant( + x_q_block: torch.Tensor, + x_s: torch.Tensor, + block_size: List[int], +) -> Tuple[torch.Tensor, torch.Tensor]: + """This function converts block-wise quantization to tensor-wise quantization. + The inputs are block-wise quantization tensor `x_q_block`, block-wise quantization scale + and the block size. + The outputs are tensor-wise quantization tensor and tensor-wise quantization scale. + Note only float8 is supported for now. + """ + block_n, block_k = block_size[0], block_size[1] + n, k = x_q_block.shape + n_tiles = (n + block_n - 1) // block_n + k_tiles = (k + block_k - 1) // block_k + assert n_tiles == x_s.shape[0] + assert k_tiles == x_s.shape[1] + + x_dq_block = x_q_block.to(torch.float32) + + x_dq_block_tiles = [ + [ + x_dq_block[ + j * block_n : min((j + 1) * block_n, n), + i * block_k : min((i + 1) * block_k, k), + ] + for i in range(k_tiles) + ] + for j in range(n_tiles) + ] + + for i in range(k_tiles): + for j in range(n_tiles): + x_dq_block_tiles[j][i][:, :] = x_dq_block_tiles[j][i] * x_s[j][i] + + x_q_tensor, scale = ( + scaled_fp8_quant(x_dq_block) + if _is_cuda + else input_to_float8(x_dq_block, dtype=x_q_block.dtype) + ) + return x_q_tensor, scale + + +def block_quant_dequant( + x_q_block: torch.Tensor, + x_s: torch.Tensor, + block_size: List[int], + dtype: torch.dtype, +) -> torch.Tensor: + """This function converts block-wise quantization to unquantized. + The inputs are block-wise quantization tensor `x_q_block`, block-wise quantization scale + and the block size. + The output is an unquantized tensor with dtype. + """ + block_n, block_k = block_size[0], block_size[1] + *_, n, k = x_q_block.shape + + # ... n_scale k_scale -> ... (n_scale block_n) (k_scale block_k) + x_scale_repeat = x_s.repeat_interleave(block_n, dim=-2).repeat_interleave( + block_k, dim=-1 + ) + x_scale_repeat = x_scale_repeat[..., :n, :k] + + return (x_q_block.to(torch.float32) * x_scale_repeat).to(dtype) + + +def requant_weight_ue8m0_inplace(weight, weight_scale_inv, weight_block_size): + assert isinstance(weight, torch.nn.Parameter) + assert isinstance(weight_scale_inv, torch.nn.Parameter) + + new_weight, new_weight_scale_inv = requant_weight_ue8m0( + weight.to(weight_scale_inv.device), weight_scale_inv, weight_block_size + ) + + offloader.update_param(weight, new_weight) + weight_scale_inv.data = new_weight_scale_inv + + +def requant_weight_ue8m0( + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + weight_block_size: List[int], +): + assert weight_block_size == [128, 128] + + *_, n, k = weight.shape + + weight_dequant = block_quant_dequant( + weight, + weight_scale_inv, + weight_block_size, + torch.bfloat16, + ) + + out_w, out_s = quant_weight_ue8m0( + weight_dequant=weight_dequant, + weight_block_size=weight_block_size, + ) + out_s = transform_scale_ue8m0(out_s, mn=out_w.shape[-2]) + + return out_w, out_s + + +def quant_weight_ue8m0( + weight_dequant: torch.Tensor, + weight_block_size: List[int], +): + assert weight_block_size == [128, 128] + assert ( + weight_dequant.dtype == torch.bfloat16 + ), f"{weight_dequant.dtype=} {weight_dequant.shape=}" + + *batch_dims, n, k = weight_dequant.shape + + weight_dequant_flat = weight_dequant.view((-1, k)) + out_w_flat, out_s_flat = per_block_cast_to_fp8(weight_dequant_flat) + + out_w = out_w_flat.view((*batch_dims, n, k)) + out_s = out_s_flat.view( + ( + *batch_dims, + ceil_div(n, weight_block_size[0]), + ceil_div(k, weight_block_size[1]), + ) + ) + + return out_w, out_s + + +def transform_scale_ue8m0_inplace(param, mn): + param.data = transform_scale_ue8m0(param.data, mn=mn) + + +# NOTE copy and modified from DeepGEMM +def transform_scale_ue8m0(sf, mn, use_torch_impl: bool = False): + import deep_gemm.utils.layout + + get_mn_major_tma_aligned_packed_ue8m0_tensor = ( + _get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl + if use_torch_impl + else deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor + ) + + sf = sf.index_select(-2, torch.arange(mn, device=sf.device) // 128) + sf = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf) + return sf + + +# Copied from DeepGEMM tests +def _get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl( + x: torch.Tensor, +) -> torch.Tensor: + from deep_gemm.utils import align, get_tma_aligned_size + + assert x.dtype == torch.float and x.dim() in (2, 3) + + # First, convert into UE8M0 `uint8_t` + ue8m0_tensor = (x.view(torch.int) >> 23).to(torch.uint8) + + # Second, make padded packed tensors + mn, k = x.shape[-2], x.shape[-1] + remove_dim = False + if x.dim() == 2: + x, remove_dim = x.unsqueeze(0), True + b = x.shape[0] + aligned_mn = get_tma_aligned_size(mn, 4) + aligned_k = align(k, 4) + padded = torch.zeros((b, aligned_mn, aligned_k), device=x.device, dtype=torch.uint8) + padded[:, :mn, :k] = ue8m0_tensor + padded = padded.view(-1).view(dtype=torch.int).view(b, aligned_mn, aligned_k // 4) + + # Finally, transpose + transposed = torch.zeros( + (b, aligned_k // 4, aligned_mn), device=x.device, dtype=torch.int + ).mT + transposed[:, :, :] = padded + aligned_x = transposed[:, :mn, :] + return aligned_x.squeeze(0) if remove_dim else aligned_x + + +def inverse_transform_scale_ue8m0(sf_packed, mn): + sf_fp32 = _inverse_transform_scale_ue8m0_impl(sf_packed) + # Can call consistency check every time since this is only called on startup + sf_packed_recreated = transform_scale_ue8m0(sf_fp32, mn=mn, use_torch_impl=True) + assert torch.all( + sf_packed == sf_packed_recreated + ), f"{sf_packed=} {sf_packed_recreated=} {sf_fp32=}" + return sf_fp32 + + +# Inverse impl can refer to DeepGEMM's torch impl in get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl +def _inverse_transform_scale_ue8m0_impl(sf_packed): + """ + NOTE: We assume k is aligned + :param sf_packed: (scale_mn, scale_k/4) int32 + :return: (scale_mn, scale_k), float32 + """ + if len(sf_packed.shape) == 3: + return torch.stack( + [_inverse_transform_scale_ue8m0_impl(x) for x in sf_packed], dim=0 + ) + + block_size = 128 + assert len(sf_packed.shape) == 2, f"{sf_packed.shape=}" + assert sf_packed.dtype == torch.int32 + + mn_repeat_128, k_div_4 = sf_packed.shape + mn = mn_repeat_128 // block_size + k = k_div_4 * 4 + + # packed u8 -> fp32 + sf_u8 = sf_packed.contiguous().flatten().view(torch.uint8).view(mn_repeat_128, k) + sf_fp32 = (sf_u8.to(torch.int32) << 23).view(torch.float32) + + # remove repeat + sf_reshaped = sf_fp32.view(mn, block_size, k) + sf_unrepeated = sf_reshaped[:, 0:1, :] + if not torch.all(sf_unrepeated == sf_reshaped): + from sglang.srt.debug_utils.dumper import get_tensor_info + + raise AssertionError( + f"sf_unrepeated != sf_reshaped ({get_tensor_info(sf_unrepeated)=} {get_tensor_info(sf_reshaped)=})" + ) + sf_unrepeated = sf_unrepeated.squeeze(1).contiguous() + + assert sf_unrepeated.shape == (mn, k) + return sf_unrepeated + + +# COPIED FROM DeepGEMM +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_align(m, 128), ceil_align(n, 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) + sf = ceil_to_ue8m0(x_amax / 448.0) + x_scaled = (x_view * (1.0 / sf)).to(torch.float8_e4m3fn) + return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view( + x_view.size(0), x_view.size(2) + ) + + +# COPIED FROM DeepGEMM +def ceil_to_ue8m0(x: torch.Tensor): + return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) + + +def channel_quant_to_tensor_quant( + x_q_channel: torch.Tensor, + x_s: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + x_dq_channel = x_q_channel.to(torch.float32) * x_s + x_q_tensor, scale = ( + scaled_fp8_quant(x_dq_channel) + if _is_cuda + else input_to_float8(x_dq_channel, dtype=x_q_channel.dtype) + ) + return x_q_tensor, scale + + +def _process_scaled_mm_output(output, input_2d_shape, output_shape): + if type(output) is tuple and len(output) == 2: + output = output[0] + return torch.narrow(output, 0, 0, input_2d_shape[0]).view(*output_shape) + + +def _apply_fallback_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + input_2d_shape, + output_shape, + bias, + input_dtype, +): + global TORCH_DEVICE_IDENTITY + if TORCH_DEVICE_IDENTITY is None: + TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32, device=weight.device) + + output = torch._scaled_mm( + qinput, + weight, + scale_a=TORCH_DEVICE_IDENTITY, + scale_b=TORCH_DEVICE_IDENTITY, + out_dtype=torch.float32, + ) + + output = _process_scaled_mm_output(output, input_2d_shape, output_shape) + x_scale = torch.narrow(x_scale, 0, 0, input_2d_shape[0]) + + output = output * x_scale * weight_scale.t() + if bias is not None: + output = output + bias + return output.to(dtype=input_dtype) + + +def apply_fp8_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + input_scale_ub: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + cutlass_fp8_supported: bool = cutlass_fp8_supported(), + use_per_token_if_dynamic: bool = False, + pad_output: Optional[bool] = None, + compressed_tensor_quant: bool = False, +) -> torch.Tensor: + # Note: we pad the input because torch._scaled_mm is more performant + # for matrices with batch dimension > 16. + # This could change in the future. + # We also don't pad when using torch.compile, + # as it breaks with dynamic shapes. + if pad_output is None: + pad_output = not cutlass_fp8_supported and not get_bool_env_var( + "SGLANG_ENABLE_TORCH_COMPILE" + ) + output_padding = 17 if pad_output else None + + # View input as 2D matrix for fp8 methods + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[1]] + + if compressed_tensor_quant: + # Maybe apply padding to output, see comment in __init__ + num_token_padding = output_padding + if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]: + num_token_padding = None + qinput, x_scale = scaled_fp8_quant( + input_2d, + input_scale, + num_token_padding=num_token_padding, + use_per_token_if_dynamic=use_per_token_if_dynamic, + ) + else: + # cutlass w8a8 fp8 sgl-kernel only supports per-token scale + if input_scale is not None: + assert input_scale.numel() == 1 + # broadcast per-tensor scale to per-token scale when supporting cutlass + qinput, x_scale = static_quant_fp8( + input_2d, input_scale, repeat_scale=cutlass_fp8_supported + ) + else: + # default use per-token quantization if dynamic + if _is_cuda: + qinput, x_scale = sglang_per_token_quant_fp8(input_2d) + else: + # TODO(kkhuang): temporarily enforce per-tensor activation scaling if weight is per-tensor scaling + # final solution should be: 1. add support to per-tensor activation scaling. + # 2. solve the torch.compile error from weight_scale.numel() == 1 and x_scale.numel() > 1 (below line#308) + if _is_hip and weight_scale.numel() == 1: + qinput, x_scale = scaled_fp8_quant( + input_2d, + input_scale, + use_per_token_if_dynamic=use_per_token_if_dynamic, + ) + else: + qinput, x_scale = per_token_group_quant_fp8( + input_2d, group_size=input_2d.shape[1] + ) + + if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]: + cutlass_compatible_b = weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0 + if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel: + # Massage the input to be 2D + qinput = qinput.view(-1, qinput.shape[-1]) + output = triton_scaled_mm( + qinput, weight, x_scale, weight_scale, input.dtype, bias + ) + else: + output = fp8_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + out_dtype=input.dtype, + bias=bias, + ) + return output.view(*output_shape) + + # torch.scaled_mm supports per tensor weights + activations only + # so fallback to naive if per channel or per token + per_tensor_weights = weight_scale.numel() == 1 + # When the number of token is 1, + # per-token scale has shape (1, 1), per-tensor scale has shape (1) or (). + per_tensor_activations = (x_scale.numel() == 1) and x_scale.dim() < 2 + + if ( + use_per_token_if_dynamic + and not per_tensor_weights + and not per_tensor_activations + and (USE_ROWWISE_TORCH_SCALED_MM or _use_aiter) + ): + # into this sector means use dynamic per-token-per-channel quant + # per-token scale quant for input matrix, every row(one token) have one scale factor + # per-channel scale quant for weight matrix, every col(one channel) have one scale factor + if _use_aiter: + # gemm_a8w8_bpreshuffle(XQ, WQ, x_scale, w_scale, dtype) + # XQ -> input tensor, shape = (m, k) + # WQ -> weight tensor, shape = (n, k), with preshuffe get better perf + # x_scale -> input scale tensor, shape = (m, 1) + # w_scale -> weight scale tensor, shape = (n ,1) + # dtype -> output dtype + output = gemm_a8w8_bpreshuffle( + XQ=qinput, + WQ=weight.T, + x_scale=x_scale, + w_scale=weight_scale, + dtype=input.dtype, + ) + if bias is not None: + output += bias + return _process_scaled_mm_output(output, input_2d.shape, output_shape) + else: + # For now validated on ROCm platform + # fp8 rowwise scaling in torch._scaled_mm is introduced in + # https://github.com/pytorch/pytorch/pull/144432 using hipBLASLt + # and ROCm 6.3, which only exists in torch 2.7 and above. + # For CUDA platform please validate if the + # torch._scaled_mm support rowwise scaled GEMM + # Fused GEMM_DQ Rowwise GEMM + output = torch._scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale.t(), + bias=bias, + ) + return _process_scaled_mm_output(output, input_2d.shape, output_shape) + + if per_tensor_weights and per_tensor_activations: + # Fused GEMM_DQ; _scaled_mm with torch.compile requires len(weight_scale.shape) == len(x_scale.shape) + if weight_scale.ndim == 0 and x_scale.ndim == 1: + weight_scale = weight_scale.unsqueeze(0) + output = torch._scaled_mm( + qinput, + weight, + out_dtype=input.dtype, + scale_a=x_scale, + scale_b=weight_scale, + bias=bias, + ) + return _process_scaled_mm_output(output, input_2d.shape, output_shape) + + # Fallback for channelwise case, where we use unfused DQ + # due to limitations with scaled_mm + + # Symmetric quantized GEMM by definition computes the following: + # C = (s_x * X) (s_w * W) + bias + # This is equivalent to dequantizing the weights and activations + # before applying a GEMM. + # + # In order to compute quantized operands, a quantized kernel + # will rewrite the above like so: + # C = s_w * s_x * (X * W) + bias + # + # For the scaled_mm fallback case, we break this down, since it + # does not support s_w being a vector. + return _apply_fallback_scaled_mm( + qinput, + weight, + x_scale, + weight_scale, + input_2d.shape, + output_shape, + bias, + input.dtype, + ) + + +def can_auto_enable_marlin_fp8() -> bool: + # P800 (Kunlunxin) reports an NVIDIA-compatible SM but does not support + # the Marlin FP8 repacking path. Always return False to fall back to + # the native FP8 / xtorch_ops kernels. + return False + + +def apply_fp8_ptpc_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + input_scale_ub: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + cutlass_fp8_supported: bool = cutlass_fp8_supported(), + use_per_token_if_dynamic: bool = False, + pad_output: Optional[bool] = None, + compressed_tensor_quant: bool = False, +) -> torch.Tensor: + # View input as 2D matrix for fp8 methods + input_2d = input.view(-1, input.shape[-1]) + + # weight is transposed (K, N) + output_shape = [*input.shape[:-1], weight.shape[1]] + + q_input, x_scale = aiter.per_token_quant_hip(input_2d, quant_dtype=aiter.dtypes.fp8) + + per_tensor_weights = (weight_scale.numel() == 1) and weight_scale.dim() < 2 + per_tensor_activations = (x_scale.numel() == 1) and x_scale.dim() < 2 + + if not (per_tensor_weights and per_tensor_activations): + # weight is in (N, K) + output_shape = [*input.shape[:-1], weight.shape[0]] + + output = aiter.gemm_a8w8_bpreshuffle( + q_input, weight, x_scale, weight_scale, None, input.dtype + ) + if bias is not None: + output = output + bias + return output.view(*output_shape) + + +def validate_fp8_block_shape( + layer: torch.nn.Module, + input_size: int, + output_size: int, + input_size_per_partition: int, + output_partition_sizes: list[int], + block_size: list[int], +) -> None: + """Validate block quantization shapes for tensor parallelism.""" + from sglang.srt.distributed import get_tensor_model_parallel_world_size + + tp_size = getattr(layer, "tp_size", get_tensor_model_parallel_world_size()) + block_n, block_k = block_size[0], block_size[1] + + # Required by row parallel + if ( + tp_size > 1 + and input_size // input_size_per_partition == tp_size + and input_size_per_partition % block_k != 0 + ): + raise ValueError( + f"Weight input_size_per_partition = {input_size_per_partition} " + f"is not divisible by weight quantization block_k = {block_k}." + ) + + # Required by column parallel or enabling merged weights + is_tp_split = tp_size > 1 and output_size // sum(output_partition_sizes) == tp_size + is_merged_gemm = len(output_partition_sizes) > 1 + if is_tp_split or is_merged_gemm: + sizes_to_check = output_partition_sizes + if not is_tp_split and is_merged_gemm: + # In case of merged matrices, we allow the last + # matrix to not be a multiple of block size + sizes_to_check = output_partition_sizes[:-1] + for output_partition_size in sizes_to_check: + if output_partition_size % block_n != 0: + raise ValueError( + f"Weight output_partition_size = " + f"{output_partition_size} is not divisible by " + f"weight quantization block_n = {block_n}." + ) diff --git a/platforms/patches/kunlun_p800/hf_transformers_utils.py.patched b/platforms/patches/kunlun_p800/hf_transformers_utils.py.patched new file mode 100644 index 0000000..91901f3 --- /dev/null +++ b/platforms/patches/kunlun_p800/hf_transformers_utils.py.patched @@ -0,0 +1,1451 @@ +# 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 json +import logging +import os +import tempfile +import warnings +from functools import lru_cache +from pathlib import Path +from typing import Any, Dict, List, Literal, Optional, Type, Union + +import torch +from huggingface_hub import snapshot_download + +from sglang.srt.environ import envs +from sglang.srt.utils import get_bool_env_var +from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri + +# Compatibility shim: flash-attn-4 registers a bare ``flash_attn`` namespace +# that makes ``is_flash_attn_2_available()`` return True, but lacks the v2 API +# (``flash_attn_func``, etc.). HuggingFace remote model code (e.g. Kimi-VL) +# guarded by that check will crash with ImportError at module load time. +# Force it to False when the real v2 API is absent. +try: + import flash_attn as _flash_attn_mod + + if not hasattr(_flash_attn_mod, "flash_attn_func"): + import transformers.utils as _hf_utils + import transformers.utils.import_utils as _hf_import_utils + + _hf_import_utils.is_flash_attn_2_available = lambda: False + _hf_utils.is_flash_attn_2_available = lambda: False + del _flash_attn_mod +except ImportError: + pass + +# Conditional import based on SGLANG_USE_MODELSCOPE environment variable +if get_bool_env_var("SGLANG_USE_MODELSCOPE"): + from modelscope import AutoConfig, GenerationConfig +else: + from transformers import AutoConfig, GenerationConfig + +from transformers import ( + AutoProcessor, + AutoTokenizer, + PretrainedConfig, + PreTrainedTokenizer, + PreTrainedTokenizerBase, + PreTrainedTokenizerFast, +) +from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES + +from sglang.srt.configs import ( + AfmoeConfig, + BailingHybridConfig, + ChatGLMConfig, + DbrxConfig, + DeepseekVL2Config, + DotsOCRConfig, + DotsVLMConfig, + ExaoneConfig, + FalconH1Config, + GraniteMoeHybridConfig, + JetNemotronConfig, + JetVLMConfig, + KimiK25Config, + KimiLinearConfig, + KimiVLConfig, + LongcatFlashConfig, + MultiModalityConfig, + NemotronH_Nano_VL_V2_Config, + NemotronHConfig, + Olmo3Config, + Qwen3_5Config, + Qwen3_5MoeConfig, + Qwen3NextConfig, + Step3p5Config, + Step3VLConfig, +) +from sglang.srt.configs.deepseek_ocr import DeepseekVLV2Config +from sglang.srt.configs.internvl import InternVLChatConfig +from sglang.srt.connector import create_remote_connector +from sglang.srt.multimodal.customized_mm_processor_utils import _CUSTOMIZED_MM_PROCESSOR +from sglang.srt.utils import is_remote_url, logger, lru_cache_frozenset, mistral_utils + +_CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [ + AfmoeConfig, + BailingHybridConfig, + ChatGLMConfig, + DbrxConfig, + ExaoneConfig, + DeepseekVL2Config, + MultiModalityConfig, + KimiVLConfig, + InternVLChatConfig, + Step3VLConfig, + LongcatFlashConfig, + Olmo3Config, + KimiLinearConfig, + Qwen3NextConfig, + FalconH1Config, + GraniteMoeHybridConfig, + DotsVLMConfig, + DotsOCRConfig, + NemotronH_Nano_VL_V2_Config, + NemotronHConfig, + DeepseekVLV2Config, + Qwen3_5Config, + Qwen3_5MoeConfig, + JetNemotronConfig, + JetVLMConfig, + KimiK25Config, + Step3p5Config, +] + +_CONFIG_REGISTRY = { + config_cls.model_type: config_cls for config_cls in _CONFIG_REGISTRY +} + +for name, cls in _CONFIG_REGISTRY.items(): + with contextlib.suppress(ValueError): + AutoConfig.register(name, cls) + + +def download_from_hf( + model_path: str, + allow_patterns: Optional[Union[str, list]] = None, +): + if os.path.exists(model_path): + return model_path + + if not allow_patterns: + allow_patterns = ["*.json", "*.bin", "*.model"] + + return snapshot_download(model_path, allow_patterns=allow_patterns) + + +def get_rope_config(config): + """Get (rope_theta, rope_scaling) from config, supporting both v4 and v5. + + In transformers v5, rope_theta/rope_scaling are accessed via the computed + property config.rope_parameters. Trust-remote-code configs or parent configs + passed to sub-models may not have this property or may return None. + Falls back to the v4-style config.rope_theta / config.rope_scaling attributes. + """ + rope_params = getattr(config, "rope_parameters", None) + if rope_params is not None and "rope_theta" in rope_params: + return rope_params["rope_theta"], rope_params + rope_theta = getattr(config, "rope_theta", None) + if rope_theta is None and rope_params is not None: + rope_theta = rope_params.get("rope_theta", None) + return rope_theta, getattr(config, "rope_scaling", None) + + +def _patch_text_config(parent_config: PretrainedConfig, text_config): + """Synchronize standard attributes between parent config and text sub-config. + + In transformers v5, the "untangle config" refactor removed automatic + inheritance of top-level PretrainedConfig attributes (pad_token_id, + tie_word_embeddings, etc.) from sub-configs. Downstream code expects + these attributes to be present on both configs (some models pass the + parent directly to the language model, others pass the text sub-config), + so we propagate in both directions when an attribute is missing. + (See https://github.com/huggingface/transformers/pull/41541) + """ + _ATTRS_TO_PROPAGATE = [ + "pad_token_id", + "bos_token_id", + "eos_token_id", + "tie_word_embeddings", + ] + for attr in _ATTRS_TO_PROPAGATE: + parent_has = hasattr(parent_config, attr) + text_has = hasattr(text_config, attr) + if parent_has and not text_has: + setattr(text_config, attr, getattr(parent_config, attr)) + elif text_has and not parent_has: + setattr(parent_config, attr, getattr(text_config, attr)) + return text_config + + +def get_hf_text_config(config: PretrainedConfig): + """Get the "sub" config relevant to llm for multi modal models. + No op for pure text models. + """ + if config.architectures is not None: + class_name = config.architectures[0] + if class_name.startswith("Llava") and class_name.endswith("ForCausalLM"): + # We support non-hf version of llava models, so we do not want to + # read the wrong values from the unused default text_config. + # NOTE(HandH1998): We set `torch_dtype` of config to `torch.float16` for the weights, as + # `torch.float16` is default used for image features in `python/sglang/srt/models/llava.py`. + setattr(config, "dtype", torch.float16) + return config + + text_config = None + + # Some models (e.g. DeepSeek-OCR) store sub-configs as plain dicts. + # Convert to PretrainedConfig early so hasattr() checks and asserts work. + for _attr in ("text_config", "llm_config", "language_config", "thinker_config"): + _sub = getattr(config, _attr, None) + if isinstance(_sub, dict): + _converted = PretrainedConfig(**_sub) + # Propagate torch_dtype from parent so weight loading uses correct precision. + if ( + getattr(_converted, "torch_dtype", None) is None + and getattr(config, "torch_dtype", None) is not None + ): + _converted.torch_dtype = config.torch_dtype + setattr(config, _attr, _converted) + + # Priority: thinker_config > llm_config > language_config > text_config + if hasattr(config, "thinker_config"): + # qwen2.5 omni + thinker_config = config.thinker_config + if hasattr(thinker_config, "text_config"): + setattr( + thinker_config.text_config, + "torch_dtype", + getattr(thinker_config, "torch_dtype", None), + ) + text_config = thinker_config.text_config + else: + text_config = thinker_config + elif hasattr(config, "llm_config"): + # PointsV1.5 Chat Model + assert hasattr(config.llm_config, "num_attention_heads") + text_config = config.llm_config + elif hasattr(config, "language_config"): + text_config = config.language_config + elif hasattr(config, "text_config"): + # The code operates under the assumption that text_config should have + # `num_attention_heads` (among others). Assert here to fail early + # if transformers config doesn't align with this assumption. + assert hasattr(config.text_config, "num_attention_heads") + text_config = config.text_config + + # Ensure rope_scaling dicts have "type" for remote-code compat (v5). + normalize_rope_scaling_compat(config) + + if text_config is not None: + return _patch_text_config(config, text_config) + return config + + +# Temporary hack for DeepSeek-V3.2 model +def _load_deepseek_temp_model( + model_path: str, + model_type: Literal["deepseek_v32", "deepseek_ref"], + architecture: Literal["DeepseekV3ForCausalLM", "DeepseekV4ForCausalLM"], + trust_remote_code: bool = False, + revision: Optional[str] = None, + **kwargs, +): + # first get the local path + local_path = download_from_hf(model_path) + # then load the config file in json + backup_mode = envs.SGLANG_APPLY_CONFIG_BACKUP.get() + if backup_mode == "auto": + real_config_file = os.path.join(local_path, "config.json") + if not os.path.exists(real_config_file): + raise RuntimeError( + f"SGLANG_APPLY_CONFIG_BACKUP=auto requires the checkpoint's " + f"config.json at {real_config_file} to read num_hidden_layers." + ) + with open(real_config_file, "r") as f: + num_hidden_layers = json.load(f).get("num_hidden_layers") + if not isinstance(num_hidden_layers, int): + raise RuntimeError( + f"SGLANG_APPLY_CONFIG_BACKUP=auto could not read a numeric " + f"num_hidden_layers from {real_config_file} (got {num_hidden_layers!r})." + ) + backup_mode = "small" if num_hidden_layers <= 50 else "large" + logger.warning( + f"SGLANG_APPLY_CONFIG_BACKUP=auto: checkpoint has " + f"num_hidden_layers={num_hidden_layers}, dispatching to {backup_mode!r}." + ) + if backup_mode != "none": + backup_file = { + "small": "config_backup_small.json", + "large": "config_backup_large.json", + }.get(backup_mode) + if backup_file is None: + raise ValueError( + f"SGLANG_APPLY_CONFIG_BACKUP={backup_mode!r} is not recognized; " + f"use 'none' (off), 'small', 'large', or 'auto'." + ) + config_file = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "configs", + backup_file, + ) + logger.warning( + f"SGLANG_APPLY_CONFIG_BACKUP={backup_mode}: using packaged {config_file} " + f"instead of the checkpoint's config.json at {local_path}." + ) + else: + config_file = os.path.join(local_path, "config.json") + if not os.path.exists(config_file): + raise RuntimeError(f"Can't find config file at {config_file}.") + + with open(config_file, "r") as f: + config_json = json.load(f) + + config_json["architectures"] = [architecture] + config_json["model_type"] = "deepseek_v3" + + tmp_path = os.path.join(tempfile.gettempdir(), "_tmp_config_folder") + os.makedirs(tmp_path, exist_ok=True) + + unique_path = os.path.join(tmp_path, f"{model_type}_{os.getpid()}") + with open(unique_path, "w") as f: + json.dump(config_json, f) + + return AutoConfig.from_pretrained( + unique_path, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) + + +# Temporary hack for Mistral Large +@lru_cache(maxsize=2) +def _load_mistral_large_3_for_causal_LM( + model_path: str, + trust_remote_code: bool = False, + revision: Optional[str] = None, +): + # first get the local path + local_path = download_from_hf(model_path) + # then load the config file in json + parser = mistral_utils.MistralConfigParser() + config_dict, _ = parser.parse(local_path) + + with tempfile.NamedTemporaryFile(mode="w+", suffix=".json") as f: + json.dump(config_dict, f) + f.flush() + loaded_config = AutoConfig.from_pretrained( + f.name, trust_remote_code=trust_remote_code, revision=revision + ) + text_config = getattr(loaded_config, "text_config", None) + if text_config is not None and isinstance(text_config, dict): + text_config = AutoConfig.for_model(**text_config) + setattr(loaded_config, "text_config", text_config) + vision_config = getattr(loaded_config, "vision_config", None) + if vision_config is not None and isinstance(vision_config, dict): + vision_config = AutoConfig.for_model(**vision_config) + setattr(loaded_config, "vision_config", vision_config) + + return loaded_config + + +def _is_deepseek_ocr_model(config: PretrainedConfig) -> bool: + # TODO: Remove this workaround related when AutoConfig correctly identifies deepseek-ocr. + # Hugging Face's AutoConfig currently misidentifies it as deepseekvl2. + auto_map = getattr(config, "auto_map", None) or {} + return auto_map.get("AutoModel") == "modeling_deepseekocr.DeepseekOCRForCausalLM" + + +def _is_deepseek_ocr2_model(config: PretrainedConfig) -> bool: + auto_map = getattr(config, "auto_map", None) or {} + return auto_map.get("AutoModel") == "modeling_deepseekocr2.DeepseekOCR2ForCausalLM" + + +def _override_deepseek_ocr_v_head_dim(config: DeepseekVLV2Config) -> None: + # FIXME: deepseek-ocr's v_head_dim is set to 0 in its config file. + # https://huggingface.co/deepseek-ai/DeepSeek-OCR/blob/main/config.json#L116 + if config.text_config.v_head_dim == 0: + V_HEAD_DIM_PATCH = 128 + config.text_config.v_head_dim = V_HEAD_DIM_PATCH + # Also fix language_config so get_hf_text_config (which may prefer it + # over text_config) stays consistent. + lc = getattr(config, "language_config", None) + if isinstance(lc, dict): + lc["v_head_dim"] = V_HEAD_DIM_PATCH + elif hasattr(lc, "v_head_dim"): + lc.v_head_dim = V_HEAD_DIM_PATCH + logger.warning( + f"Overriding deepseek-ocr's v_head_dim from 0 to {V_HEAD_DIM_PATCH} to avoid potential issues." + ) + + +def _override_v_head_dim_if_zero(config: PretrainedConfig, patch: int = 128) -> None: + text_config = getattr(config, "text_config", None) + language_config = getattr(config, "language_config", None) + target = text_config or language_config + if target is None: + return + if getattr(target, "v_head_dim", None) == 0: + setattr(target, "v_head_dim", patch) + logger.warning( + f"Overriding v_head_dim from 0 to {patch} to avoid potential issues." + ) + + +def _ensure_clean_up_tokenization_compat() -> None: + """Re-add ``clean_up_tokenization`` removed in transformers v5. + + Remote-code tokenizers (e.g. InternLM2Tokenizer) call + ``self.clean_up_tokenization()`` which was a static method on + ``PreTrainedTokenizerBase`` in v4 but removed in v5. Patch it back + so existing HuggingFace Hub tokenizer code keeps working. + """ + if hasattr(PreTrainedTokenizerBase, "clean_up_tokenization"): + return + + @staticmethod + def clean_up_tokenization(out_string: str) -> str: + out_string = ( + out_string.replace(" .", ".") + .replace(" ?", "?") + .replace(" !", "!") + .replace(" ,", ",") + .replace(" ' ", "'") + .replace(" n't", "n't") + .replace(" 'm", "'m") + .replace(" 's", "'s") + .replace(" 've", "'ve") + .replace(" 're", "'re") + ) + return out_string + + PreTrainedTokenizerBase.clean_up_tokenization = clean_up_tokenization + + +# Apply immediately so all code paths (get_tokenizer, get_processor, +# and any external callers) benefit without needing an explicit call. +_ensure_clean_up_tokenization_compat() + + +def _ensure_is_torch_fx_available_compat() -> None: + """Re-add ``is_torch_fx_available`` removed in transformers v5. + + Remote-code models (e.g. MiniCPM-V) import ``is_torch_fx_available`` + from ``transformers.utils.import_utils``. The function was removed + in v5. Patch it back so existing HuggingFace Hub model code keeps + working. torch.fx is always available in PyTorch >= 2.0. + """ + import transformers.utils.import_utils as _import_utils + + if hasattr(_import_utils, "is_torch_fx_available"): + return + + _import_utils.is_torch_fx_available = lambda: True + + +_ensure_is_torch_fx_available_compat() + + +def normalize_rope_scaling_compat(config: "PretrainedConfig") -> None: + """Ensure rope_scaling dicts have ``"type"`` alongside ``"rope_type"``. + + Transformers v5 standardises rope_scaling to use ``"rope_type"`` and may + omit the legacy ``"type"`` key. Remote-code models (e.g. Kimi-VL) still + read ``rope_scaling["type"]``, causing a ``KeyError``. This helper adds + ``"type"`` from ``"rope_type"`` whenever it is missing, recursively across + the config and all its sub-configs. + """ + + def _patch(cfg): + try: + rs = getattr(cfg, "rope_scaling", None) + except AttributeError: + rs = None + if isinstance(rs, dict) and "rope_type" in rs and "type" not in rs: + rs["type"] = rs["rope_type"] + # Recurse into sub-configs + for attr in ( + "text_config", + "llm_config", + "language_config", + "vision_config", + "thinker_config", + ): + sub = getattr(cfg, attr, None) + if sub is not None: + _patch(sub) + + _patch(config) + + +def _ensure_llama_flash_attention2_compat() -> None: + """Ensure LlamaFlashAttention2 symbol exists for remote code compatibility.""" + try: + from transformers.models.llama import modeling_llama + except (ImportError, ModuleNotFoundError): + return + if not hasattr(modeling_llama, "LlamaFlashAttention2"): + if hasattr(modeling_llama, "LlamaAttention"): + modeling_llama.LlamaFlashAttention2 = modeling_llama.LlamaAttention + + +def _ensure_gguf_version(): + """Workaround for transformers v5 bug where is_gguf_available() fails + when the gguf package lacks __version__ and metadata lookup also fails, + resulting in packaging.version.InvalidVersion: Invalid version: 'N/A'.""" + try: + import gguf + + if not hasattr(gguf, "__version__"): + import importlib.metadata + + try: + gguf.__version__ = importlib.metadata.version("gguf") + except Exception: + gguf.__version__ = "0.0.0" + except ImportError: + pass + + +@lru_cache_frozenset(maxsize=32) +def get_config( + model: str, + trust_remote_code: bool, + revision: Optional[str] = None, + model_override_args: Optional[dict] = None, + **kwargs, +): + is_gguf = check_gguf_file(model) + if is_gguf: + _ensure_gguf_version() + kwargs["gguf_file"] = model + model = Path(model).parent + + if is_runai_obj_uri(model): + model = ObjectStorageModel.get_path(model) + + if is_remote_url(model): + # 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(model) + client.pull_files(ignore_pattern=["*.pt", "*.safetensors", "*.bin"]) + model = client.get_local_dir() + + if ( + "mistral-large-3" in str(model).lower() + or "mistral-small-4" in str(model).lower() + or "leanstral" in str(model).lower() + ): + config = _load_mistral_large_3_for_causal_LM( + model, trust_remote_code=trust_remote_code, revision=revision + ) + elif envs.SGLANG_APPLY_CONFIG_BACKUP.get() != "none": + config = _load_deepseek_temp_model( + model, + model_type="deepseek_ref", + architecture="DeepseekV4ForCausalLM", + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + else: + _ensure_llama_flash_attention2_compat() + try: + config = AutoConfig.from_pretrained( + model, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) + except ValueError as e: + if "deepseek_ref" in str(e): + config = _load_deepseek_temp_model( + model, + model_type="deepseek_ref", + architecture="DeepseekV4ForCausalLM", + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + elif "deepseek_v32" in str(e): + config = _load_deepseek_temp_model( + model, + model_type="deepseek_v32", + architecture="DeepseekV3ForCausalLM", + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + else: + raise e + except KeyError as e: + # Transformers v5 may register a built-in config class that + # conflicts with sglang's custom one (e.g. NemotronHConfig + # doesn't handle '-' in hybrid_override_pattern). Fall back + # to loading the raw config dict and using sglang's class. + # Also handle deepseek_v32 which v5 doesn't recognize. + if "deepseek_v32" in str(e): + config = _load_deepseek_temp_model( + model, + model_type="deepseek_v32", + architecture="DeepseekV3ForCausalLM", + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + else: + config_dict, _ = PretrainedConfig.get_config_dict( + model, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + model_type = config_dict.get("model_type") + if model_type in _CONFIG_REGISTRY: + config = _CONFIG_REGISTRY[model_type].from_dict(config_dict) + config._name_or_path = model + else: + raise + + if ( + config.architectures is not None + and config.architectures[0] == "Phi4MMForCausalLM" + ): + # Phi4MMForCausalLM uses a hard-coded vision_config. See: + # https://github.com/vllm-project/vllm/blob/6071e989df1531b59ef35568f83f7351afb0b51e/vllm/model_executor/models/phi4mm.py#L71 + # We set it here to support cases where num_attention_heads is not divisible by the TP size. + from transformers import SiglipVisionConfig + + vision_config = { + "hidden_size": 1152, + "image_size": 448, + "intermediate_size": 4304, + "model_type": "siglip_vision_model", + "num_attention_heads": 16, + "num_hidden_layers": 26, + # Model is originally 27-layer, we only need the first 26 layers for feature extraction. + "patch_size": 14, + } + config.vision_config = SiglipVisionConfig(**vision_config) + + if config.architectures in [ + ["LongcatCausalLM"], + ["LongcatFlashForCausalLM"], + ["LongcatFlashNgramForCausalLM"], + ]: + config.model_type = "longcat_flash" + + text_config = get_hf_text_config(config=config) + + if isinstance(model, str) and text_config is not None: + items = ( + text_config.items() + if hasattr(text_config, "items") + else vars(text_config).items() + ) + for key, val in items: + if not hasattr(config, key) and val is not None: + setattr(config, key, val) + + if _is_deepseek_ocr2_model(config): + _override_v_head_dim_if_zero(config) + # Temporary hack for load deepseek-ocr2 + config.model_type = "deepseek-ocr" + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) + config = DeepseekVLV2Config.from_pretrained(model, revision=revision) + _override_v_head_dim_if_zero(config) + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) + setattr(config, "_name_or_path", model) + elif config.model_type in _CONFIG_REGISTRY: + model_type = config.model_type + if model_type == "deepseek_vl_v2": + if _is_deepseek_ocr_model(config) or _is_deepseek_ocr2_model(config): + model_type = "deepseek-ocr" + config_class = _CONFIG_REGISTRY[model_type] + config = config_class.from_pretrained(model, revision=revision) + + if _is_deepseek_ocr_model(config): + _override_deepseek_ocr_v_head_dim(config) + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) + elif _is_deepseek_ocr2_model(config): + _override_v_head_dim_if_zero(config) + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) + + # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`. + setattr(config, "_name_or_path", model) + + if isinstance(model, str) and config.model_type == "internvl_chat": + for key, val in config.llm_config.__dict__.items(): + if not hasattr(config, key): + setattr(config, key, val) + + if config.model_type == "multi_modality": + config.update({"architectures": ["MultiModalityCausalLM"]}) + + if config.model_type == "gemma4": + # Gemma4 configs use base attributes for SWA layers and `global_*` + # variants for full-attention layers. SGLang expects the opposite: + # base = full-attention, `swa_*` = sliding-window overrides. + # Remap here so the rest of the stack sees a uniform convention. + text_config = config.text_config + global_head_dim = getattr(text_config, "global_head_dim", None) + global_kv_heads = getattr(text_config, "num_global_key_value_heads", None) + + swa_head_dim = text_config.head_dim + swa_kv_heads = text_config.num_key_value_heads + + text_config.swa_head_dim = swa_head_dim + text_config.swa_v_head_dim = swa_head_dim + text_config.swa_num_key_value_heads = swa_kv_heads + + if global_head_dim is not None: + text_config.head_dim = global_head_dim + if global_kv_heads is not None: + text_config.num_key_value_heads = global_kv_heads + + if not hasattr(text_config, "v_head_dim"): + text_config.v_head_dim = text_config.head_dim + if not hasattr(text_config, "swa_v_head_dim"): + text_config.swa_v_head_dim = text_config.swa_head_dim + + if config.model_type == "longcat_flash": + config.update({"architectures": ["LongcatFlashForCausalLM"]}) + + 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 + + +@lru_cache_frozenset(maxsize=32) +def get_generation_config( + model: str, + trust_remote_code: bool, + revision: Optional[str] = None, + **kwargs, +): + try: + return GenerationConfig.from_pretrained( + model, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) + except OSError as e: + return None + + +# Qwen-1M related +def get_sparse_attention_config( + model: str, + sparse_attention_config_filename: str = "sparse_attention_config.json", +) -> Dict[str, Any]: + is_local = os.path.isdir(model) + if not is_local: + # Download the config files. + model = download_from_hf(model, allow_patterns=["*.json"]) + + config_file = os.path.join(model, sparse_attention_config_filename) + if not os.path.exists(config_file): + return {} + + # Load the sparse attention config. + with open(config_file) as f: + config = json.load(f) + return config + + +# Models don't use the same configuration key for determining the maximum +# context length. Store them here so we can sanely check them. +# NOTE: The ordering here is important. Some models have two of these and we +# have a preference for which value gets used. +CONTEXT_LENGTH_KEYS = [ + "max_sequence_length", + "seq_length", + "max_seq_len", + "model_max_length", + "max_position_embeddings", +] + + +def get_context_length(config): + """Get the context length of a model from a huggingface model configs.""" + text_config = config + rope_scaling = getattr(text_config, "rope_scaling", None) + if rope_scaling: + rope_scaling_factor = rope_scaling.get("factor", 1) + if "original_max_position_embeddings" in rope_scaling: + rope_scaling_factor = 1 + if rope_scaling.get("rope_type", None) == "llama3": + rope_scaling_factor = 1 + else: + rope_scaling_factor = 1 + + for key in CONTEXT_LENGTH_KEYS: + val = getattr(text_config, key, None) + if val is not None: + return int(rope_scaling_factor * val) + return 2048 + + +# A fast LLaMA tokenizer with the pre-processed `tokenizer.json` file. +_FAST_LLAMA_TOKENIZER = "hf-internal-testing/llama-tokenizer" + + +# Filter warnings like: https://github.com/sgl-project/sglang/issues/8082 +class TokenizerWarningsFilter(logging.Filter): + def filter(self, record: logging.LogRecord) -> bool: + return "Calling super().encode with" not in record.getMessage() + + +_is_base_mistral_patched = False + +# transformers version where _patch_mistral_regex calls model_info() on every tokenizer load +_TRANSFORMERS_PATCHED_VERSION = "5.3.0" + + +def _patch_is_base_mistral_in_ci(): + """Patch transformers' _patch_mistral_regex to avoid HF API calls in CI. + + transformers defines is_base_mistral as a local function inside + _patch_mistral_regex, so it cannot be patched via module attribute. + Instead we replace the entire _patch_mistral_regex classmethod with a + version that simply returns the tokenizer unchanged. + + In CI this prevents exhausting the 3000 req/5min HF API rate limit. + """ + global _is_base_mistral_patched + if _is_base_mistral_patched: + return + + from sglang.srt.environ import envs + + if not envs.SGLANG_IS_IN_CI.get(): + return + + import transformers + + if transformers.__version__ != _TRANSFORMERS_PATCHED_VERSION: + logger.warning( + "transformers version changed to %s (expected %s), " + "_patch_mistral_regex patch skipped — may need update if 429 errors recur", + transformers.__version__, + _TRANSFORMERS_PATCHED_VERSION, + ) + _is_base_mistral_patched = True # don't warn repeatedly + return + + from transformers import PreTrainedTokenizerFast + + if hasattr(PreTrainedTokenizerFast, "_patch_mistral_regex"): + + @classmethod + def _noop_patch_mistral_regex(cls, tokenizer, *args, **kwargs): + return tokenizer + + PreTrainedTokenizerFast._patch_mistral_regex = _noop_patch_mistral_regex + logger.info("CI: patched _patch_mistral_regex to skip HF API calls") + + _is_base_mistral_patched = True + + +def get_tokenizer( + tokenizer_name: str, + *args, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + tokenizer_revision: Optional[str] = None, + **kwargs, +) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: + """Gets a tokenizer for the given model name via Huggingface.""" + if tokenizer_name.endswith(".json"): + from sglang.srt.tokenizer.tiktoken_tokenizer import TiktokenTokenizer + + return TiktokenTokenizer(tokenizer_name) + + if tokenizer_mode == "slow": + if kwargs.get("use_fast", False): + raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.") + kwargs["use_fast"] = False + elif tokenizer_mode == "auto": + # In Transformers v5, the default for use_fast changed from True to False. + # Explicitly set use_fast=True for "auto" mode to maintain previous behavior + # and avoid issues with models that have incorrect tokenizer_class values. + if "use_fast" not in kwargs: + kwargs["use_fast"] = True + + # TODO(Xinyuan): Remove this once we have a proper tokenizer for Devstral + if tokenizer_name == "mistralai/Devstral-Small-2505": + tokenizer_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503" + + is_gguf = check_gguf_file(tokenizer_name) + if is_gguf: + _ensure_gguf_version() + kwargs["gguf_file"] = tokenizer_name + tokenizer_name = Path(tokenizer_name).parent + + if is_runai_obj_uri(tokenizer_name): + tokenizer_name = ObjectStorageModel.get_path(tokenizer_name) + + if is_remote_url(tokenizer_name): + # 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(tokenizer_name) + client.pull_files(ignore_pattern=["*.pt", "*.safetensors", "*.bin"]) + tokenizer_name = client.get_local_dir() + + _patch_is_base_mistral_in_ci() + + try: + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + tokenizer_revision=tokenizer_revision, + clean_up_tokenization_spaces=False, + **kwargs, + ) + # Filter tokenizer warnings + logging.getLogger(tokenizer.__class__.__module__).addFilter( + TokenizerWarningsFilter() + ) + except TypeError as e: + # The LLaMA tokenizer causes a protobuf error in some environments. + err_msg = ( + "Failed to load the tokenizer. If you are using a LLaMA V1 model " + f"consider using '{_FAST_LLAMA_TOKENIZER}' instead of the " + "original tokenizer." + ) + raise RuntimeError(err_msg) from e + except ValueError as e: + # If the error pertains to the tokenizer class not existing or not + # currently being imported, suggest using the --trust-remote-code flag. + if not trust_remote_code and ( + "does not exist or is not currently imported." in str(e) + or "requires you to execute the tokenizer file" in str(e) + ): + err_msg = ( + "Failed to load the tokenizer. If the tokenizer is a custom " + "tokenizer not yet available in the HuggingFace transformers " + "library, consider setting `trust_remote_code=True` in LLM " + "or using the `--trust-remote-code` flag in the CLI." + ) + raise RuntimeError(err_msg) from e + else: + raise e + + # Transformers v5 may silently fall back to a generic TokenizersBackend + # when trust_remote_code=False and the model requires a custom tokenizer. + # Detect this and auto-retry with trust_remote_code=True. + if not trust_remote_code and type(tokenizer).__name__ == "TokenizersBackend": + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=True, + tokenizer_revision=tokenizer_revision, + clean_up_tokenization_spaces=False, + **kwargs, + ) + + _fix_v5_tokenizer_components(tokenizer, tokenizer_name, tokenizer_revision) + _fix_v5_add_bos_eos_token(tokenizer, tokenizer_name, tokenizer_revision) + + if not isinstance(tokenizer, PreTrainedTokenizerFast): + warnings.warn( + "Using a slow tokenizer. This might cause a significant " + "slowdown. Consider using a fast tokenizer instead." + ) + + _fix_special_tokens_pattern(tokenizer) + attach_additional_stop_token_ids(tokenizer) + + return tokenizer + + +def _resolve_local_or_cached_file(model_name_or_path, filename, revision=None): + """Resolve a file from a local directory or HF hub cache (no network).""" + local_path = Path(model_name_or_path) / filename + if local_path.is_file(): + return str(local_path) + from huggingface_hub import hf_hub_download + + return hf_hub_download( + model_name_or_path, filename, revision=revision, local_files_only=True + ) + + +def _fix_v5_tokenizer_components(tokenizer, model_name_or_path, revision=None): + """Fix pre_tokenizer/decoder when a v5 tokenizer class overwrites them. + + In transformers v5, some tokenizer classes (e.g. LlamaTokenizer) have a + custom __init__ that rebuilds the pre_tokenizer and decoder from scratch + with class-specific components, discarding the originals from tokenizer.json. + This breaks models that specify LlamaTokenizerFast but actually use a + different tokenizer architecture (e.g. DeepSeek-V3.2 uses ByteLevel). + + Detects the mismatch by comparing against the raw tokenizer.json and + restores the original components when they differ. + """ + backend = getattr(tokenizer, "_tokenizer", None) + if backend is None: + return + + try: + from tokenizers import Tokenizer as RawTokenizer + + tok_file = _resolve_local_or_cached_file( + model_name_or_path, "tokenizer.json", revision + ) + raw = RawTokenizer.from_file(tok_file) + except Exception as e: + logger.debug( + "_fix_v5_tokenizer_components: could not load tokenizer.json for %s: %s", + model_name_or_path, + e, + ) + return + + raw_pre = type(raw.pre_tokenizer).__name__ if raw.pre_tokenizer else None + loaded_pre = type(backend.pre_tokenizer).__name__ if backend.pre_tokenizer else None + + if raw_pre and loaded_pre and raw_pre != loaded_pre: + logger.info( + "Fixing v5 tokenizer component mismatch for %s: " + "pre_tokenizer %s -> %s, decoder %s -> %s", + model_name_or_path, + loaded_pre, + raw_pre, + type(backend.decoder).__name__ if backend.decoder else None, + type(raw.decoder).__name__ if raw.decoder else None, + ) + backend.pre_tokenizer = raw.pre_tokenizer + backend.decoder = raw.decoder + + +def _fix_v5_add_bos_eos_token(tokenizer, model_name_or_path, revision=None): + """Restore add_bos_token/add_eos_token stripped by transformers v5. + + In transformers v5, _from_pretrained() strips add_bos_token and + add_eos_token from init kwargs when a tokenizer.json file is present, + assuming the tokenizer.json post-processor handles BOS/EOS addition. + However, many models (e.g. DeepSeek-V3) have a tokenizer.json whose + post-processor does NOT add BOS/EOS, and rely on the add_bos_token flag + from tokenizer_config.json instead. This causes silent accuracy regressions. + + This function reads the tokenizer_config.json and restores the values, + but only for tokenizer classes that actually supported these flags in v4. + Classes like Qwen2Tokenizer did not support add_bos_token/add_eos_token + in v4, so restoring them would change behavior. + """ + # In transformers v4, only certain tokenizer classes supported + # add_bos_token / add_eos_token as init parameters. Restoring these + # flags for classes that never supported them (e.g. Qwen2Tokenizer) + # would incorrectly change tokenization behavior. + _V4_CLASSES_WITH_BOS_EOS_FLAGS = frozenset( + { + "LlamaTokenizer", + "LlamaTokenizerFast", + "CodeLlamaTokenizer", + "CodeLlamaTokenizerFast", + "GemmaTokenizer", + "GemmaTokenizerFast", + "CohereTokenizerFast", + } + ) + + try: + config_file = _resolve_local_or_cached_file( + model_name_or_path, "tokenizer_config.json", revision + ) + with open(config_file) as f: + config = json.load(f) + except Exception as e: + logger.debug( + "_fix_v5_add_bos_eos_token: could not read tokenizer_config.json " + "for %s: %s", + model_name_or_path, + e, + ) + return + + tokenizer_class = config.get("tokenizer_class", "") + if tokenizer_class not in _V4_CLASSES_WITH_BOS_EOS_FLAGS: + logger.debug( + "_fix_v5_add_bos_eos_token: skipping %s (tokenizer_class=%s " + "did not support add_bos/eos_token in v4)", + model_name_or_path, + tokenizer_class, + ) + return + + # In v4, Llama/Gemma tokenizers defaulted add_bos_token=True. + # When the config omits the key or has null, use the v4 default so that + # update_post_processor() doesn't drop BOS/EOS that was there before. + _V4_DEFAULTS = {"add_bos_token": True, "add_eos_token": False} + + changed = False + for attr in ("add_bos_token", "add_eos_token"): + config_val = config.get(attr) + if config_val is None: + # Key missing or null → use v4 default for this tokenizer class + config_val = _V4_DEFAULTS.get(attr, False) + # Fast tokenizers in v4 used tokenizer.json post-processor for EOS — + # the add_eos_token Python attribute was set but the post-processor + # came from tokenizer.json, not from the attribute. In v5, the flag is + # stripped and both sglang and HF reference end up with add_eos_token=False. + # Restoring add_eos_token for fast tokenizers makes sglang diverge from + # the HF reference (which doesn't restore it), breaking embedding models + # like intfloat/e5-mistral-7b-instruct (cosine similarity drops to ~0.33). + if attr == "add_eos_token" and isinstance(tokenizer, PreTrainedTokenizerFast): + config_val = _V4_DEFAULTS["add_eos_token"] # False + current_val = getattr(tokenizer, attr, None) + if current_val != config_val: + logger.info( + "Restoring %s=%s for %s (was %s after v5 loading)", + attr, + config_val, + model_name_or_path, + current_val, + ) + setattr(tokenizer, f"_{attr}", config_val) + changed = True + + # Rebuild the post-processor so it respects the restored flags + if changed and hasattr(tokenizer, "update_post_processor"): + tokenizer.update_post_processor() + + +def _fix_special_tokens_pattern(tokenizer): + """Fix https://github.com/huggingface/transformers/pull/42563 which defaults + special_tokens_pattern to "cls_sep", inserting None into token IDs when + cls_token/sep_token are undefined (e.g. Kimi-VL's TikTokenTokenizer). + """ + pattern = getattr(tokenizer, "special_tokens_pattern", None) + if pattern == "cls_sep" and ( + tokenizer.cls_token_id is None or tokenizer.sep_token_id is None + ): + tokenizer.special_tokens_pattern = "none" + + +def _fix_added_tokens_encoding(tokenizer): + """Ensure special tokens encode as single tokens in transformers v5. + + Some model tokenizers (e.g. MiniCPM-V-4) define special tokens like , + as attributes on the tokenizer class with corresponding IDs in the + vocabulary (via tokenizer.json's added_tokens). In transformers v5, these + tokens may not appear in get_added_vocab() and encode() splits them into + subwords, breaking multimodal pipelines that rely on finding them in input_ids. + + This function discovers such tokens by scanning tokenizer attributes, checks + if they encode correctly, and re-registers any that don't. + """ + # Discover special token strings from tokenizer attributes. + # Model tokenizers (e.g. MiniCPMVTokenizerFast) store them as attributes + # like im_start="", slice_start="", etc. + candidates = {} + for attr in dir(tokenizer): + if attr.startswith("_"): + continue + try: + val = getattr(tokenizer, attr) + except Exception: + continue + if ( + not isinstance(val, str) + or not val.startswith("<") + or not val.endswith(">") + or len(val) > 20 + ): + continue + token_id = tokenizer.convert_tokens_to_ids(val) + if token_id is not None and token_id != tokenizer.unk_token_id: + candidates[val] = token_id + + if not candidates: + return + + # Check which tokens fail to encode as single tokens. + broken = [] + for token_str, expected_id in candidates.items(): + try: + ids = tokenizer.encode(token_str, add_special_tokens=False) + if len(ids) != 1 or ids[0] != expected_id: + broken.append(token_str) + except Exception: + broken.append(token_str) + + if not broken: + return + + from transformers import AddedToken + + tokens_to_add = [AddedToken(tok, special=True, normalized=False) for tok in broken] + tokenizer.add_tokens(tokens_to_add, special_tokens=True) + logger.info( + "Re-registered %d special tokens for correct v5 encoding: %s", + len(broken), + broken[:10], + ) + + +# Some models doesn't have an available processor, e.g.: InternVL +def get_tokenizer_from_processor(processor): + if isinstance(processor, PreTrainedTokenizerBase): + return processor + return processor.tokenizer + + +def _build_processor_manually( + model_path, config, trust_remote_code, revision, **kwargs +): + """Build processor when AutoProcessor fails to resolve feature_extractor_type. + + In transformers v5, AutoProcessor.from_pretrained calls + AutoFeatureExtractor.from_pretrained which fails if + preprocessor_config.json lacks 'feature_extractor_type'. This loads the + processor class from the hub and constructs it with individually-loaded + components. + """ + import transformers + from transformers import AutoImageProcessor, AutoTokenizer + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + # Resolve processor class from auto_map — check both the model config + # and the preprocessor_config.json (some models like MiniCPM-o only + # declare AutoProcessor in the latter). + auto_map = getattr(config, "auto_map", None) or {} + proc_ref = auto_map.get("AutoProcessor") + if not proc_ref: + try: + pp_file = _resolve_local_or_cached_file( + model_path, "preprocessor_config.json", revision + ) + with open(pp_file) as f: + pp_auto_map = json.load(f).get("auto_map", {}) + proc_ref = pp_auto_map.get("AutoProcessor") + except Exception as e: + logger.debug( + "_build_processor_manually: could not read preprocessor_config.json " + "for %s: %s", + model_path, + e, + ) + if not proc_ref: + raise ValueError(f"Cannot determine processor class for {model_path}") + + proc_cls = get_class_from_dynamic_module( + proc_ref, model_path, code_revision=revision + ) + + # Load sub-components individually (these succeed) + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=trust_remote_code, revision=revision + ) + init_kwargs = {"tokenizer": tokenizer} + + if "image_processor" in getattr(proc_cls, "attributes", []): + try: + init_kwargs["image_processor"] = AutoImageProcessor.from_pretrained( + model_path, trust_remote_code=trust_remote_code, revision=revision + ) + except Exception as e: + logger.warning("Failed to load image_processor for %s: %s", model_path, e) + + # Instantiate feature extractor from its declared class + fe_class_name = getattr(proc_cls, "feature_extractor_class", None) + if fe_class_name: + fe_class = getattr(transformers, fe_class_name, None) + if fe_class is not None: + init_kwargs["feature_extractor"] = fe_class() + + return proc_cls(**init_kwargs) + + +def get_processor( + tokenizer_name: str, + *args, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + tokenizer_revision: Optional[str] = None, + use_fast: Optional[bool] = True, + **kwargs, +): + # pop 'revision' from kwargs if present. + revision = kwargs.pop("revision", tokenizer_revision) + if ( + "mistral-large-3" in str(tokenizer_name).lower() + or "mistral-small-4" in str(tokenizer_name).lower() + or "leanstral" in str(tokenizer_name).lower() + ): + config = _load_mistral_large_3_for_causal_LM( + tokenizer_name, + trust_remote_code=trust_remote_code, + revision=revision, + ) + else: + _ensure_llama_flash_attention2_compat() + config = AutoConfig.from_pretrained( + tokenizer_name, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + if _is_deepseek_ocr_model(config): + # Temporary hack for load deepseek-ocr + config.model_type = "deepseek-ocr" + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) + elif _is_deepseek_ocr2_model(config): + # Temporary hack for load deepseek-ocr2 + config.model_type = "deepseek-ocr" + config.update({"architectures": ["DeepseekOCRForCausalLM"]}) + _override_v_head_dim_if_zero(config) + + # fix: for Qwen2-VL and Sarashina2Vision models, inject default 'size' if not provided. + if config.model_type in {"qwen2_vl", "sarashina2_vision"}: + if "size" not in kwargs: + kwargs["size"] = {"shortest_edge": 3136, "longest_edge": 1003520} + + if config.model_type not in {"llava", "clip"}: + kwargs["use_fast"] = use_fast + try: + if "InternVL3_5" in tokenizer_name: + processor = AutoTokenizer.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + else: + if config.model_type in _CUSTOMIZED_MM_PROCESSOR: + processor = _CUSTOMIZED_MM_PROCESSOR[config.model_type].from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + else: + processor = AutoProcessor.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + + except ValueError as e: + error_message = str(e) + if "does not have a slow version" in error_message: + logger.info( + f"Processor {tokenizer_name} does not have a slow version. Automatically use fast version" + ) + kwargs["use_fast"] = True + processor = AutoProcessor.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + elif "Unrecognized feature extractor" in error_message: + logger.info( + "AutoProcessor failed on feature extractor for %s, " + "constructing processor manually", + tokenizer_name, + ) + processor = _build_processor_manually( + tokenizer_name, + config, + trust_remote_code, + revision, + **kwargs, + ) + else: + raise e + # If processor is a bare tokenizer (e.g. Mistral-Small-4 has no processor_config.json) + # and the model is a vision model (pixtral), wrap it in a proper PixtralProcessor + # so that image data is actually processed through the image processor. + if ( + isinstance(processor, PreTrainedTokenizerBase) + and getattr(config, "model_type", None) == "pixtral" + ): + from transformers.models.pixtral.image_processing_pixtral import ( + PixtralImageProcessor, + ) + from transformers.models.pixtral.processing_pixtral import ( + PixtralProcessor as HFPixtralProcessor, + ) + + vision_config = config.vision_config + patch_size = vision_config.patch_size + image_size = vision_config.image_size + spatial_merge_size = getattr(vision_config, "spatial_merge_size", 1) + + effective_patch = patch_size * spatial_merge_size + image_processor = PixtralImageProcessor( + do_resize=True, + size={"longest_edge": image_size}, + patch_size={"height": effective_patch, "width": effective_patch}, + ) + processor = HFPixtralProcessor( + image_processor=image_processor, + tokenizer=processor, + patch_size=patch_size, + spatial_merge_size=spatial_merge_size, + ) + + tokenizer = get_tokenizer_from_processor(processor) + + if tokenizer.chat_template is None: + local_path = download_from_hf( + tokenizer_name, allow_patterns=["*.json", "*.jinja", "*.model"] + ) + jinja_path = Path(local_path) / "chat_template.jinja" + if jinja_path.is_file(): + tokenizer.chat_template = jinja_path.read_text() + logger.info("Loaded chat_template from %s", jinja_path) + + _fix_special_tokens_pattern(tokenizer) + _fix_added_tokens_encoding(tokenizer) + attach_additional_stop_token_ids(tokenizer) + return processor + + +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: Union[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" diff --git a/platforms/patches/kunlun_p800/nic_priority_matrix_test.json b/platforms/patches/kunlun_p800/nic_priority_matrix_test.json new file mode 100644 index 0000000..f8ef41d --- /dev/null +++ b/platforms/patches/kunlun_p800/nic_priority_matrix_test.json @@ -0,0 +1,98 @@ +{ + "cpu:0": [ + [ + "mlx5_0" + ], + [] + ], + "cpu:1": [ + [ + "mlx5_0" + ], + [] + ], + "cpu:2": [ + [ + "mlx5_0" + ], + [] + ], + "cpu:3": [ + [ + "mlx5_0" + ], + [] + ], + "cpu:4": [ + [ + "mlx5_0" + ], + [] + ], + "cpu:5": [ + [ + "mlx5_0" + ], + [] + ], + "cpu:6": [ + [ + "mlx5_0" + ], + [] + ], + "cpu:7": [ + [ + "mlx5_0" + ], + [] + ], + "cuda:0": [ + [ + "mlx5_1" + ], + [] + ], + "cuda:1": [ + [ + "mlx5_1" + ], + [] + ], + "cuda:2": [ + [ + "mlx5_2" + ], + [] + ], + "cuda:3": [ + [ + "mlx5_2" + ], + [] + ], + "cuda:4": [ + [ + "mlx5_3" + ], + [] + ], + "cuda:5": [ + [ + "mlx5_3" + ], + [] + ], + "cuda:6": [ + [ + "mlx5_4" + ], + [] + ], + "cuda:7": [ + [ + "mlx5_4" + ], + [] + ] +} diff --git a/platforms/patches/kunlun_p800/parallel_state.py b/platforms/patches/kunlun_p800/parallel_state.py new file mode 100644 index 0000000..e21b824 --- /dev/null +++ b/platforms/patches/kunlun_p800/parallel_state.py @@ -0,0 +1,2384 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/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. +"""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/pipeline + parallelism, you can skip the model parallel initialization and destruction + steps. +""" + +import contextlib +import gc +import logging +import os +import pickle +import weakref +from collections import namedtuple +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass +from datetime import timedelta +from multiprocessing import shared_memory +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from unittest.mock import patch + +import torch +import torch.distributed +from torch.distributed import Backend, ProcessGroup + +from sglang.srt.compilation.compilation_config import register_split_op +from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph +from sglang.srt.distributed.utils import set_global_tcp_store +from sglang.srt.environ import envs +from sglang.srt.utils import ( + get_bool_env_var, + get_current_device_stream_fast, + get_int_env_var, + is_cpu, + is_cuda_alike, + is_hip, + is_musa, + is_npu, + is_shm_available, + is_xpu, +) +from sglang.srt.utils.custom_op import register_custom_op +from sglang.srt.utils.network import get_local_ip_auto + +_is_npu = is_npu() +_is_cpu = is_cpu() +_is_xpu = is_xpu() +_is_musa = is_musa() +_USE_ALLGATHER_ADD = ( + os.environ.get("XSGL_ALLREDUCE_USE_ALLGATHER_ADD", "0") == "1" +) # 排除 allgather不稳定因素！ +_is_cuda_alike = is_cuda_alike() + +TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) + +# use int value instead of ReduceOp.SUM to support torch compile +REDUCE_OP_SUM = int(torch.distributed.ReduceOp.SUM) + + +def get_torch_distributed_pg_options(group_name=None): + if not _is_npu: + return None + + # Only create HCCL options for default group or MoE-related groups + if group_name is not None and "moe" not in group_name: + return None + + import torch_npu + + options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options() + hccl_buffer_size = int( + os.environ.get("DEEPEP_HCCL_BUFFSIZE") or os.environ.get("HCCL_BUFFSIZE") or 200 + ) + options.hccl_config = {"hccl_buffer_size": hccl_buffer_size} + return options + + +@dataclass +class GraphCaptureContext: + stream: torch.get_device_module().Stream + + +@dataclass +class P2PWork: + work: Optional[torch.distributed.Work] + payload: Optional[torch.Tensor] + + +def _split_tensor_dict( + tensor_dict: Dict[str, Union[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 + + +_group_name_counter: Dict[str, int] = {} + + +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 + + +_groups: Dict[str, Callable[[], Optional["GroupCoordinator"]]] = {} + + +def _register_group(group: "GroupCoordinator") -> None: + _groups[group.unique_name] = weakref.ref(group) + + +@register_custom_op(mutates_args=["tensor"]) +@register_split_op() +def inplace_all_reduce(tensor: torch.Tensor, group_name: str) -> None: + 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.") + group._all_reduce_in_place(tensor) + + +@register_custom_op(out_shape="tensor") +def outplace_all_reduce( + tensor: torch.Tensor, group_name: str, outplace_all_reduce_method: 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, outplace_all_reduce_method) + + +@register_custom_op(mutates_args=["output"]) +def reg_all_gather_into_tensor( + output: torch.Tensor, input: torch.Tensor, group_name: str +) -> None: + 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.") + group._all_gather_into_tensor(output, input) + + +@register_custom_op(mutates_args=["output"]) +def reg_reduce_scatter_tensor( + output: torch.Tensor, input: torch.Tensor, group_name: str +) -> None: + 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.") + group._reduce_scatter_tensor(output, input) + + +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 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_pynccl: bool # a hint of whether to use PyNccl + use_pymscclpp: bool # a hint of whether to use PyMsccl + use_custom_allreduce: bool # a hint of whether to use CustomAllreduce + use_torch_symm_mem_all_reduce: ( + bool # a hint of whether to use TorchSymmMemAllReduce + ) + use_message_queue_broadcaster: ( + bool # a hint of whether to use message queue broadcaster + ) + # communicators are only created for world size > 1 + pynccl_comm: Optional[Any] # PyNccl communicator + ca_comm: Optional[Any] # Custom allreduce communicator + torch_symm_mem_comm: Optional[Any] # Torch symm mem communicator + mq_broadcaster: Optional[Any] # shared memory broadcaster + + def __init__( + self, + group_ranks: List[List[int]], + local_rank: int, + torch_distributed_backend: Union[str, Backend], + use_pynccl: bool, + use_pymscclpp: bool, + use_custom_allreduce: bool, + use_torch_symm_mem_all_reduce: bool, + use_hpu_communicator: bool, + use_xpu_communicator: bool, + use_npu_communicator: bool, + use_message_queue_broadcaster: bool = False, + group_name: Optional[str] = None, + gloo_timeout: timedelta = timedelta(seconds=120 * 60), + ): + # Set group info + group_name = group_name or "anonymous" + self.unique_name = _get_unique_name(group_name) + _register_group(self) + + # Set rank info + self.rank = torch.distributed.get_rank() + self.local_rank = local_rank + self.device_group = None + self.cpu_group = None + self.local_size = get_int_env_var("LOCAL_SIZE", 0) + + if is_cuda_alike(): + device_id = ( + 0 if envs.SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS.get() else local_rank + ) + self.device = torch.device(f"cuda:{device_id}") + elif _is_npu: + self.device = torch.device(f"npu:{local_rank}") + elif _is_xpu: + self.device = torch.device(f"xpu:{local_rank}") + elif _is_musa: + self.device = torch.device(f"musa:{local_rank}") + else: + self.device = torch.device("cpu") + self.device_module = torch.get_device_module(self.device) + + for ranks in group_ranks: + active_ranks = torch.ones(len(ranks), dtype=torch.int32, device=self.device) + active_ranks_cpu = torch.ones(len(ranks), dtype=torch.int32) + if "mooncake" in torch_distributed_backend: + from mooncake.ep import MooncakeBackendOptions + + device_group = torch.distributed.new_group( + ranks, + backend="mooncake", + pg_options=MooncakeBackendOptions(active_ranks), + ) + cpu_group = torch.distributed.new_group( + ranks, + backend="mooncake-cpu", + pg_options=MooncakeBackendOptions(active_ranks_cpu), + ) + else: + pg_options = get_torch_distributed_pg_options(group_name) + device_group = torch.distributed.new_group( + ranks, backend=torch_distributed_backend, pg_options=pg_options + ) + # a group with `gloo` backend, to allow direct coordination + # between processes through the CPU. + cpu_group = torch.distributed.new_group( + ranks, backend="gloo", timeout=gloo_timeout + ) + 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 + self.active_ranks = active_ranks + self.active_ranks_cpu = active_ranks_cpu + + assert self.cpu_group is not None + assert self.device_group is not None + + # Import communicators + self.use_pynccl = use_pynccl + self.use_pymscclpp = use_pymscclpp + self.use_custom_allreduce = use_custom_allreduce + self.use_torch_symm_mem_all_reduce = use_torch_symm_mem_all_reduce + self.use_hpu_communicator = use_hpu_communicator + self.use_xpu_communicator = use_xpu_communicator + self.use_npu_communicator = use_npu_communicator + self.use_message_queue_broadcaster = use_message_queue_broadcaster + + # Lazy import to avoid documentation build error + from sglang.srt.distributed.device_communicators.custom_all_reduce import ( + dispatch_custom_allreduce, + ) + from sglang.srt.distributed.device_communicators.pymscclpp import ( + PyMscclppCommunicator, + ) + from sglang.srt.distributed.device_communicators.pynccl import ( + PyNcclCommunicator, + ) + from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + is_symmetric_memory_enabled, + use_symmetric_memory, + ) + from sglang.srt.distributed.device_communicators.torch_symm_mem import ( + TorchSymmMemCommunicator, + ) + from sglang.srt.layers.dp_attention import is_allocation_symmetric + + self.is_symmetric_memory_enabled = is_symmetric_memory_enabled + self.use_symmetric_memory = use_symmetric_memory + self.is_allocation_symmetric = is_allocation_symmetric + if is_hip(): + from sglang.srt.distributed.device_communicators.quick_all_reduce import ( + QuickAllReduce, + qr_rocm_arch_available, + ) + + self.pynccl_comm: Optional[PyNcclCommunicator] = None + if use_pynccl and self.world_size > 1: + self.pynccl_comm = PyNcclCommunicator( + group=self.cpu_group, + device=self.device, + ) + + self.pymscclpp_comm: Optional[PyMscclppCommunicator] = None + if use_pymscclpp and self.world_size > 1: + self.pymscclpp_comm = PyMscclppCommunicator( + group=self.cpu_group, + device=self.device, + ) + + self.ca_comm: Optional[Any] = None + self.qr_comm: Optional[QuickAllReduce] = None + if use_custom_allreduce and self.world_size > 1: + # Initialize a custom fast all-reduce implementation. + try: + CAClass = dispatch_custom_allreduce() + self.ca_comm = CAClass( + group=self.cpu_group, + device=self.device, + ) + except Exception as e: + logger.warning( + f"Setup Custom allreduce failed with {e}. To silence this " + "warning, specify --disable-custom-all-reduce explicitly." + ) + + if is_hip(): + try: + # Initialize a custom quick all-reduce implementation for AMD + # when rocm >= gfx942. Quick reduce is designed as a + # complement to custom allreduce. + # Based on quickreduce (https://github.com/mk1-project/quickreduce). + if qr_rocm_arch_available(): + self.qr_comm = QuickAllReduce( + group=self.cpu_group, device=self.device + ) + except Exception as e: + logger.warning(f"Failed to initialize QuickAllReduce: {e}") + elif self.world_size > 1 and is_hip(): + logger.info("[AR] All-reduce call path: NCCL (custom AR disabled)") + + self.torch_symm_mem_comm: Optional[TorchSymmMemCommunicator] = None + if self.use_torch_symm_mem_all_reduce and self.world_size > 1: + self.torch_symm_mem_comm = TorchSymmMemCommunicator( + group=self.cpu_group, + device=self.device, + ) + + # Create communicator for other hardware backends + from sglang.srt.distributed.device_communicators.hpu_communicator import ( + HpuCommunicator, + ) + from sglang.srt.distributed.device_communicators.npu_communicator import ( + NpuCommunicator, + ) + from sglang.srt.distributed.device_communicators.xpu_communicator import ( + XpuCommunicator, + ) + + self.hpu_communicator: Optional[HpuCommunicator] = None + if use_hpu_communicator and self.world_size > 1: + self.hpu_communicator = HpuCommunicator(group=self.device_group) + + self.xpu_communicator: Optional[XpuCommunicator] = None + if use_xpu_communicator and self.world_size > 1: + self.xpu_communicator = XpuCommunicator(group=self.device_group) + + self.npu_communicator: Optional[NpuCommunicator] = None + if use_npu_communicator and self.world_size > 1: + self.npu_communicator = NpuCommunicator(group=self.device_group) + + # Create message queue + from sglang.srt.distributed.device_communicators.shm_broadcast import ( + MessageQueue, + ) + + self.mq_broadcaster: Optional[MessageQueue] = None + if use_message_queue_broadcaster and self.world_size > 1: + self.mq_broadcaster = MessageQueue.create_from_process_group( + self.cpu_group, 1 << 22, 6 + ) + + def __repr__(self): + return ( + f"ranks={self.ranks} rank={self.rank} local_rank={self.local_rank} use_pynccl={self.use_pynccl} " + f"device_group={self.device_group} cpu_group={self.cpu_group} unique_name={self.unique_name} " + f"world_size={self.world_size} rank_in_group={self.rank_in_group}" + ) + + @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] + + @contextmanager + def graph_capture( + self, + graph_capture_context: Optional[GraphCaptureContext] = None, + stream: Optional[torch.cuda.Stream] = None, + ): + if graph_capture_context is None: + if stream is None: + stream = self.device_module.Stream() + graph_capture_context = GraphCaptureContext(stream) + else: + stream = graph_capture_context.stream + # We don't need the context of custom quick allreduce because the ipc access + # is already collected in init() and we can capture the quick allreduce directly. + ca_comm = self.ca_comm + maybe_ca_context = nullcontext() if ca_comm is None else ca_comm.capture() + + # ensure all initialization operations complete before attempting to + # capture the graph on another stream + curr_stream = get_current_device_stream_fast() + if curr_stream != stream: + stream.wait_stream(curr_stream) + + with self.device_module.stream(stream), maybe_ca_context: + # In graph mode, we have to be very careful about the collective + # operations. The current status is: + # allreduce \ Mode | Eager | Graph | + # -------------------------------------------- + # quick allreduce | enabled | enabled | + # custom allreduce | enabled | enabled | + # PyNccl | disabled| enabled | + # PyMscclpp | disabled| enabled | + # TorchSymmMem | disabled| enabled | + # torch.distributed | enabled | disabled| + # + # Note: When custom quick allreduce is enabled, a runtime check + # will be performed. If the tensor size is too small, it will + # automatically fall back to the next available option. + # Note that custom allreduce will have a runtime check, if the + # tensor size is too large, it will fallback to the next + # available option. + # Note that the PyMsccl needs to register the tensor in ahead, + # which will introduce large overhead in the eager case, + # therefore it is only supported in the graph case. + # In summary: We select the appropriate allreduce method for + # each mode based on the algorithm order in the table and + # their usage conditions. + pynccl_comm = self.pynccl_comm + maybe_pynccl_context: Any + if not pynccl_comm: + maybe_pynccl_context = nullcontext() + else: + maybe_pynccl_context = pynccl_comm.change_state(enable=True) + + pymscclpp_comm = self.pymscclpp_comm + maybe_pymscclpp_context: Any + if not pymscclpp_comm: + maybe_pymscclpp_context = nullcontext() + else: + maybe_pymscclpp_context = pymscclpp_comm.change_state(enable=True) + with maybe_pynccl_context, maybe_pymscclpp_context: + yield graph_capture_context + + def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: + """ + User-facing all-reduce function before we actually call the + all-reduce operation. + + We need this because Dynamo does not support passing an arbitrary + object (`self` in this case) to a custom op. We need to pass the + group name as a string, and then look up the group coordinator from + the group name, dispatch the all-reduce operation to the group + coordinator. + + In addition, PyTorch custom ops do not support mutation or returning + a new tensor in the same op. So we need to figure out if the op is + in-place or out-of-place ahead of time. + """ + # Bypass the function if we are using only 1 GPU. + if self.world_size == 1: + return input_ + + # Deterministic allgather+add path + if _USE_ALLGATHER_ADD: + return self._allgather_add(input_) + + if input_.is_cpu: + if is_shm_available(input_.dtype, self.world_size, self.local_size): + torch.ops.sgl_kernel.shm_allreduce(input_, REDUCE_OP_SUM) + else: + torch.distributed.all_reduce(input_, group=self.device_group) + return input_ + + if self.hpu_communicator is not None and not self.hpu_communicator.disabled: + return self.hpu_communicator.all_reduce(input_) + + if self.xpu_communicator is not None and not self.xpu_communicator.disabled: + return self.xpu_communicator.all_reduce(input_) + + if self.npu_communicator is not None and not self.npu_communicator.disabled: + return self.npu_communicator.all_reduce(input_) + + if self.pynccl_comm is not None and self.is_symmetric_memory_enabled(): + with self.pynccl_comm.change_state(enable=True): + self.pynccl_comm.all_reduce(input_) + return input_ + + outplace_all_reduce_method = None + if ( + self.ca_comm is not None + and not self.ca_comm.disabled + and self.ca_comm.should_custom_ar(input_) + ): + outplace_all_reduce_method = "ca" + elif ( + self.qr_comm is not None + and not self.qr_comm.disabled + and self.qr_comm.should_quick_allreduce(input_) + ): + outplace_all_reduce_method = "qr" + elif ( + self.pymscclpp_comm is not None + and not self.pymscclpp_comm.disabled + and self.pymscclpp_comm.should_mscclpp_allreduce(input_) + ): + outplace_all_reduce_method = "pymscclpp" + elif ( + self.torch_symm_mem_comm is not None + and not self.torch_symm_mem_comm.disabled + and self.torch_symm_mem_comm.should_torch_symm_mem_allreduce(input_) + ): + outplace_all_reduce_method = "torch_symm_mem" + elif is_in_piecewise_cuda_graph(): + # For piecewise cuda graph, we use pynccl outplace allreduce + outplace_all_reduce_method = "pynccl" + if outplace_all_reduce_method is not None: + return outplace_all_reduce( + input_, + group_name=self.unique_name, + outplace_all_reduce_method=outplace_all_reduce_method, + ) + else: + inplace_all_reduce(input_, group_name=self.unique_name) + return input_ + + def _allgather_add(self, input_: torch.Tensor) -> torch.Tensor: + """Deterministic allreduce via allgather + local sum (fixed rank order).""" + gathered = torch.empty( + (self.world_size,) + input_.shape, + dtype=input_.dtype, + device=input_.device, + ) + torch.distributed.all_gather_into_tensor( + gathered, input_.contiguous(), group=self.device_group + ) + return gathered.sum(dim=0) + + def fused_allreduce_rmsnorm( + self, + input_: torch.Tensor, + residual_inp_: torch.Tensor, + weight_: torch.Tensor, + eps: float, + ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Attempt fused all-reduce + RMSNorm via custom all-reduce communicator.""" + ca_comm = self.ca_comm + if ca_comm is None or getattr(ca_comm, "disabled", True): + return None + + # Prefer communicator-native fused API when provided. + if hasattr(ca_comm, "fused_allreduce_rmsnorm"): + try: + return ca_comm.fused_allreduce_rmsnorm( + input_, residual_inp_, weight_, eps + ) + except Exception: + # Fall back to custom_fused_ar_rms path below. + pass + + if not hasattr(ca_comm, "custom_fused_ar_rms"): + return None + + # 1-stage policy for fused AR+RMSNorm: + # 1) Explicit env override wins. + # 2) Deterministic inference forces 1-stage for reproducibility. + # 3) Otherwise follow AITER's heuristic (small payloads only). + if envs.SGLANG_USE_1STAGE_ALLREDUCE.is_set(): + use_1stage_ar = envs.SGLANG_USE_1STAGE_ALLREDUCE.get() + elif envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get(): + use_1stage_ar = True + else: + total_bytes = input_.numel() * input_.element_size() + hidden_dim = input_.shape[-1] + use_1stage_ar = total_bytes <= 128 * 1024 and hidden_dim in { + 512, + 1024, + 2048, + 2880, + 4096, + } + + fused_outputs = ca_comm.custom_fused_ar_rms( + input_, + residual_inp_, + weight_, + eps, + use_1stage_ar, + ) + return fused_outputs + + def _all_reduce_out_place( + self, input_: torch.Tensor, outplace_all_reduce_method: str + ) -> torch.Tensor: + ca_comm = self.ca_comm + qr_comm = self.qr_comm + pymscclpp_comm = self.pymscclpp_comm + torch_symm_mem_comm = self.torch_symm_mem_comm + pynccl_comm = self.pynccl_comm + assert any([qr_comm, ca_comm, pymscclpp_comm, torch_symm_mem_comm, pynccl_comm]) + if outplace_all_reduce_method == "ca": + assert not ca_comm.disabled + out = ca_comm.custom_all_reduce(input_) + elif outplace_all_reduce_method == "qr": + assert not qr_comm.disabled + out = qr_comm.quick_all_reduce(input_) + elif outplace_all_reduce_method == "torch_symm_mem": + assert not torch_symm_mem_comm.disabled + out = torch_symm_mem_comm.all_reduce(input_) + elif outplace_all_reduce_method == "pymscclpp": + assert not pymscclpp_comm.disabled + out = pymscclpp_comm.all_reduce(input_) + elif outplace_all_reduce_method == "pynccl": + with pynccl_comm.change_state(enable=True): + out = pynccl_comm.outplace_all_reduce(input_) + assert out is not None + return out + + def _all_reduce_in_place(self, input_: torch.Tensor) -> None: + pynccl_comm = self.pynccl_comm + torch_symm_mem_comm = self.torch_symm_mem_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.all_reduce(input_) + elif torch_symm_mem_comm is not None and not torch_symm_mem_comm.disabled: + torch_symm_mem_comm.all_reduce(input_) + else: + torch.distributed.all_reduce(input_, group=self.device_group) + + def _reduce_scatter_tensor( + self, + output: torch.Tensor, + input: torch.Tensor, + ) -> torch.Tensor: + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and ( + not pynccl_comm.disabled or self.is_symmetric_memory_enabled() + ): + with pynccl_comm.change_state(enable=True): + pynccl_comm.reduce_scatter(output, input) + else: + torch.distributed.reduce_scatter_tensor( + output, input, group=self.device_group + ) + return output + + def reduce_scatter_tensor(self, output: torch.Tensor, input: torch.Tensor): + if _is_npu or _is_cuda_alike: + self._reduce_scatter_tensor(output, input) + else: + reg_reduce_scatter_tensor(output, input, group_name=self.unique_name) + + def reduce_scatter( + self, + output: torch.Tensor, + input_list: List[torch.Tensor], + ) -> None: + # TODO(ch-wan): support other backends + torch.distributed.reduce_scatter(output, input_list, group=self.device_group) + return output + + def reduce_scatterv( + self, + input_: torch.Tensor, + output: Optional[torch.Tensor] = None, + sizes: Optional[List[int]] = None, + ) -> torch.Tensor: + world_size = self.world_size + pynccl_comm = self.pynccl_comm + + with pynccl_comm.change_state(enable=True): + assert ( + pynccl_comm is not None and not pynccl_comm.disabled + ), "pynccl is required for reduce_scatterv" + + if sizes is not None: + assert len(sizes) == world_size + assert input_.shape[0] == sum(sizes) + chunk_size = sizes[self.rank_in_group] + else: + assert input_.shape[0] % world_size == 0 + chunk_size = input_.shape[0] // world_size + output_shape = (chunk_size,) + input_.shape[1:] + + if output is None: + output = torch.empty( + output_shape, dtype=input_.dtype, device=input_.device + ) + else: + assert output.shape == output_shape + + pynccl_comm.reduce_scatter(output, input_, sizes=sizes) + return output + + def _all_gather_into_tensor(self, output: torch.Tensor, input: torch.Tensor): + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and ( + not pynccl_comm.disabled or self.is_symmetric_memory_enabled() + ): + with pynccl_comm.change_state(enable=True): + pynccl_comm.all_gather(output, input) + else: + torch.distributed.all_gather_into_tensor( + output, input, group=self.device_group + ) + + def all_gather_into_tensor(self, output: torch.Tensor, input: torch.Tensor): + if _is_npu or _is_xpu: + self._all_gather_into_tensor(output, input) + else: + reg_all_gather_into_tensor(output, input, group_name=self.unique_name) + + def cp_all_gather_into_tensor_async( + self, output: torch.Tensor, input: torch.Tensor, stream: torch.cuda.Stream + ): + """ + Implement an asynchronous `allgather` operation on a specified stream. + (the default `torch.distributed.all_gather_into_tensor` will trigger event synchronization), + eliminating the CPU-side launch-kernel blocking issue caused by synchronization problems. + The specific implementation uses the interface provided by pynccl to remove the synchronization logic of events. + """ + pynccl_comm = self.pynccl_comm + if pynccl_comm is None or pynccl_comm.disabled: + self.all_gather_into_tensor(output, input) + else: + pynccl_comm.cp_all_gather_into_tensor(output, input, stream=stream) + + def all_to_all(self, output: torch.Tensor, input: torch.Tensor) -> None: + torch.distributed.all_to_all_single(output, input, group=self.device_group) + + def all_gather( + self, + input_: torch.Tensor, + dim: int = -1, + output_tensor_list: Optional[List[torch.Tensor]] = None, + ) -> torch.Tensor: + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + if output_tensor_list is not None: + logger.warning( + "Performing in-place all-gather with a group size of 1. " + "This may be unnecessary; consider bypassing it for better efficiency." + ) + output_tensor_list[0].copy_(input_) + return None + else: + return input_ + + if output_tensor_list is not None: + # TODO(ch-wan): support other backends + return torch.distributed.all_gather( + output_tensor_list, input_, group=self.device_group + ) + + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + + # For HPUs, use HPU communicator. + hpu_comm = self.hpu_communicator + if hpu_comm is not None and not hpu_comm.disabled: + return hpu_comm.all_gather(input_, dim) + + # For NPUs, use NPU communicator. + npu_comm = self.npu_communicator + if npu_comm is not None and not npu_comm.disabled: + return npu_comm.all_gather(input_, dim) + + # For XPUs / P800, avoid all_gather_into_tensor which can segfault on BKCL; + # use the list-based all_gather + cat fallback instead. + xpu_comm = self.xpu_communicator + if ( + xpu_comm is not None and not xpu_comm.disabled + ) or os.environ.get("SGLANG_P800_ALL_GATHER_FALLBACK"): + input_ = input_.contiguous() + output_tensor_list = [ + torch.empty_like(input_) for _ in range(world_size) + ] + torch.distributed.all_gather( + output_tensor_list, input_, group=self.device_group + ) + return torch.cat(output_tensor_list, dim=dim) + + 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] * world_size,) + input_size[1:] + # Allocate output tensor. + with self.use_symmetric_memory( + self, disabled=not self.is_allocation_symmetric() + ): + output_tensor = torch.empty( + output_size, dtype=input_.dtype, device=input_.device + ) + + # All-gather. + if input_.is_cpu: + if is_shm_available(input_.dtype, self.world_size, self.local_size): + return torch.ops.sgl_kernel.shm_allgather(input_, dim) + else: + torch.distributed.all_gather_into_tensor( + output_tensor, input_, group=self.device_group + ) + else: + self.all_gather_into_tensor(output_tensor, input_) + + # Reshape + 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 + + def all_gatherv( + self, + input_: Union[torch.Tensor, List[torch.Tensor]], + sizes: Optional[List[int]] = None, + ) -> Union[torch.Tensor, List[torch.Tensor]]: + """ + Supports varying sizes per rank and input tensor list. + `sizes`: a list of len(world_size) with the number of items per rank to gather. + """ + world_size = self.world_size + pynccl_comm = self.pynccl_comm + + with pynccl_comm.change_state(enable=True): + assert ( + pynccl_comm is not None and not pynccl_comm.disabled + ), "pynccl is required for all_gatherv" + + def _all_gather_allocate_output( + input_: torch.Tensor, sizes: Optional[List[int]] = None + ): + input_size = input_.size() + if sizes is not None: + assert len(sizes) == world_size + assert input_.shape[0] == sizes[self.rank_in_group] + output_size = (sum(sizes),) + input_size[1:] + # 'sizes' is not needed if all inputs in the same group have the same shape + if all(s == sizes[0] for s in sizes): + sizes = None + else: + output_size = (input_size[0] * world_size,) + input_size[1:] + # Allocate output tensor. + with self.use_symmetric_memory(self, disabled=sizes is not None): + output_tensor = torch.empty( + output_size, dtype=input_.dtype, device=input_.device + ) + return output_tensor, sizes + + if isinstance(input_, torch.Tensor): + input_ = [input_] + + output_list = [] + size_list = [] + for inp in input_: + output_tensor, s = _all_gather_allocate_output(inp, sizes=sizes) + output_list.append(output_tensor) + size_list.append(s) + + pynccl_comm.group_start() + for i, inp in enumerate(input_): + pynccl_comm.all_gather(output_list[i], inp, sizes=size_list[i]) + pynccl_comm.group_end() + + return output_list + + def gather( + self, input_: torch.Tensor, dst: int = 0, dim: int = -1 + ) -> Optional[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() + if self.xpu_communicator is not None and not self.xpu_communicator.disabled: + return self.xpu_communicator.gather(input_, self.rank_in_group, dst, 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): + """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 + ) + 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.mq_broadcaster is not None: + assert src == 0, "Message queue broadcaster only supports src=0" + return self.mq_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 all_gather_object(self, obj: Any) -> List[Any]: + objs = [None] * self.world_size + torch.distributed.all_gather_object(objs, obj, group=self.cpu_group) + return objs + + def send_object( + self, + obj: Any, + dst: int, + async_send: bool = False, + ) -> List[P2PWork]: + """ + Send the input object list to the destination rank. + This function uses the CPU group for all communications. + + TODO: If you want to use GPU communication, please add a new argument (e.g., data_group, group), + use other functions (e.g., send), or implement a new function (e.g., send_object_device). + + NOTE: `dst` is the local rank of the destination rank. + """ + + assert dst < self.world_size, f"Invalid dst rank ({dst})" + assert dst != self.rank_in_group, ( + "Invalid destination rank. Destination rank is the same " + "as the current rank." + ) + send_func = torch.distributed.isend if async_send else torch.distributed.send + + # 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 + p2p_work = [] + size_work = send_func( + size_tensor, + self.ranks[dst], + group=self.cpu_group, + ) + if async_send: + p2p_work.append(P2PWork(size_work, size_tensor)) + + object_work = send_func( + object_tensor, + self.ranks[dst], + group=self.cpu_group, + ) + if async_send: + p2p_work.append(P2PWork(object_work, object_tensor)) + + return p2p_work + + 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_in_group + ), "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 + # We have to use irecv here to make it work for both isend and send. + work = torch.distributed.irecv( + size_tensor, src=self.ranks[src], group=self.cpu_group + ) + work.wait() + + # Tensor to receive serialized objects into. + object_tensor: Any = torch.empty( # type: ignore[call-overload] + size_tensor.item(), # type: ignore[arg-type] + dtype=torch.uint8, + device="cpu", + ) + + work = torch.distributed.irecv( + object_tensor, src=self.ranks[src], group=self.cpu_group + ) + work.wait() + + obj = pickle.loads(object_tensor.numpy()) + 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})" + + rank_in_group = self.rank_in_group + if rank_in_group == 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=self.ranks[src], group=metadata_group, async_op=True + ) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[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. + tensor_dict[key] = tensor + continue + if tensor.is_cpu: + # use metadata_group for CPU tensors + handle = torch.distributed.broadcast( + tensor, + src=self.ranks[src], + group=metadata_group, + async_op=True, + ) + else: + # use group for GPU tensors + handle = torch.distributed.broadcast( + tensor, src=self.ranks[src], group=group, async_op=True + ) + async_handles.append(handle) + tensor_dict[key] = tensor + else: + 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, + all_gather_group: Optional["GroupCoordinator"] = None, + async_send: bool = False, + ) -> Optional[List[P2PWork]]: + """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 self.world_size == 1: + return tensor_dict + + all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size + all_gather_rank = ( + 0 if all_gather_group is None else all_gather_group.rank_in_group + ) + + group = self.device_group + metadata_group = self.cpu_group + + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + assert dst < self.world_size, f"Invalid dst rank ({dst})" + + assert isinstance( + tensor_dict, dict + ), f"Expecting a dictionary, got {type(tensor_dict)}" + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + # Note: While switching to Device-to-Device (D2D) would introduce an extra + # Device-to-Host (D2H) memory copy overhead for serialization, our benchmarks + # show better overall transmission performance with D2D due to: + # 1. Superior D2D transfer bandwidth + # 2. Ability to overlap send and recv operations + # Thus the net performance gain justifies this approach. + + send_func = torch.distributed.isend if async_send else torch.distributed.send + p2p_works = self.send_object(metadata_list, dst=dst, async_send=async_send) + + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip sending empty tensors. + continue + + # send-allgather: send only a slice, then do allgather. + if all_gather_group is not None and tensor.numel() % all_gather_size == 0: + tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] + + comm_group = metadata_group if tensor.is_cpu else group + work = send_func(tensor, self.ranks[dst], group=comm_group) + if async_send: + p2p_works.append(P2PWork(work, tensor)) + return p2p_works + + def recv_tensor_dict( + self, + src: Optional[int] = None, + all_gather_group: Optional["GroupCoordinator"] = 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 + + all_gather_size = 1 if all_gather_group is None else all_gather_group.world_size + all_gather_rank = ( + 0 if all_gather_group is None else all_gather_group.rank_in_group + ) + + group = self.device_group + metadata_group = self.cpu_group + + if src is None: + src = (self.rank_in_group - 1) % self.world_size + 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. + tensor_dict[key] = tensor + continue + + # send-allgather: send only a slice, then do allgather. + use_all_gather = ( + all_gather_group is not None + and tensor.numel() % all_gather_size == 0 + ) + + if use_all_gather: + orig_shape = tensor.shape + tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] + + # We have to use irecv here to make it work for both isend and send. + comm_group = metadata_group if tensor.is_cpu else group + work = torch.distributed.irecv( + tensor, src=self.ranks[src], group=comm_group + ) + work.wait() + + if use_all_gather: + tensor = all_gather_group.all_gather(tensor, dim=0) + tensor = tensor.reshape(orig_shape) + + tensor_dict[key] = tensor + else: + 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 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: Optional[int] = 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): + 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.pynccl_comm is not None: + self.pynccl_comm = None + if self.ca_comm is not None: + self.ca_comm = None + if self.mq_broadcaster is not None: + self.mq_broadcaster = None + + +_WORLD: Optional[GroupCoordinator] = 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_pynccl=False, + use_pymscclpp=False, + use_custom_allreduce=False, + use_torch_symm_mem_all_reduce=False, + use_hpu_communicator=False, + use_xpu_communicator=False, + use_npu_communicator=False, + group_name="world", + ) + + +def init_model_parallel_group( + group_ranks: List[List[int]], + local_rank: int, + backend: str, + use_pynccl: Optional[bool] = None, + use_custom_allreduce: Optional[bool] = None, + use_message_queue_broadcaster: bool = False, + group_name: Optional[str] = None, + use_mscclpp_allreduce: Optional[bool] = None, + use_torch_symm_mem_allreduce: Optional[bool] = None, +) -> GroupCoordinator: + if use_custom_allreduce is None: + use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE + if use_mscclpp_allreduce is None: + use_mscclpp_allreduce = _ENABLE_MSCCLPP_ALL_REDUCE + if use_torch_symm_mem_allreduce is None: + use_torch_symm_mem_allreduce = _ENABLE_TORCH_SYMM_MEM_ALL_REDUCE + return GroupCoordinator( + group_ranks=group_ranks, + local_rank=local_rank, + torch_distributed_backend=backend, + use_pynccl=False, + use_pymscclpp=use_mscclpp_allreduce, + use_custom_allreduce=use_custom_allreduce, + use_torch_symm_mem_all_reduce=use_torch_symm_mem_allreduce, + use_hpu_communicator=True, + use_xpu_communicator=True, + use_npu_communicator=True, + use_message_queue_broadcaster=use_message_queue_broadcaster, + group_name=group_name, + ) + + +_TP: Optional[GroupCoordinator] = None +_ATTN_TP: Optional[GroupCoordinator] = None +_ATTN_CP: Optional[GroupCoordinator] = None + +# duplicate GroupCoordinator for prefill in PD-Multiplexing +_PDMUX_PREFILL_TP_GROUP: Optional[GroupCoordinator] = None + +_ENABLE_PDMUX_P_TP: bool = False + + +def set_pdmux_status(enable_prefill_multiplexing: bool): + global _ENABLE_PDMUX_P_TP + _ENABLE_PDMUX_P_TP = enable_prefill_multiplexing + + +def get_tp_group() -> GroupCoordinator: + if _ENABLE_PDMUX_P_TP: + assert ( + _PDMUX_PREFILL_TP_GROUP is not None + ), "tensor model parallel group for PD-Multiplexing Prefill is not initialized" + return _PDMUX_PREFILL_TP_GROUP + assert _TP is not None, "tensor model parallel group is not initialized" + return _TP + + +def get_attn_tp_group() -> GroupCoordinator: + assert ( + _ATTN_TP is not None + ), "attention tensor model parallel group is not initialized" + return _ATTN_TP + + +def get_attn_cp_group() -> GroupCoordinator: + assert ( + _ATTN_CP is not None + ), "attention context model parallel group is not initialized" + return _ATTN_CP + + +_MOE_DP: Optional[GroupCoordinator] = None +_MOE_EP: Optional[GroupCoordinator] = None +_MOE_TP: Optional[GroupCoordinator] = None + + +def get_moe_dp_group() -> GroupCoordinator: + assert _MOE_DP is not None, "moe data parallel group is not initialized" + return _MOE_DP + + +def get_moe_ep_group() -> GroupCoordinator: + assert _MOE_EP is not None, "expert model parallel group is not initialized" + return _MOE_EP + + +def get_moe_tp_group() -> GroupCoordinator: + assert _MOE_TP is not None, "expert model parallel group is not initialized" + return _MOE_TP + + +# kept for backward compatibility +get_tensor_model_parallel_group = get_tp_group + +_PP: Optional[GroupCoordinator] = None + + +def get_pp_group() -> GroupCoordinator: + assert _PP is not None, "pipeline model parallel group is not initialized" + return _PP + + +# kept for backward compatibility +get_pipeline_model_parallel_group = get_pp_group + + +def get_mooncake_transfer_engine(): + """ + Return the shared MooncakeTransferEngine if initialized in device_communicators, + else None. Used by disaggregation mooncake backend and mem_cache mooncake_store. + """ + from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import ( + get_mooncake_transfer_engine as _get_engine, + ) + + return _get_engine() + + +@contextmanager +def graph_capture(stream: Optional[torch.cuda.Stream] = None): + """ + `graph_capture` is a context manager which should surround the code that + is capturing the CUDA graph. Its main purpose is to ensure that the + some operations will be run after the graph is captured, before the graph + is replayed. It returns a `GraphCaptureContext` object which contains the + necessary data for the graph capture. Currently, it only contains the + stream that the graph capture is running on. This stream is set to the + current CUDA stream when the context manager is entered and reset to the + default stream when the context manager is exited. This is to ensure that + the graph capture is running on a separate stream from the default stream, + in order to explicitly distinguish the kernels to capture + from other kernels possibly launched on background in the default stream. + """ + with get_tp_group().graph_capture( + stream=stream + ) as context, get_pp_group().graph_capture(context): + with contextlib.ExitStack() as stack: + seen = {id(_TP)} + for group in (_MOE_EP, _MOE_TP): + if group is not None and id(group) not in seen: + seen.add(id(group)) + stack.enter_context(group.graph_capture(context)) + yield context + + +logger = logging.getLogger(__name__) + +_ENABLE_CUSTOM_ALL_REDUCE = True +_ENABLE_MSCCLPP_ALL_REDUCE = False +_ENABLE_TORCH_SYMM_MEM_ALL_REDUCE = False + + +def set_custom_all_reduce(enable: bool): + global _ENABLE_CUSTOM_ALL_REDUCE + _ENABLE_CUSTOM_ALL_REDUCE = enable + + +def set_mscclpp_all_reduce(enable: bool): + global _ENABLE_MSCCLPP_ALL_REDUCE + _ENABLE_MSCCLPP_ALL_REDUCE = enable + + +def set_torch_symm_mem_all_reduce(enable: bool): + global _ENABLE_TORCH_SYMM_MEM_ALL_REDUCE + _ENABLE_TORCH_SYMM_MEM_ALL_REDUCE = enable + + +_DEVICE_TO_DISTRIBUTED_BACKEND = { + "cuda": "nccl", + "xpu": "xccl", + "hpu": "hccl", + "cpu": "gloo", + "npu": "hccl", + "musa": "mccl", +} + + +def get_default_distributed_backend(device: str) -> str: + return _DEVICE_TO_DISTRIBUTED_BACKEND.get(device, "gloo") + + +def _create_global_tcp_store(rank: int, world_size: int) -> None: + """Create a global TCPStore for coordination across ranks. + + This function creates a TCPStore that all ranks can use for coordination + (e.g., for NIXL buffer setup). + """ + from torch.distributed import TCPStore + + master_ip = os.environ.get("MASTER_ADDR") + + if not master_ip: + logger.warning( + "Could not determine master IP for global TCPStore. " + "Broadcasting from rank 0 to all ranks." + ) + + base_store_port = envs.SGLANG_TCP_STORE_PORT.get() + + # Rank 0 gets its local IP and broadcasts it to all ranks + # Use broadcast_object_list which works with any backend (handles CPU/GPU automatically) + if not master_ip: + if rank == 0: + master_ip = get_local_ip_auto() + ip_list = [master_ip] + else: + ip_list = [None] + + torch.distributed.broadcast_object_list(ip_list, src=0) + master_ip = ip_list[0] + + try: + tcp_store = TCPStore( + host_name=master_ip, + port=base_store_port, + world_size=world_size, + is_master=(rank == 0), + ) + set_global_tcp_store(tcp_store) + logger.info( + "Created global TCPStore at %s:%d (rank=%d, world_size=%d)", + master_ip, + base_store_port, + rank, + world_size, + ) + except Exception as e: + logger.warning( + "Failed to create global TCPStore at %s:%d: %s. " + "Components requiring TCPStore (like NIXL) may not work.", + master_ip, + base_store_port, + e, + ) + + +def init_distributed_environment( + world_size: int = -1, + rank: int = -1, + distributed_init_method: str = "env://", + local_rank: int = -1, + backend: str = "nccl", + timeout: Optional[int] = None, + moe_a2a_backend: Optional[str] = None, +): + 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 "mooncake" in backend: + try: + from mooncake import ep as mooncake_ep + except ImportError as e: + raise ImportError( + "Please install mooncake by following the instructions at " + "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md " # noqa: E501 + "to run SGLang with Mooncake Backend." + ) from e + mooncake_ep.set_host_ip(get_local_ip_auto()) + + if not torch.distributed.is_initialized(): + assert distributed_init_method is not None, ( + "distributed_init_method must be provided when initializing " + "distributed environment" + ) + if timeout is not None: + assert isinstance(timeout, (int)), "timeout must be a number" + assert timeout > 0, "timeout must be positive" + timeout = timedelta(seconds=timeout) + + pg_options = get_torch_distributed_pg_options() + + # this backend is used for WORLD + torch.distributed.init_process_group( + backend=backend, + init_method=distributed_init_method, + world_size=world_size, + rank=rank, + timeout=timeout, + pg_options=pg_options, + ) + + # Create a global TCPStore for coordination (used by NIXL) + if moe_a2a_backend == "nixl": + _create_global_tcp_store(rank, world_size) + + # 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 = int(os.environ.get("LOCAL_RANK", "0")) + else: + local_rank = rank + global _WORLD + if _WORLD is None: + ranks = list(range(torch.distributed.get_world_size())) + _WORLD = init_world_group(ranks, local_rank, backend) + else: + assert ( + _WORLD.world_size == torch.distributed.get_world_size() + ), "world group already initialized with a different world size" + + +def initialize_model_parallel( + tensor_model_parallel_size: int = 1, + expert_model_parallel_size: int = 1, + pipeline_model_parallel_size: int = 1, + attention_data_parallel_size: int = 1, + attention_context_model_parallel_size: int = 1, + moe_data_model_parallel_size: int = 1, + backend: Optional[str] = None, + duplicate_tp_group: bool = False, +) -> None: + """ + Initialize model parallel groups. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model + parallelism. + expert_model_parallel_size: number of GPUs used for expert model + parallelism. + pipeline_model_parallel_size: number of GPUs used for pipeline model + parallelism. + attention_data_parallel_size: number of GPUs used for attention data + parallelism. + attention_context_model_parallel_size: number of GPUs used for attention context + parallelism. + moe_data_model_parallel_size: number of GPUs used for moe data + parallelism. + + Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we + use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize + the model pipeline. The present function will + create 4 tensor model-parallel groups and 2 pipeline model-parallel groups: + 4 tensor model-parallel groups: + [g0, g1], [g2, g3], [g4, g5], [g6, g7] + 2 pipeline model-parallel groups: + [g0, g2, g4, g6], [g1, g3, g5, g7] + + Let's say we use 2 GPUs for attention context parallelism (attn_cp_size=2) and 4 GPUs for + attention tensor parallelism (attn_tp_size=4). As for MoE part, we use 2 GPUs for moe data + parallelism (moe_dp_size=2) and 4 GPUs for moe expert parallelism (moe_ep_size=4). The present + function will create the following groups: + 2 tensor model-parallel groups: + [g0, g1, g2, g3], [g4, g5, g6, g7] + 4 attention context-parallel groups: + [g0, g4], [g1, g5], [g2, g6], [g3, g7] + 2 moe expert-parallel groups: + [g0, g1, g2, g3], [g4, g5, g6, g7] + 4 moe data-parallel groups: + [g0, g4], [g1, g5], [g2, g6], [g3, g7] + + 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. + """ + # 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) + + if world_size != tensor_model_parallel_size * pipeline_model_parallel_size: + raise RuntimeError( + f"world_size ({world_size}) is not equal to " + f"tensor_model_parallel_size ({tensor_model_parallel_size}) x " + f"pipeline_model_parallel_size ({pipeline_model_parallel_size})" + ) + + # Build the tensor model-parallel groups. + 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 tp_group_idx in range(num_tensor_model_parallel_groups): + ranks = list( + range( + tp_group_idx * tensor_model_parallel_size, + (tp_group_idx + 1) * tensor_model_parallel_size, + ) + ) + group_ranks.append(ranks) + + # message queue broadcaster is only used in tensor model parallel group + _TP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=get_bool_env_var( + "SGLANG_USE_MESSAGE_QUEUE_BROADCASTER", "true" + ), + group_name="tp", + ) + + if duplicate_tp_group: + global _PDMUX_PREFILL_TP_GROUP + assert ( + _PDMUX_PREFILL_TP_GROUP is None + ), "tensor model parallel group for PD-Multiplexing Prefill is already initialized" + _PDMUX_PREFILL_TP_GROUP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=get_bool_env_var( + "SGLANG_USE_MESSAGE_QUEUE_BROADCASTER", "true" + ), + group_name="pdmux_prefill_tp", + ) + if _TP.pynccl_comm: + _TP.pynccl_comm.disabled = False + _PDMUX_PREFILL_TP_GROUP.pynccl_comm.disabled = False + + attn_dp_size = attention_data_parallel_size + attn_cp_size = attention_context_model_parallel_size + attn_tp_size = tensor_model_parallel_size // attn_cp_size // attn_dp_size + + global _ATTN_CP + assert ( + _ATTN_CP is None + ), "attention context model parallel group is already initialized" + if attn_cp_size == tensor_model_parallel_size: + _ATTN_CP = _TP + else: + group_ranks = [] + for tp_group_idx in range(num_tensor_model_parallel_groups): + for dp_idx in range(attn_dp_size): + for attn_tp_idx in range(attn_tp_size): + st = ( + tp_group_idx * tensor_model_parallel_size + + dp_idx * attn_tp_size * attn_cp_size + + attn_tp_idx + ) + en = ( + tp_group_idx * tensor_model_parallel_size + + (dp_idx + 1) * attn_tp_size * attn_cp_size + + attn_tp_idx + ) + ranks = list(range(st, en, attn_tp_size)) + group_ranks.append(ranks) + _ATTN_CP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="attn_cp", + ) + + from sglang.srt.layers.sampler import SYNC_TOKEN_IDS_ACROSS_TP + + global _ATTN_TP + assert ( + _ATTN_TP is None + ), "attention tensor model parallel group is already initialized" + if attn_tp_size == tensor_model_parallel_size: + _ATTN_TP = _TP + else: + group_ranks = [] + for tp_group_idx in range(num_tensor_model_parallel_groups): + for cp_dp_combined_idx in range(attn_cp_size * attn_dp_size): + st = ( + tp_group_idx * tensor_model_parallel_size + + cp_dp_combined_idx * attn_tp_size + ) + en = ( + tp_group_idx * tensor_model_parallel_size + + (cp_dp_combined_idx + 1) * attn_tp_size + ) + ranks = list(range(st, en)) + group_ranks.append(ranks) + _ATTN_TP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_pynccl=SYNC_TOKEN_IDS_ACROSS_TP, + use_mscclpp_allreduce=False, + use_custom_allreduce=False, + use_torch_symm_mem_allreduce=False, + group_name="attention_tp", + ) + + moe_ep_size = expert_model_parallel_size + moe_dp_size = moe_data_model_parallel_size + moe_tp_size = tensor_model_parallel_size // moe_ep_size // moe_dp_size + + global _MOE_DP + assert _MOE_DP is None, "moe data parallel group is already initialized" + # gpus_per_pp_stage = tensor_model_parallel_size * attention_context_model_parallel_size + if moe_dp_size == tensor_model_parallel_size: + _MOE_DP = _TP + else: + group_ranks = [] + for tp_group_idx in range(num_tensor_model_parallel_groups): + for tp_ep_combined_idx in range(moe_tp_size * moe_ep_size): + st = tp_group_idx * tensor_model_parallel_size + tp_ep_combined_idx + en = ( + tp_group_idx + 1 + ) * tensor_model_parallel_size + tp_ep_combined_idx + ranks = list(range(st, en, moe_tp_size * moe_ep_size)) + group_ranks.append(ranks) + _MOE_DP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="moe_dp", + ) + + global _MOE_EP + assert _MOE_EP is None, "expert model parallel group is already initialized" + if moe_ep_size == tensor_model_parallel_size: + _MOE_EP = _TP + else: + # TODO(ch-wan): use split_group to save memory + group_ranks = [] + for tp_group_idx in range(num_tensor_model_parallel_groups): + for moe_dp_idx in range(moe_dp_size): + for moe_tp_idx in range(moe_tp_size): + st = ( + tp_group_idx * tensor_model_parallel_size + + moe_dp_idx * moe_ep_size * moe_tp_size + + moe_tp_idx + ) + en = st + moe_ep_size * moe_tp_size + ranks = list(range(st, en, moe_tp_size)) + group_ranks.append(ranks) + _MOE_EP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="moe_ep", + ) + + global _MOE_TP + assert _MOE_TP is None, "expert model parallel group is already initialized" + if moe_tp_size == tensor_model_parallel_size: + _MOE_TP = _TP + else: + # TODO(ch-wan): use split_group to save memory + group_ranks = [] + for tp_group_idx in range(num_tensor_model_parallel_groups): + for ep_dp_combined_idx in range(moe_ep_size * moe_dp_size): + st = ( + tp_group_idx * tensor_model_parallel_size + + ep_dp_combined_idx * moe_tp_size + ) + en = ( + tp_group_idx * tensor_model_parallel_size + + (ep_dp_combined_idx + 1) * moe_tp_size + ) + ranks = list(range(st, en)) + group_ranks.append(ranks) + _MOE_TP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + group_name="moe_tp", + ) + + # Build the pipeline model-parallel groups. + num_pipeline_model_parallel_groups: int = world_size // pipeline_model_parallel_size + global _PP + assert _PP is None, "pipeline model parallel group is already initialized" + group_ranks = [] + for pp_group_idx in range(num_pipeline_model_parallel_groups): + ranks = list( + range(pp_group_idx, world_size, num_pipeline_model_parallel_groups) + ) + group_ranks.append(ranks) + # pipeline parallel does not need custom allreduce + _PP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_custom_allreduce=False, + group_name="pp", + ) + + +def create_custom_parallel_group( + group_ranks: List[int], backend: str = "gloo" +) -> Optional[torch.distributed.ProcessGroup]: + """ + Create a custom parallel group based on the provided ranks. + + Args: + group_ranks: The list of ranks that the CURRENT process wants to join. + (e.g., Rank 0 passes [0...7], Rank 8 passes [8...15]) + backend: The communication backend (default: "gloo"). + + Returns: + The ProcessGroup if the current rank is in group_ranks, else None. + """ + assert torch.distributed.is_initialized() + + world_size = torch.distributed.get_world_size() + rank = torch.distributed.get_rank() + + local_config = sorted(list(set(group_ranks))) + gathered_configs = [None for _ in range(world_size)] + + torch.distributed.all_gather_object(gathered_configs, local_config) + + unique_groups = [] + seen_signatures = set() + + for config in gathered_configs: + config_tuple = tuple(config) + if config_tuple not in seen_signatures: + seen_signatures.add(config_tuple) + unique_groups.append(list(config_tuple)) + + unique_groups.sort(key=lambda x: x[0]) + + my_new_group = None + + for g_ranks in unique_groups: + group = torch.distributed.new_group(ranks=g_ranks, backend=backend) + + if set(g_ranks) == set(local_config): + my_new_group = group + logger.debug( + f"Rank {rank} successfully created/joined custom group: {g_ranks}" + ) + + return my_new_group + + +def ensure_model_parallel_initialized( + tensor_model_parallel_size: int, + expert_model_parallel_size: int, + pipeline_model_parallel_size: int, + backend: Optional[str] = None, +) -> None: + """Helper to initialize model parallel groups if they are not initialized, + or ensure tensor-parallel and pipeline-parallel sizes are equal to expected + values if the model parallel groups are initialized. + """ + backend = backend or torch.distributed.get_backend(get_world_group().device_group) + if not model_parallel_is_initialized(): + initialize_model_parallel( + tensor_model_parallel_size, + expert_model_parallel_size, + pipeline_model_parallel_size, + backend, + ) + return + + assert get_tensor_model_parallel_world_size() == tensor_model_parallel_size, ( + "tensor parallel group already initialized, but of unexpected size: " + f"{get_tensor_model_parallel_world_size()=} vs. " + f"{tensor_model_parallel_size=}" + ) + pp_world_size = get_pp_group().world_size + assert pp_world_size == pipeline_model_parallel_size, ( + "pipeline parallel group already initialized, but of unexpected size: " + f"{pp_world_size=} vs. " + f"{pipeline_model_parallel_size=}" + ) + + +def model_parallel_is_initialized(): + """Check if tensor and pipeline parallel groups are initialized.""" + return _TP is not None and _PP 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_world_size(): + """Return world size for the world group.""" + return get_world_group().world_size + + +def get_world_rank(): + """Return my rank for the world group.""" + return get_world_group().rank_in_group + + +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 + + +# ATTN_TP +def get_attn_tensor_model_parallel_world_size(): + """Return world size for the attention tensor model parallel group.""" + return get_attn_tp_group().world_size + + +def get_attn_tensor_model_parallel_rank(): + """Return my rank for the attention tensor model parallel group.""" + return get_attn_tp_group().rank_in_group + + +# ATTN_CP +def get_attn_context_model_parallel_world_size(): + """Return world size for the attention context model parallel group.""" + return get_attn_cp_group().world_size + + +def get_attn_context_model_parallel_rank(): + """Return my rank for the attention context model parallel group.""" + return get_attn_cp_group().rank_in_group + + +def get_pipeline_model_parallel_world_size(): + """Return world size for the pipeline model parallel group.""" + return get_pp_group().world_size + + +def get_pipeline_model_parallel_rank(): + """Return my rank for the pipeline model parallel group.""" + return get_pp_group().rank_in_group + + +# MOE_DP +def get_moe_data_parallel_world_size(): + """Return world size for the moe data parallel group.""" + return get_moe_dp_group().world_size + + +def get_moe_data_parallel_rank(): + """Return my rank for the moe data parallel group.""" + return get_moe_dp_group().rank_in_group + + +# MOE_EP +def get_moe_expert_parallel_world_size(): + """Return world size for the moe expert parallel group.""" + return get_moe_ep_group().world_size + + +def get_moe_expert_parallel_rank(): + """Return my rank for the moe expert parallel group.""" + return get_moe_ep_group().rank_in_group + + +# MOE_TP +def get_moe_tensor_parallel_world_size(): + """Return world size for the moe tensor parallel group.""" + return get_moe_tp_group().world_size + + +def get_moe_tensor_parallel_rank(): + """Return my rank for the moe tensor parallel group.""" + return get_moe_tp_group().rank_in_group + + +def destroy_model_parallel(): + """Set the groups to none and destroy them.""" + global _TP + if _TP: + _TP.destroy() + _TP = None + + global _PP + if _PP: + _PP.destroy() + _PP = None + + global _MOE_EP + if _MOE_EP: + _MOE_EP.destroy() + _MOE_EP = None + + global _MOE_TP + if _MOE_TP: + _MOE_TP.destroy() + _MOE_TP = None + + global _ATTN_CP + if _ATTN_CP: + _ATTN_CP.destroy() + _ATTN_CP = None + + global _ATTN_TP + if _ATTN_TP: + _ATTN_TP.destroy() + _ATTN_TP = None + + global _MOE_DP + if _MOE_DP: + _MOE_DP.destroy() + _MOE_DP = None + + global _PDMUX_PREFILL_TP_GROUP + if _PDMUX_PREFILL_TP_GROUP: # type: ignore[union-attr] + _PDMUX_PREFILL_TP_GROUP.destroy() + _PDMUX_PREFILL_TP_GROUP = None + + +def destroy_distributed_environment(): + 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() + gc.collect() + if not _is_cpu: + if hasattr(torch, "cuda") and torch.cuda.is_available(): + torch.cuda.empty_cache() + if hasattr(torch._C, "_host_emptyCache"): + torch._C._host_emptyCache() + else: + logger.warning( + "torch._C._host_emptyCache() only available in Pytorch >=2.5" + ) + elif hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.empty_cache() + elif hasattr(torch, "npu") and torch.npu.is_available(): + torch.npu.empty_cache() + elif hasattr(torch, "musa") and torch.musa.is_available(): + torch.musa.empty_cache() + + +def in_the_same_node_as(pg: ProcessGroup, source_rank: int = 0) -> List[bool]: + """ + 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). + """ + 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) + + # local tensor in each process to store the result + is_in_the_same_node = torch.tensor([0] * world_size, dtype=torch.int32) + + # global ranks of the processes in the group + ranks = torch.distributed.get_process_group_ranks(pg) + + 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 + torch.distributed.broadcast_object_list( + [shm.name], src=ranks[source_rank], group=pg + ) + is_in_the_same_node[rank] = 1 + else: + # try to open the shared memory segment + recv = [None] + torch.distributed.broadcast_object_list( + recv, src=ranks[source_rank], group=pg + ) + name = recv[0] + # 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() + + torch.distributed.barrier(group=pg) + + # clean up the shared memory segment + with contextlib.suppress(OSError): + if rank == source_rank and shm: + shm.unlink() + torch.distributed.all_reduce(is_in_the_same_node, group=pg) + + return [x == 1 for x in is_in_the_same_node.tolist()] + + +vllm_get_pp_group = None +vllm_get_tp_group = None +vllm_get_world_group = None + + +def monkey_patch_vllm_parallel_state(reverse: bool = False): + try: + import vllm.distributed.parallel_state as vllm_parrlel_state + except ImportError: + return + + global vllm_get_pp_group, vllm_get_tp_group, vllm_get_world_group + if vllm_get_pp_group is None: + vllm_get_pp_group = vllm_parrlel_state.get_pp_group + vllm_get_tp_group = vllm_parrlel_state.get_tp_group + vllm_get_world_group = vllm_parrlel_state.get_world_group + if reverse: + setattr(vllm_parrlel_state, "get_pp_group", vllm_get_pp_group) + setattr(vllm_parrlel_state, "get_tp_group", vllm_get_tp_group) + setattr(vllm_parrlel_state, "get_world_group", vllm_get_world_group) + else: + setattr(vllm_parrlel_state, "get_pp_group", get_pp_group) + setattr(vllm_parrlel_state, "get_tp_group", get_tp_group) + setattr(vllm_parrlel_state, "get_world_group", get_world_group) diff --git a/platforms/patches/kunlun_p800/sitecustomize_xpu.py b/platforms/patches/kunlun_p800/sitecustomize_xpu.py new file mode 100644 index 0000000..b8d8756 --- /dev/null +++ b/platforms/patches/kunlun_p800/sitecustomize_xpu.py @@ -0,0 +1,7 @@ +import torch +# DeepSeek-V4-Flash safetensors contain F8_E8M0FN dtype metadata, but +# PyTorch 2.5.1 does not define this dtype. Alias it to uint8 so that +# safetensors deserialization can load the scales; sglang reinterprets +# them on P800 via its own FP8 kernels. +if not hasattr(torch, "float8_e8m0fnu"): + torch.float8_e8m0fnu = torch.uint8 diff --git a/platforms/patches/kunlun_p800/vocab_parallel_embedding.py b/platforms/patches/kunlun_p800/vocab_parallel_embedding.py new file mode 100644 index 0000000..2a89968 --- /dev/null +++ b/platforms/patches/kunlun_p800/vocab_parallel_embedding.py @@ -0,0 +1,680 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.3.post1/vllm/model_executor/layers/vocab_parallel_embedding.py + +import logging +from dataclasses import dataclass +from typing import List, Optional, Sequence, Tuple + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter, UninitializedParameter + +from sglang.srt.distributed import ( + divide, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, +) +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.layers.amx_utils import PackWeightMethod +from sglang.srt.layers.communicator import get_attn_tp_context +from sglang.srt.layers.dp_attention import ( + attn_tp_all_reduce, + get_attention_tp_rank, + get_attention_tp_size, + is_allocation_symmetric, +) +from sglang.srt.layers.parameter import BasevLLMParameter +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, + method_has_implemented_embedding, +) +from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod +from sglang.srt.utils import ( + cpu_has_amx_support, + is_cpu, + is_npu, + set_weight_attrs, +) + +DEFAULT_VOCAB_PADDING_SIZE = 64 + +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() +_is_npu = is_npu() + +logger = logging.getLogger(__name__) + + +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 + + +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: Optional[torch.dtype] = None, + org_num_embeddings: Optional[int] = None, + padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + enable_tp: bool = True, + use_attn_tp_group: bool = False, + use_presharded_weights: bool = False, + enable_over_encoding: bool = False, + ): + super().__init__() + self.quant_config = quant_config + self.enable_over_encoding = enable_over_encoding + + self.enable_tp = enable_tp + self.use_attn_tp_group = use_attn_tp_group + if self.enable_tp: + if use_attn_tp_group: + tp_rank = get_attention_tp_rank() + self.tp_size = get_attention_tp_size() + else: + tp_rank = get_tensor_model_parallel_rank() + self.tp_size = get_tensor_model_parallel_world_size() + else: + assert use_attn_tp_group is False + tp_rank = 0 + self.tp_size = 1 + + self.num_embeddings = num_embeddings + self.org_vocab_size = org_num_embeddings or num_embeddings + + # Support the case where the vocab size is not divisible by the TP size. + if ( + _is_cpu + and pad_vocab_size(self.org_vocab_size, padding_size) % self.tp_size != 0 + ): + padding_size *= self.tp_size + self.padding_size = padding_size + + num_added_embeddings = num_embeddings - self.org_vocab_size + self.use_presharded_weights = use_presharded_weights + if use_presharded_weights: + assert ( + num_added_embeddings == 0 + ), "Lora is not supported with presharded weights." + + 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, + host_tensor=( + True + if ( + self.__class__ is VocabParallelEmbedding + and self.enable_over_encoding + ) + else False + ), + ) + + @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) -> Optional[List[int]]: + """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] = shape[output_dim] // self.tp_size + 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.packed_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 + // (self.tp_size if self.use_presharded_weights else 1) + ), f"{self.org_vocab_size=} {self.use_presharded_weights=} {loaded_weight.shape[output_dim]=}" + + # Copy the data. + if not self.use_presharded_weights: + 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. + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): + output_parallel = self.quant_method.embedding(self, masked_input.long()) + + if self.tp_size > 1: + # Mask the output embedding. + output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) + if not get_attn_tp_context().input_scattered: + if self.use_attn_tp_group: + output_parallel = attn_tp_all_reduce(output_parallel) + else: + # Reduce across all the model parallel GPUs. + output_parallel = tensor_model_parallel_all_reduce(output_parallel) + return output_parallel + + 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}" + if self.enable_tp: + s += f", tp_size={self.tp_size}" + return s + + +class KLXVocabParallelEmbedding(torch.nn.Module): + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + params_dtype: Optional[torch.dtype] = None, + org_num_embeddings: Optional[int] = None, + padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + enable_tp: bool = True, + use_attn_tp_group: bool = False, + use_presharded_weights: bool = False, + ): + super().__init__() + self.quant_config = quant_config + + self.enable_tp = enable_tp + self.use_attn_tp_group = use_attn_tp_group + if self.enable_tp: + if use_attn_tp_group: + self.tp_rank = get_attention_tp_rank() + self.tp_size = get_attention_tp_size() + else: + self.tp_rank = get_tensor_model_parallel_rank() + self.tp_size = get_tensor_model_parallel_world_size() + else: + assert use_attn_tp_group is False + self.tp_rank = 0 + self.tp_size = 1 + # vocal_size, embedding_dim -> vocal_size, embedding_dim // tp_size + # P800 workaround: keep the embedding weight replicated on every TP rank + # to avoid the unstable all_gather across the BKCL/XCCL stack. + self.local_dim = embedding_dim + + 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.use_presharded_weights = use_presharded_weights + if use_presharded_weights: + assert ( + num_added_embeddings == 0 + ), "Lora is not supported with presharded weights." + + self.embedding_dim = embedding_dim + + """Create weights for embedding layer.""" + weight = Parameter( + torch.empty( + (self.num_embeddings, self.local_dim), + dtype=params_dtype, + ), + requires_grad=False, + ) + set_weight_attrs( + weight, + {"input_dim": 1, "output_dim": 0, "weight_loader": self.weight_loader}, + ) + self.register_parameter("weight", weight) + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): + # With the replicated-weight workaround, every rank loads the full weight. + param.copy_(loaded_weight) + + def forward(self, input_): + hidden_states = F.embedding(input_, self.weight) + # No all_gather needed because the full embedding weight is replicated. + return hidden_states + + def extra_repr(self) -> str: + s = f"num_embeddings={self.num_embeddings}" + s += f", embedding_dim={self.embedding_dim}" + s += f", org_vocab_size={self.org_vocab_size}" + if self.enable_tp: + s += f", tp_size={self.tp_size}" + return s + + +class ParallelLMHead(VocabParallelEmbedding): + """Parallelized LM head. + + Output logits weight matrices used in the Sampler. The weight and bias + tensors are padded to make sure they are divisible by the number of + model parallel GPUs. + + Args: + num_embeddings: vocabulary size. + embedding_dim: size of hidden state. + bias: whether to use bias. + params_dtype: type of the parameters. + org_num_embeddings: original vocabulary size (without LoRA). + padding_size: padding size for the vocabulary. + """ + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + *, + bias: bool = False, + params_dtype: Optional[torch.dtype] = None, + org_num_embeddings: Optional[int] = None, + padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + use_attn_tp_group: bool = False, + use_presharded_weights: bool = False, + enable_over_encoding: bool = False, + ): + super().__init__( + num_embeddings, + embedding_dim, + params_dtype=params_dtype, + org_num_embeddings=org_num_embeddings, + padding_size=padding_size, + quant_config=quant_config, + prefix=prefix, + use_attn_tp_group=use_attn_tp_group, + use_presharded_weights=use_presharded_weights, + enable_over_encoding=enable_over_encoding, + ) + self.quant_config = quant_config + + # We only support pack LMHead if it's not quantized. + if _is_cpu and _is_cpu_amx_available: + if hasattr(self, "weight") and self.weight.dtype in [ + torch.bfloat16, + torch.float16, + ]: + self.quant_method = PackWeightMethod(weight_names=["weight"]) + + if bias: + self.bias = Parameter( + torch.empty(self.num_embeddings_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 tie_weights(self, embed_tokens: VocabParallelEmbedding): + """Tie the weights with word embeddings.""" + # GGUF quantized embed_tokens. + if self.quant_config and self.quant_config.get_name() == "gguf": + return embed_tokens + else: + self.weight = embed_tokens.weight + return self + + def forward(self, input_): + del input_ + raise RuntimeError("LMHead's weights should be used in the sampler.") diff --git a/scripts/analysis/compare_experiments.py b/scripts/analysis/compare_experiments.py new file mode 100644 index 0000000..0c4d4ee --- /dev/null +++ b/scripts/analysis/compare_experiments.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""Load results.json from multiple experiments and produce a comparison report. + +Usage: + python3 scripts/analysis/compare_experiments.py + python3 scripts/analysis/compare_experiments.py --output comparison.md +""" + +import argparse +import json +from collections import defaultdict +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def load_results(): + """Glob experiments/*/results/*/results.json and load them.""" + results = [] + for path in sorted(REPO_ROOT.glob("experiments/*/results/*/results.json")): + try: + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + data["_source"] = str(path.relative_to(REPO_ROOT)) + results.append(data) + except (json.JSONDecodeError, OSError) as e: + print(f"WARN: failed to load {path}: {e}") + return results + + +def main(): + parser = argparse.ArgumentParser(description="Compare benchmark experiments.") + parser.add_argument("--output", "-o", default="comparison_report.md", help="Output markdown file") + args = parser.parse_args() + + results = load_results() + if not results: + print("No experiments found under experiments/*/results/*/results.json") + return + + # Group by scenario name. + by_scenario = defaultdict(list) + for data in results: + meta = data.get("metadata", {}) + for scenario in data.get("scenarios", []): + key = scenario.get("name", "unknown") + by_scenario[key].append((meta, scenario)) + + output_path = Path(args.output) + with open(output_path, "w", encoding="utf-8") as f: + f.write("# Cross-Experiment Comparison\n\n") + f.write("| Experiment | Chip | Engine | Scenario | Conc | In/Out | Req/s | OutTok/s | TTFT p99 | TPOT p99 | E2E p99 |\n") + f.write("|---|---|---|---|---:|---:|---:|---:|---:|---:|---:|\n") + + for scenario_name in sorted(by_scenario.keys()): + for meta, scenario in by_scenario[scenario_name]: + lat = scenario.get("latencies", {}) + f.write( + f"| {meta.get('experiment', '')} " + f"| {meta.get('chip', '')} " + f"| {meta.get('engine', '')} " + f"| {scenario_name} " + f"| {scenario.get('concurrency', '')} " + f"| {scenario.get('input_len', '')}/{scenario.get('output_len', '')} " + f"| {scenario.get('request_throughput') or ''} " + f"| {scenario.get('output_token_throughput') or ''} " + f"| {lat.get('ttft_ms', {}).get('p99') or ''} " + f"| {lat.get('tpot_ms', {}).get('p99') or ''} " + f"| {lat.get('e2e_ms', {}).get('p99') or ''} |\n" + ) + + print(f"Loaded {len(results)} experiment runs.") + print(f"Wrote {output_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/common/bench_client_docker.sh b/scripts/common/bench_client_docker.sh new file mode 100755 index 0000000..504412d --- /dev/null +++ b/scripts/common/bench_client_docker.sh @@ -0,0 +1,76 @@ +#!/usr/bin/env bash +# Helpers to run sglang.bench_serving inside the P800 Docker container. +# Usage: source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/../scripts/common/bench_client_docker.sh" + +set -Eeuo pipefail + +_BENCH_CLIENT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null +source "${_BENCH_CLIENT_DIR}/lib.sh" +# shellcheck source=/dev/null +source "${_BENCH_CLIENT_DIR}/platform.sh" + +# Required platform variables: CONTAINER_NAME, CONTAINER_PYTHON + +# Run sglang.bench_serving inside the running container. +# All arguments are forwarded to bench_serving. +run_bench_in_container() { + local container="${CONTAINER_NAME}" + + if ! docker inspect "$container" >/dev/null 2>&1; then + log "ERROR: container ${container} is not running" + return 1 + fi + + log "running bench_serving in container ${container}" + docker exec "$container" \ + env HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 HF_DATASETS_OFFLINE=1 \ + "${CONTAINER_PYTHON}" -m sglang.bench_serving "$@" +} + +# Convenience wrapper for a single random-dataset case. +# Args: +# $1: backend (e.g. sglang) +# $2: port +# $3: model path or served model name +# $4: output jsonl path (inside the container) +# $5: concurrency +# $6: input length +# $7: output length +# $8: num prompts (optional, default 512) +# $9: dataset path inside container (optional) +run_random_case() { + local backend="$1" + local port="$2" + local model="$3" + local output_file="$4" + local concurrency="$5" + local input_len="$6" + local output_len="$7" + local num_prompts="${8:-512}" + local dataset_path="${9:-}" + local warmup="${WARMUP:-100}" + + # Ensure the output directory exists inside the container. + docker exec "${CONTAINER_NAME}" mkdir -p "$(dirname "$output_file")" + + local args=( + --backend "$backend" + --host 127.0.0.1 + --port "$port" + --model "$model" + --dataset-name random + --random-input-len "$input_len" + --random-output-len "$output_len" + --num-prompts "$num_prompts" + --max-concurrency "$concurrency" + --warmup-requests "$warmup" + --output-file "$output_file" + --output-details + ) + if [[ -n "$dataset_path" ]]; then + args+=(--dataset-path "$dataset_path") + fi + + run_bench_in_container "${args[@]}" +} diff --git a/scripts/common/lib.sh b/scripts/common/lib.sh new file mode 100755 index 0000000..c77b775 --- /dev/null +++ b/scripts/common/lib.sh @@ -0,0 +1,153 @@ +#!/usr/bin/env bash +# Common helpers for benchmark orchestrators. +# Usage: source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/../scripts/common/lib.sh" + +set -Eeuo pipefail + +# Resolve repository root relative to this file. +_COMMON_LIB_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "${_COMMON_LIB_DIR}/../.." && pwd)" + +# ----------------------------------------------------------------------------- +# Logging +# ----------------------------------------------------------------------------- + +LOG_FILE="" + +log_init() { + local log_path="$1" + LOG_FILE="$log_path" + mkdir -p "$(dirname "$LOG_FILE")" + : > "$LOG_FILE" +} + +log() { + local msg="[$(date --iso-8601=seconds)] $*" + echo "$msg" + if [[ -n "${LOG_FILE:-}" ]]; then + echo "$msg" >> "$LOG_FILE" + fi +} + +# ----------------------------------------------------------------------------- +# Result directories +# ----------------------------------------------------------------------------- + +ensure_result_root() { + local result_root="$1" + mkdir -p "${result_root}/logs" + mkdir -p "${result_root}/raw_outputs" + echo "$result_root" +} + +# ----------------------------------------------------------------------------- +# Server health check +# ----------------------------------------------------------------------------- + +health_check() { + local host="${1:-127.0.0.1}" + local port="${2:-30000}" + local max_wait="${3:-120}" + + for ((i = 1; i <= max_wait; i++)); do + if curl --fail --silent --show-error --max-time 5 "http://${host}:${port}/health" >/dev/null 2>&1; then + return 0 + fi + sleep 1 + done + return 1 +} + +# ----------------------------------------------------------------------------- +# Git metadata +# ----------------------------------------------------------------------------- + +git_commit() { + cd "$ROOT_DIR" || return 1 + git rev-parse --short HEAD 2>/dev/null || echo "unknown" +} + +git_dirty() { + cd "$ROOT_DIR" || return 1 + if [[ -n "$(git status --porcelain 2>/dev/null)" ]]; then + echo "dirty" + else + echo "clean" + fi +} + +# ----------------------------------------------------------------------------- +# results.json metadata +# ----------------------------------------------------------------------------- + +write_metadata_json() { + local output_path="$1" + local experiment="$2" + local run_id="$3" + local model="$4" + local backend="$5" + local engine="$6" + local hardware="$7" + local accelerator="$8" + local chip="$9" + local script="${10}" + local env_path="${11:-}" + local description="${12:-}" + + mkdir -p "$(dirname "$output_path")" + + cat > "$output_path" <&2 + return 1 + fi + + # Use Python for safe JSON manipulation. + "${PYTHON:-python3}" - "$json_path" "$scenario_json" <<'PY' +import json +import sys + +json_path = sys.argv[1] +scenario_json = sys.argv[2] + +with open(json_path, "r", encoding="utf-8") as f: + data = json.load(f) + +scenario = json.loads(scenario_json) +data["scenarios"].append(scenario) + +with open(json_path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) +PY +} diff --git a/scripts/common/platform.sh b/scripts/common/platform.sh new file mode 100755 index 0000000..26823a6 --- /dev/null +++ b/scripts/common/platform.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +# Platform loader. +# Usage: source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/../scripts/common/platform.sh" + +set -Eeuo pipefail + +_PLATFORM_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "${_PLATFORM_DIR}/../.." && pwd)" + +# Auto-detect platform if not set. +if [[ -z "${PLATFORM:-}" ]]; then + if lspci 2>/dev/null | grep -qiE "XPU|Kunlun"; then + PLATFORM="kunlun_p800" + elif command -v nvidia-smi >/dev/null 2>&1; then + PLATFORM="nvidia_h200" + else + echo "ERROR: Could not auto-detect PLATFORM. Set PLATFORM env var explicitly." >&2 + exit 1 + fi + echo "Auto-detected platform: ${PLATFORM}" +fi + +PLATFORM_FILE="${ROOT_DIR}/platforms/${PLATFORM}.env" +if [[ ! -f "$PLATFORM_FILE" ]]; then + echo "ERROR: Platform config not found: ${PLATFORM_FILE}" >&2 + exit 1 +fi + +# shellcheck source=/dev/null +source "$PLATFORM_FILE" + +# Export key variables if not already set. +export CHIP="${CHIP:-$PLATFORM}" +export ACCELERATOR="${ACCELERATOR:-$PLATFORM}" +export HARDWARE="${HARDWARE:-$PLATFORM}" +export DEFAULT_PORT="${DEFAULT_PORT:-30000}" +export MODEL_ROOT="${MODEL_ROOT:-/data1/models}" diff --git a/scripts/common/server_docker.sh b/scripts/common/server_docker.sh new file mode 100755 index 0000000..7cc08f8 --- /dev/null +++ b/scripts/common/server_docker.sh @@ -0,0 +1,205 @@ +#!/usr/bin/env bash +# Docker-based SGLang server lifecycle helpers for Kunlun P800. +# Usage: source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/../scripts/common/server_docker.sh" + +set -Eeuo pipefail + +_DOCKER_COMMON_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null +source "${_DOCKER_COMMON_DIR}/lib.sh" +# shellcheck source=/dev/null +source "${_DOCKER_COMMON_DIR}/platform.sh" + +# Required platform variables: DOCKER_IMAGE, CONTAINER_NAME, DEFAULT_PORT, PATCH_ROOT + +docker_server_status() { + docker inspect -f '{{.State.Status}}' "$CONTAINER_NAME" 2>/dev/null || echo "missing" +} + +docker_server_stop() { + local container="${1:-$CONTAINER_NAME}" + log "stopping container ${container}" + docker rm -f "$container" 2>/dev/null || true + + # Clean up leftover shared memory on the host. + ipcs -m 2>/dev/null | awk '$4 == 666 {print $2}' | while read -r shmid; do + ipcrm -m "$shmid" 2>/dev/null || true + done + + # Kill lingering sglang processes just in case. + pkill -9 -f 'sglang.launch_server' 2>/dev/null || true + pkill -9 -f 'multiprocessing.spawn' 2>/dev/null || true + pkill -9 -f 'multiprocessing.resource_tracker' 2>/dev/null || true + sleep 2 +} + +# Build device args for docker run (8 XPU devices). +_docker_device_args() { + local args="" + for i in 0 1 2 3 4 5 6 7; do + args="${args} --device /dev/xpu${i}:/dev/xpu${i}" + done + args="${args} --device /dev/xpuctrl:/dev/xpuctrl" + echo "$args" +} + +# Start the SGLang server container for DeepSeek-V4 on P800. +# Args: +# $1: model path on host (e.g. /data1/models/DeepSeek-V4-Flash-INT8) +# $2: port (optional, defaults to DEFAULT_PORT) +# $3: server log path (optional) +# $4: server mode (optional): "fp8" or "w8a8_int8" (default "fp8") +docker_server_start() { + local model_path="$1" + local port="${2:-$DEFAULT_PORT}" + local server_log="${3:-${RESULT_ROOT:-/tmp}/logs/server.outer.log}" + local mode="${4:-fp8}" + local container="${CONTAINER_NAME}" + local image="${DOCKER_IMAGE}" + + mkdir -p "$(dirname "$server_log")" + + log "starting container ${container} from ${image}" + log "model path: ${model_path} -> /models" + log "port: ${port}" + log "mode: ${mode}" + log "server log: ${server_log}" + + docker_server_stop "$container" + + local device_args + device_args="$(_docker_device_args)" + + # Ensure patch files exist under PATCH_ROOT. Fall back to /tmp if the repo + # copy is not present (legacy path). + local patch_root="${PATCH_ROOT:-/tmp}" + + # Common environment variables shared by all modes. + local env_args=( + -e XPU_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + -e CUDA_DEVICE_ORDER=OAM_ID + -e SGLANG_USE_TRANSFORMERS_V5_TOKENIZER=1 + -e XMLIR_FORCE_USE_XPU_GRAPH=1 + -e SGLANG_DSV4_MODE=2604 + -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + -e SGLANG_NSA_DUAL_STREAM=true + -e SGLANG_NSA_QUANT_WQ_B_WK=false + -e SGLANG_FP8_PAGED_MQA_LOGITS_TORCH=1 + -e SGLANG_OPT_DEEPGEMM_HC_PRENORM=false + -e SGLANG_OPT_USE_TILELANG_MHC_PRE=1 + -e SGLANG_OPT_USE_TILELANG_MHC_POST=1 + -e SGLANG_CLEAN_REQUEST_WHEN_RETRACT=1 + -e SGLANG_SET_CPU_AFFINITY=1 + -e SGLANG_OPT_USE_KLX_TOPK_KERNEL=1 + -e XSGL_INTERTYPE_BFP16=1 + -e ENABLE_FAST_BFP16_ATTN=1 + -e XSGL_USE_DEEP_GEMM_BMM=1 + -e XSGL_XDNN_QUANT=1 + -e XSGL_FUSE_RMS_NORM_QUANT=1 + -e XSGL_TRANSPOSE_MATMUL_WEIGHT=1 + -e XINFER_QUANT_SDNN=1 + -e XSGL_USE_MOE_SIGMOID_GROUP_TOPK_NORM=1 + -e XSGL_EARLY_FIRST_TOKEN=1 + -e USE_FAST_BFP16_MOE=1 + -e XSGL_ENABLE_TGEMM_FP16=1 + -e SGLANG_ENABLE_SPEC_V2=True + -e SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 + -e PYTHONDONTWRITEBYTECODE=1 + -e XTORCH_OPS_LIB_DIR=/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/xtorch_ops + -e XPU_RUNTIME_LIB_DIR=/root/miniconda/envs/python310_torch25_cuda/xcudart/lib + -e BKCL_TREE_THRESHOLD=1048576 + -e CUDA_ENABLE_P2P_NO_UVA=1 + -e NCCL_IB_GID_INDEX=3 + -e IS_DSV4=1 + -e MC_CUSTOM_TOPO_JSON=/workspace/nic_priority_matrix_test.json + ) + + # Mode-specific environment variables. + if [[ "$mode" == "w8a8_int8" ]]; then + env_args+=( + -e SGLANG_DSV4_FP4_EXPERTS=false + -e SGLANG_APPLY_CONFIG_BACKUP=auto + -e BKCL_ENABLE_XDR=1 + -e BKCL_RDMA_NICS=eth1,eth1,eth3,eth3,eth5,eth5,eth7,eth7 + -e BKCL_RDMA_VERBS=1 + -e XSGL_INT8_LM_HEAD=1 + -e SGLANG_P800_ALL_GATHER_FALLBACK=0 + ) + else + env_args+=( + -e SGLANG_DSV4_FP4_EXPERTS=true + -e SGLANG_OPT_FUSE_WQA_WKV=false + -e SGLANG_DSV4_2604_SUBMODE=2604B + -e SGLANG_APPLY_CONFIG_BACKUP=small + -e BKCL_ENABLE_XDR=0 + -e BKCL_FORCE_PCIE_P2P=1 + -e BKCL_ENABLE_PCIE_P2P=1 + -e BKCL_SOCKET_IFNAME=bond0 + -e BKCL_NET_ENABLE_INTRA=1 + -e XSHMEM_MODE=0 + -e SGLANG_P800_ALL_GATHER_FALLBACK=1 + ) + fi + + # Mode-specific launch arguments. + local launch_args + if [[ "$mode" == "w8a8_int8" ]]; then + launch_args="--host 0.0.0.0 --port ${port} --model-path /models --attention-backend nsa --nsa-prefill klxdsa --nsa-decode klxdsa --trust-remote-code --disable-custom-all-reduce --chunked-prefill-size 8192 --page-size 64 --mem-fraction-static 0.8 --max-prefill-tokens 16384 --max-running-requests 64 --tensor-parallel-size 8 --ep-size 8 --disable-shared-experts-fusion --quantization w8a8_int8 --kv-cache-dtype float16 --disable-piecewise-cuda-graph --cuda-graph-max-bs 32 --watchdog-timeout 3000000 --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --constrained-json-disable-any-whitespace --enable-metrics --enable-request-time-stats-logging" + else + launch_args="--host 0.0.0.0 --port ${port} --model-path /models --quantization fp8 --attention-backend nsa --nsa-prefill klxdsa --nsa-decode klxdsa --trust-remote-code --chunked-prefill-size 8192 --page-size 64 --mem-fraction-static 0.85 --max-prefill-tokens 16384 --context-length 65536 --max-running-requests 8 --max-total-tokens 524288 --tensor-parallel-size 8 --ep-size 8 --moe-runner-backend deep_gemm --disable-shared-experts-fusion --kv-cache-dtype float16 --disable-piecewise-cuda-graph --disable-cuda-graph --watchdog-timeout 3000000 --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 --enable-metrics --enable-request-time-stats-logging" + fi + + # Build the full server bootstrap command and base64-encode it to avoid + # host-shell quoting hell. + local server_cmd + server_cmd=$(cat </dev/null || true +/root/miniconda/envs/python310_torch25_cuda/bin/pip install --upgrade safetensors -q +/root/miniconda/envs/python310_torch25_cuda/bin/pip install https://files.pythonhosted.org/packages/14/8b/2a1333a6455c6fad401c2285dee6f58016c55b1cb44cae3a31f8a9cc7d83/apache_tvm_ffi-0.1.0b2-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl -q +/root/miniconda/envs/python310_torch25_cuda/bin/python -c "import torch; torch.float8_e8m0fnu = torch.uint8; import runpy, sys; sys.argv[0] = 'sglang.launch_server'; runpy.run_module('sglang.launch_server', run_name='__main__')" ${launch_args} +EOF +) + local server_cmd_b64 + server_cmd_b64=$(printf '%s' "$server_cmd" | base64 -w0) + + # Patch mounts. For w8a8_int8 the working container relies on patches baked + # into the image, so we only mount the runtime config, nic matrix, and a + # local dummy dataset for offline benchmarking. + local patch_mounts=( + -v "${patch_root}/nic_priority_matrix_test.json:/workspace/nic_priority_matrix_test.json:ro" + -v "${patch_root}/dummy_sharegpt.json:/workspace/dummy_sharegpt.json:ro" + ) + if [[ "$mode" == "fp8" ]]; then + patch_mounts+=( + -v "${patch_root}/sitecustomize_xpu.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sitecustomize.py:ro" + -v "${patch_root}/hf_transformers_utils.py.patched:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/utils/hf_transformers_utils.py:ro" + -v "${patch_root}/fp8_utils.py.patched:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/layers/quantization/fp8_utils.py:ro" + -v "${patch_root}/config_backup_small_fp8.json:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/configs/config_backup_small.json:ro" + -v "${patch_root}/parallel_state.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/distributed/parallel_state.py:ro" + -v "${patch_root}/vocab_parallel_embedding.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/layers/vocab_parallel_embedding.py:ro" + ) + fi + + docker run -d \ + --name "${container}" \ + --privileged \ + --network host \ + --ipc host \ + ${device_args} \ + -v "${model_path}:/models:ro" \ + "${patch_mounts[@]}" \ + "${env_args[@]}" \ + "${image}" \ + bash -c "echo '${server_cmd_b64}' | base64 -d | bash" \ + >> "${server_log}" 2>&1 + + log "container ${container} started, waiting for health" + if health_check 127.0.0.1 "$port" 600; then + log "container ${container} is healthy" + else + log "ERROR: container ${container} failed health check" + return 1 + fi +}