diff --git a/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py
index 81a71cb96..de702db73 100644
--- a/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py
+++ b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py
@@ -541,7 +541,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
"fused_flashmla_metadata",
),
rationale_hint=(
- "NSA replay metadata copies are already fused into one-kernel" " families."
+ "NSA replay metadata copies are already fused into one-kernel families."
),
min_share=0.02,
likely_share=0.2,
@@ -787,7 +787,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
("softmax", "sampling"),
),
rationale_hint=(
- "Decode-time sampling already has fused temperature and softmax" " kernels."
+ "Decode-time sampling already has fused temperature and softmax kernels."
),
min_share=0.05,
likely_share=0.5,
@@ -1218,8 +1218,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
FusionPatternSpec(
pattern="vLLM fused residual add + RMSNorm",
candidate_path=(
- "vllm/_custom_ops.py"
- "
vllm/compilation/passes/fusion/rms_quant_fusion.py"
+ "vllm/_custom_ops.py
vllm/compilation/passes/fusion/rms_quant_fusion.py"
),
active_keywords=(
"fused_add_rms_norm",
@@ -1236,8 +1235,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
FusionPatternSpec(
pattern="vLLM fused activation-and-mul",
candidate_path=(
- "vllm/_custom_ops.py"
- "
vllm/compilation/passes/fusion/act_quant_fusion.py"
+ "vllm/_custom_ops.py
vllm/compilation/passes/fusion/act_quant_fusion.py"
),
active_keywords=(
"silu_and_mul",
diff --git a/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_proof_generator.py b/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_proof_generator.py
index 598ac7a82..0aeb0ea5f 100644
--- a/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_proof_generator.py
+++ b/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_proof_generator.py
@@ -256,7 +256,9 @@ def _module_assign_names(text: str) -> set:
targets = (
node.targets
if isinstance(node, ast.Assign)
- else [node.target] if isinstance(node, ast.AnnAssign) else []
+ else [node.target]
+ if isinstance(node, ast.AnnAssign)
+ else []
)
names |= {t.id for t in targets if isinstance(t, ast.Name)}
return names
diff --git a/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_reproduction_utils.py b/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_reproduction_utils.py
index fb0284dba..8c2856a7e 100644
--- a/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_reproduction_utils.py
+++ b/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_reproduction_utils.py
@@ -173,9 +173,9 @@ def _find_unique_def(
if isinstance(node, definition) and node.name == name
]
assert matches, f"{name} not found in {where}"
- assert (
- len(matches) == 1
- ), f"{len(matches)} defs named {name} in {where}; pass from_class to disambiguate"
+ assert len(matches) == 1, (
+ f"{len(matches)} defs named {name} in {where}; pass from_class to disambiguate"
+ )
return matches[0]
@@ -287,9 +287,9 @@ def _lowered_call_text(text: str, node: ast.Call) -> str:
"""
receiver = node.args[0]
receiver_src = _node_slice(text, receiver)
- assert (
- "\n" not in receiver_src and "#" not in receiver_src
- ), f"receiver {receiver_src!r} must be single-line and comment-free"
+ assert "\n" not in receiver_src and "#" not in receiver_src, (
+ f"receiver {receiver_src!r} must be single-line and comment-free"
+ )
opener = _slice_span(
text,
node.func.end_lineno,
@@ -722,9 +722,9 @@ class Repro:
)
existing = [alias_text(a.name, a.asname) for a in node.names]
added = alias_text(name, asname)
- assert (
- added not in existing
- ), f"{name!r} already imported from {module!r} in {rel}"
+ assert added not in existing, (
+ f"{name!r} already imported from {module!r} in {rel}"
+ )
rebuilt = f"from {module} import " + ", ".join(existing + [added]) + nl
lines[node.lineno - 1 : node.end_lineno] = [rebuilt]
_write_source(path, "".join(lines))
@@ -825,9 +825,9 @@ class Repro:
for node in tree.body
if isinstance(node, (ast.Import, ast.ImportFrom))
]
- assert (
- imports
- ), f"no imports to anchor a new `if TYPE_CHECKING:` block in {rel}"
+ assert imports, (
+ f"no imports to anchor a new `if TYPE_CHECKING:` block in {rel}"
+ )
insert_at = imports[-1].end_lineno
lines[insert_at:insert_at] = [
nl,
@@ -864,9 +864,9 @@ class Repro:
replaced = lines[node.lineno - 1].replace(
f"from {spelled} import", f"from {new_module} import", 1
)
- assert (
- replaced != lines[node.lineno - 1]
- ), f"import spelling {spelled!r} not found on its line in {rel}"
+ assert replaced != lines[node.lineno - 1], (
+ f"import spelling {spelled!r} not found on its line in {rel}"
+ )
lines[node.lineno - 1] = replaced
changed = True
assert changed, f"nested import of {name} from {old_module} not in {rel}"
@@ -974,9 +974,9 @@ class Repro:
``self: Target`` annotation is dropped (redundant inside the class). The body is moved
verbatim; the formatter normalises the surrounding blank lines.
"""
- assert (
- before is None or after is None
- ), "move_symbol: before and after are mutually exclusive"
+ assert before is None or after is None, (
+ "move_symbol: before and after are mutually exclusive"
+ )
def op(root: Path) -> None:
src_path = root / src
@@ -1255,15 +1255,17 @@ class Repro:
targets = (
node.targets
if isinstance(node, ast.Assign)
- else [node.target] if isinstance(node, ast.AnnAssign) else []
+ else [node.target]
+ if isinstance(node, ast.AnnAssign)
+ else []
)
names = {t.id for t in targets if isinstance(t, ast.Name)}
hit = names & dropped
if not hit:
continue
- assert len(names) == len(
- targets
- ), f"drop_assigns {sorted(hit)}: non-name targets in {src}"
+ assert len(names) == len(targets), (
+ f"drop_assigns {sorted(hit)}: non-name targets in {src}"
+ )
value_src = ast.unparse(node.value) if node.value is not None else None
for dropped_name in hit:
removed_assigns[dropped_name] = value_src
@@ -1289,15 +1291,17 @@ class Repro:
else:
assign_spans.append((node.lineno, node.end_lineno))
found_assigns |= hit
- assert (
- found_assigns == dropped
- ), f"{dropped - found_assigns} not assigned in {src}"
+ assert found_assigns == dropped, (
+ f"{dropped - found_assigns} not assigned in {src}"
+ )
rederivable: dict[str, str | None] = {}
for node in tree.body:
targets = (
node.targets
if isinstance(node, ast.Assign)
- else [node.target] if isinstance(node, ast.AnnAssign) else []
+ else [node.target]
+ if isinstance(node, ast.AnnAssign)
+ else []
)
names = [t.id for t in targets if isinstance(t, ast.Name)]
if not names or set(names) & dropped:
@@ -1383,9 +1387,9 @@ class Repro:
src_text = _read_source(src_path)
assert src_text.count(body) == 1, f"block not found uniquely in {src}"
at = src_text.find(body)
- assert (
- at == 0 or src_text[at - 1] == "\n"
- ), f"block matches mid-line in {src}; it must start at a line boundary"
+ assert at == 0 or src_text[at - 1] == "\n", (
+ f"block matches mid-line in {src}; it must start at a line boundary"
+ )
_write_source(src_path, src_text.replace(body, call, 1))
dst_path = root / dst
diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/proof_generator/test_infer_moves.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/proof_generator/test_infer_moves.py
index 0b9e5ead4..8e5e4b601 100644
--- a/.claude/skills/mechanical-refactor-verify/scripts/tests/proof_generator/test_infer_moves.py
+++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/proof_generator/test_infer_moves.py
@@ -370,12 +370,7 @@ def test_infer_recipe_module_level_def_shadowed_by_method_name(repo: Path) -> No
" return foo(x=self.x)\n"
),
"util.py": (
- "def keep():\n"
- " return 1\n"
- "\n"
- "\n"
- "def foo(*, x):\n"
- " return x + 1\n"
+ "def keep():\n return 1\n\n\ndef foo(*, x):\n return x + 1\n"
),
},
)
@@ -437,9 +432,7 @@ def test_infer_recipe_move_leaving_a_forwarding_delegate(repo: Path) -> None:
_write(
repo,
**{
- "model.py": (
- "class M:\n" " def work(self, x):\n" " return x + 1\n"
- ),
+ "model.py": ("class M:\n def work(self, x):\n return x + 1\n"),
"comp.py": "class C:\n def keep(self):\n return 1\n",
},
)
@@ -448,9 +441,7 @@ def test_infer_recipe_move_leaving_a_forwarding_delegate(repo: Path) -> None:
repo,
**{
"model.py": (
- "class M:\n"
- " def work(self, x):\n"
- " return self.comp.work(x)\n"
+ "class M:\n def work(self, x):\n return self.comp.work(x)\n"
),
"comp.py": (
"class C:\n"
diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_cli/cli_testlib.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_cli/cli_testlib.py
index fb1eef19b..e9490fdfa 100644
--- a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_cli/cli_testlib.py
+++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_cli/cli_testlib.py
@@ -2,13 +2,9 @@ import subprocess
from pathlib import Path
_PASSING_PROOF = (
- "import sys\n"
- 'print("PASS: reproduces the commit byte-for-byte.")\n'
- "sys.exit(0)\n"
-)
-_FAILING_PROOF = (
- "import sys\n" 'print("RESIDUAL (2 lines):\\n+x\\n-y")\n' "sys.exit(1)\n"
+ 'import sys\nprint("PASS: reproduces the commit byte-for-byte.")\nsys.exit(0)\n'
)
+_FAILING_PROOF = 'import sys\nprint("RESIDUAL (2 lines):\\n+x\\n-y")\nsys.exit(1)\n'
def _git(repo: Path, *args: str) -> str:
diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_add_imports.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_add_imports.py
index 0cf0bd58e..3b63e3ad4 100644
--- a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_add_imports.py
+++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_add_imports.py
@@ -110,13 +110,7 @@ def test_add_typechecking_import_inserts_in_block(tmp_path: Path) -> None:
def test_add_typechecking_import_creates_missing_block(tmp_path: Path) -> None:
"""With no TYPE_CHECKING block, one is created after the trailing module import."""
(tmp_path / "m.py").write_text(
- "from typing import TYPE_CHECKING\n"
- "\n"
- "from a import X\n"
- "\n"
- "\n"
- "def f():\n"
- " pass\n"
+ "from typing import TYPE_CHECKING\n\nfrom a import X\n\n\ndef f():\n pass\n"
)
r = Repro("b", "t").add_typechecking_import("m.py", "from b import Y")
_apply(r, tmp_path)
@@ -215,12 +209,7 @@ def test_add_typechecking_import_raises_without_imports(tmp_path: Path) -> None:
def test_add_typechecking_import_drops_a_lone_pass_placeholder(tmp_path: Path) -> None:
"""Populating a `pass`-only TYPE_CHECKING block replaces the placeholder."""
(tmp_path / "m.py").write_text(
- "from typing import TYPE_CHECKING\n"
- "\n"
- "if TYPE_CHECKING:\n"
- " pass\n"
- "\n"
- "x = 1\n"
+ "from typing import TYPE_CHECKING\n\nif TYPE_CHECKING:\n pass\n\nx = 1\n"
)
r = Repro("b", "t").add_typechecking_import("m.py", "from b import Y")
_apply(r, tmp_path)
diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_extract_symbols_to_new_module.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_extract_symbols_to_new_module.py
index 26d9fae2f..7e0b1ebbf 100644
--- a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_extract_symbols_to_new_module.py
+++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_extract_symbols_to_new_module.py
@@ -111,11 +111,7 @@ def test_extract_symbols_to_new_module_drops_relocated_assigns(tmp_path: Path) -
" return _FLAG\n"
)
header = (
- "from __future__ import annotations\n"
- "\n"
- "import os\n"
- "\n"
- "_FLAG = os.cpu_count()\n"
+ "from __future__ import annotations\n\nimport os\n\n_FLAG = os.cpu_count()\n"
)
r = Repro("b", "t").extract_symbols_to_new_module(
"src.py",
diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_assign.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_assign.py
index 762eab97f..1ba4e192e 100644
--- a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_assign.py
+++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_assign.py
@@ -19,13 +19,7 @@ def test_move_assign_relocates_a_module_constant(tmp_path: Path) -> None:
_apply(r, tmp_path)
assert "LIMIT" not in (tmp_path / "src.py").read_text().split("def stay")[0]
assert (tmp_path / "dst.py").read_text() == (
- "import sys\n"
- "\n"
- "LIMIT = 480 # seconds\n"
- "\n"
- "\n"
- "def keep():\n"
- " return 1\n"
+ "import sys\n\nLIMIT = 480 # seconds\n\n\ndef keep():\n return 1\n"
)
@@ -50,13 +44,7 @@ def test_move_assign_relocates_an_annotated_constant(tmp_path: Path) -> None:
_apply(r, tmp_path)
assert "LIMIT" not in (tmp_path / "src.py").read_text().split("def stay")[0]
assert (tmp_path / "dst.py").read_text() == (
- "import sys\n"
- "\n"
- "LIMIT: int = 480\n"
- "\n"
- "\n"
- "def keep():\n"
- " return 1\n"
+ "import sys\n\nLIMIT: int = 480\n\n\ndef keep():\n return 1\n"
)
diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_symbol.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_symbol.py
index a5f74a687..fc808a998 100644
--- a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_symbol.py
+++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_symbol.py
@@ -260,13 +260,7 @@ def test_move_symbol_dedent_leaves_string_literal_interior_lines(
)
_apply(r, tmp_path)
assert (tmp_path / "dst.py").read_text() == (
- "import os\n"
- "\n"
- "def foo(self):\n"
- " s = '''raw\n"
- " partial\n"
- "'''\n"
- " return s\n"
+ "import os\n\ndef foo(self):\n s = '''raw\n partial\n'''\n return s\n"
)
diff --git a/.claude/skills/sglang-prod-incident-triage/scripts/incident_artifact_tool.py b/.claude/skills/sglang-prod-incident-triage/scripts/incident_artifact_tool.py
index 2a9dacc69..fdc4aa875 100755
--- a/.claude/skills/sglang-prod-incident-triage/scripts/incident_artifact_tool.py
+++ b/.claude/skills/sglang-prod-incident-triage/scripts/incident_artifact_tool.py
@@ -101,8 +101,7 @@ def format_summary_line(filename: str, result: Dict[str, Any]) -> str:
if result.get("ok"):
return f"{filename}: ok"
return (
- f"{filename}: failed status={result.get('status')} "
- f"error={result.get('error')}"
+ f"{filename}: failed status={result.get('status')} error={result.get('error')}"
)
@@ -617,7 +616,9 @@ def summarize_dump_file(path: Path, max_requests: int, preview_chars: int) -> st
time_span = (
max(timestamps) - min(timestamps)
if len(timestamps) >= 2
- else 0.0 if len(timestamps) == 1 else None
+ else 0.0
+ if len(timestamps) == 1
+ else None
)
lines = [
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index cc6d6c1ed..32ac18680 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -48,10 +48,7 @@ repos:
python/sglang/srt/grpc/.*_pb2\.pyi$|
python/sglang/srt/grpc/.*_pb2_grpc\.pyi$|
)$
- - repo: https://github.com/psf/black
- rev: 26.1.0
- hooks:
- - id: black-jupyter
+ - id: ruff-format
exclude: '^python/sglang/srt/grpc/.*_pb2\.py$|^python/sglang/srt/grpc/.*_pb2_grpc\.py$|^python/sglang/srt/grpc/.*_pb2\.pyi$|^python/sglang/srt/grpc/.*_pb2_grpc\.pyi$'
- repo: https://github.com/codespell-project/codespell
rev: v2.4.1
diff --git a/3rdparty/amd/tuning/benchmark_moe_rocm.py b/3rdparty/amd/tuning/benchmark_moe_rocm.py
index d7ea67c5f..71b4d1710 100644
--- a/3rdparty/amd/tuning/benchmark_moe_rocm.py
+++ b/3rdparty/amd/tuning/benchmark_moe_rocm.py
@@ -187,8 +187,10 @@ def run_grid(bs, model, method, tp_size, dtype: str):
configs = union_of_list_of_dicts(prune_configs_1, prune_configs_2)
- print(f"{bs=} || {len(full_configs)=} | {len(prune_configs_1)=} | \
- {len(prune_configs_2)=} | {len(configs)=}")
+ print(
+ f"{bs=} || {len(full_configs)=} | {len(prune_configs_1)=} | \
+ {len(prune_configs_2)=} | {len(configs)=}"
+ )
best_config = None
best_time_us = 1e20
diff --git a/benchmark/asr/bench_sglang.py b/benchmark/asr/bench_sglang.py
index 875ed952b..c60685d62 100644
--- a/benchmark/asr/bench_sglang.py
+++ b/benchmark/asr/bench_sglang.py
@@ -343,7 +343,7 @@ def run_evaluation(args):
print("\n" + "=" * 20 + " Sample Predictions " + "=" * 20)
num_to_show = min(args.print_n, len(results))
for i in range(num_to_show):
- print(f"Sample {i+1}:")
+ print(f"Sample {i + 1}:")
print(f" REF: {references[i]}")
print(f" PRED: {predictions[i]}")
print("-" * 40)
diff --git a/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py b/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py
index 5e988392e..903841d57 100644
--- a/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py
+++ b/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py
@@ -243,9 +243,7 @@ def run(task, fi, tri, device, dtype, args):
) # noqa: E731
else:
inp = make_decode_inputs(B, H, HV, K, V, pool, device, dtype)
- corr = lambda kern: call_decode(
- kern, inp, inp["ssm"].clone()
- ) # noqa: E731
+ corr = lambda kern: call_decode(kern, inp, inp["ssm"].clone()) # noqa: E731
ssm_t = inp["ssm"].clone()
timed = lambda kern: call_decode(kern, inp, ssm_t) # noqa: E731
diff --git a/benchmark/gsm8k/bench_sglang.py b/benchmark/gsm8k/bench_sglang.py
index a0e09a39c..b3989673a 100644
--- a/benchmark/gsm8k/bench_sglang.py
+++ b/benchmark/gsm8k/bench_sglang.py
@@ -53,9 +53,9 @@ def main(args):
if args.enable_thinking:
from transformers import AutoTokenizer
- assert (
- args.tokenizer_path is not None
- ), "--tokenizer-path is required when --enable-thinking is set"
+ assert args.tokenizer_path is not None, (
+ "--tokenizer-path is required when --enable-thinking is set"
+ )
tokenizer = AutoTokenizer.from_pretrained(
args.tokenizer_path, trust_remote_code=True
)
diff --git a/benchmark/hf3fs/bench_client.py b/benchmark/hf3fs/bench_client.py
index ef4a967c4..7b18a8b2c 100644
--- a/benchmark/hf3fs/bench_client.py
+++ b/benchmark/hf3fs/bench_client.py
@@ -14,11 +14,11 @@ def print_stats(x: List[int]):
x = sorted(x)
lenx = len(x)
print(
- f"mean = {sum(x)/len(x):.2f}, "
+ f"mean = {sum(x) / len(x):.2f}, "
f"min = {min(x):.2f}, "
- f"p25 = {x[int(lenx*0.25)]:.2f}, "
- f"p50 = {x[int(lenx*0.5)]:.2f}, "
- f"p75 = {x[int(lenx*0.75)]:.2f}, "
+ f"p25 = {x[int(lenx * 0.25)]:.2f}, "
+ f"p50 = {x[int(lenx * 0.5)]:.2f}, "
+ f"p75 = {x[int(lenx * 0.75)]:.2f}, "
f"max = {max(x):.2f}"
)
diff --git a/benchmark/hf3fs/bench_storage.py b/benchmark/hf3fs/bench_storage.py
index f0ce171bf..71122a41d 100644
--- a/benchmark/hf3fs/bench_storage.py
+++ b/benchmark/hf3fs/bench_storage.py
@@ -18,11 +18,11 @@ def print_stats(x: List[int]):
x = sorted(x)
lenx = len(x)
print(
- f"mean = {sum(x)/len(x):.2f}, "
+ f"mean = {sum(x) / len(x):.2f}, "
f"min = {min(x):.2f}, "
- f"p25 = {x[int(lenx*0.25)]:.2f}, "
- f"p50 = {x[int(lenx*0.5)]:.2f}, "
- f"p75 = {x[int(lenx*0.75)]:.2f}, "
+ f"p25 = {x[int(lenx * 0.25)]:.2f}, "
+ f"p50 = {x[int(lenx * 0.5)]:.2f}, "
+ f"p75 = {x[int(lenx * 0.75)]:.2f}, "
f"max = {max(x):.2f}"
)
diff --git a/benchmark/hf3fs/bench_zerocopy.py b/benchmark/hf3fs/bench_zerocopy.py
index 0f9c5e150..a4ed66912 100644
--- a/benchmark/hf3fs/bench_zerocopy.py
+++ b/benchmark/hf3fs/bench_zerocopy.py
@@ -109,7 +109,7 @@ elif hicache_mem_layout == "layer_first":
for operation in operations:
cache_controller.generic_page_backup(operation, batch_size=128)
tok = time.monotonic()
-print(f"{tok-tik:.6f} s")
+print(f"{tok - tik:.6f} s")
operations = [
PrefetchOperation(
@@ -137,4 +137,4 @@ elif hicache_mem_layout == "layer_first":
for operation in operations:
cache_controller.generic_page_transfer(operation, batch_size=128)
tok = time.monotonic()
-print(f"{tok-tik:.6f} s")
+print(f"{tok - tik:.6f} s")
diff --git a/benchmark/hicache/bench_mix.py b/benchmark/hicache/bench_mix.py
index 2a65574ea..5d2615186 100644
--- a/benchmark/hicache/bench_mix.py
+++ b/benchmark/hicache/bench_mix.py
@@ -457,7 +457,7 @@ class WorkloadGenerator:
try:
user_data, response = self.response_queue.get(timeout=10)
logger.info(
- f"{((time.perf_counter()-self.start_time)/self.duration*100):.2f}%"
+ f"{((time.perf_counter() - self.start_time) / self.duration * 100):.2f}%"
)
if not response.success:
raise ValueError(f"Request failed with error: {response.error}")
@@ -540,10 +540,10 @@ class WorkloadGenerator:
output_stats = self.user_generator.output_stats
print(f"round_ratios: {user_stats}")
print(
- f"mean_new_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in input_stats]}"
+ f"mean_new_tokens_per_round: {[int(a / b) if b > 0 else 0 for a, b in input_stats]}"
)
print(
- f"mean_return_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in output_stats]}"
+ f"mean_return_tokens_per_round: {[int(a / b) if b > 0 else 0 for a, b in output_stats]}"
)
return performance_data
diff --git a/benchmark/hicache/bench_serving.py b/benchmark/hicache/bench_serving.py
index a80c059e5..4e3285858 100644
--- a/benchmark/hicache/bench_serving.py
+++ b/benchmark/hicache/bench_serving.py
@@ -75,9 +75,9 @@ async def async_request_openai_completions(
pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
api_url = request_func_input.api_url
- assert api_url.endswith(
- "completions"
- ), "OpenAI Completions API URL must end with 'completions'."
+ assert api_url.endswith("completions"), (
+ "OpenAI Completions API URL must end with 'completions'."
+ )
async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
payload = {
diff --git a/benchmark/hicache/nextqa.py b/benchmark/hicache/nextqa.py
index 4db6caa1c..1c74dac71 100644
--- a/benchmark/hicache/nextqa.py
+++ b/benchmark/hicache/nextqa.py
@@ -120,7 +120,7 @@ class NExTQALoader(VideoLoader):
video = Video(video_path, num_frames)
prompt = entry["question"] + "?"
if self.task == "MC": # add choices
- prompt += f' a0: {entry["a0"]}, a1: {entry["a1"]}, a2: {entry["a2"]}, a3: {entry["a3"]}'
+ prompt += f" a0: {entry['a0']}, a1: {entry['a1']}, a2: {entry['a2']}, a3: {entry['a3']}"
return VideoPrompt(video_path, num_frames, prompt)
def __iter__(self):
diff --git a/benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py b/benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py
index 53e0b8c75..f3230134f 100644
--- a/benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py
+++ b/benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py
@@ -149,9 +149,9 @@ def _check_correctness():
cos = torch.nn.functional.cosine_similarity(
(q.float() * scale).flatten(), ref_deq.flatten(), dim=0
).item()
- assert (
- cos > 0.99
- ), f"{provider} h={hidden_size} residual={add_residual} cos={cos:.4f}"
+ assert cos > 0.99, (
+ f"{provider} h={hidden_size} residual={add_residual} cos={cos:.4f}"
+ )
print("correctness check passed (all fused providers vs unfused within FP8)")
diff --git a/benchmark/kernels/deepep/deepep_utils.py b/benchmark/kernels/deepep/deepep_utils.py
index 169529ef6..573d68deb 100644
--- a/benchmark/kernels/deepep/deepep_utils.py
+++ b/benchmark/kernels/deepep/deepep_utils.py
@@ -191,9 +191,9 @@ def bench_kineto(
kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names
assert all([isinstance(name, str) for name in kernel_names])
for name in kernel_names:
- assert (
- sum([name in line for line in prof_lines]) == 1
- ), f"Errors of the kernel {name} in the profiling table"
+ assert sum([name in line for line in prof_lines]) == 1, (
+ f"Errors of the kernel {name} in the profiling table"
+ )
# Save chrome traces
if trace_path is not None:
diff --git a/benchmark/kernels/deepep/tuning_deepep.py b/benchmark/kernels/deepep/tuning_deepep.py
index 191819d2c..6a1401408 100644
--- a/benchmark/kernels/deepep/tuning_deepep.py
+++ b/benchmark/kernels/deepep/tuning_deepep.py
@@ -155,7 +155,7 @@ def test_main(
for with_topk in (False, True):
if local_rank == 0:
print(
- f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
+ f"[testing] Running with {'FP8' if isinstance(current_x, tuple) else 'BF16'}, {'with' if with_topk else 'without'} top-k (async={async_mode}, previous={previous_mode}) ...",
flush=True,
end="",
)
@@ -198,9 +198,9 @@ def test_main(
# Checks
recv_gbl_rank_prefix_sum = handle[-4]
- assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
- 0
- ), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
+ assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(0), (
+ f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
+ )
assert (
gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
== recv_num_tokens_per_expert_list
@@ -325,11 +325,14 @@ def test_main(
tune_args = {"x": current_x, "handle": handle, "config": config}
t = bench(lambda: buffer.dispatch(**tune_args))[0]
if t < best_time:
- best_time, best_results = t, (
- num_sms,
- nvl_chunk_size,
- rdma_chunk_size,
- config_kwargs,
+ best_time, best_results = (
+ t,
+ (
+ num_sms,
+ nvl_chunk_size,
+ rdma_chunk_size,
+ config_kwargs,
+ ),
)
if local_rank == 0:
print(
@@ -338,7 +341,7 @@ def test_main(
)
if local_rank == 0:
print(
- f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
+ f"[tuning] Best dispatch ({'FP8' if isinstance(current_x, tuple) else 'BF16'}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
flush=True,
)
print("", flush=True)
@@ -399,11 +402,14 @@ def test_main(
flush=True,
)
if t < best_time:
- best_time, best_results = t, (
- num_sms,
- nvl_chunk_size,
- rdma_chunk_size,
- config_kwargs,
+ best_time, best_results = (
+ t,
+ (
+ num_sms,
+ nvl_chunk_size,
+ rdma_chunk_size,
+ config_kwargs,
+ ),
)
if local_rank == 0:
diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py
index 0b958de68..bae6904e2 100644
--- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py
+++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py
@@ -59,7 +59,6 @@ def tl_gemm(
bx,
by,
):
-
A_shared = T.alloc_shared(A_shared_shape, in_dtype)
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
C_shared = T.alloc_shared(C_shared_shape, out_dtype)
@@ -350,7 +349,7 @@ def get_benchmark(tp_size):
tflops = flops / (ms * 1e-3) / 1e12
# Print shape-specific results with TFLOPS
- print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}")
+ print(f"Time: {ms * 1000:.2f} ms, TFLOPS: {tflops:.2f}")
return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms
return benchmark
diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py
index 5670f3c33..280f6b97d 100644
--- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py
+++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py
@@ -224,7 +224,7 @@ def _benchmark(m, n, k, tp_size, provider):
tflops = flops / (ms * 1e-3) / 1e12
# Print shape-specific results with TFLOPS
- print(f"Time: {ms*1000:.2f} us, TFLOPS: {tflops:.2f}")
+ print(f"Time: {ms * 1000:.2f} us, TFLOPS: {tflops:.2f}")
return ms, max_ms, min_ms
diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py
index 8b1be7b88..1aace8087 100644
--- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py
+++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py
@@ -435,7 +435,7 @@ def get_benchmark(tp_size):
flops = 2 * m * n * k # multiply-adds
tflops = flops / (ms * 1e-3) / 1e12
- print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}")
+ print(f"Time: {ms * 1000:.2f} ms, TFLOPS: {tflops:.2f}")
return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms
return benchmark
diff --git a/benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py b/benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py
index 2f885b0ee..12a73faf8 100755
--- a/benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py
+++ b/benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py
@@ -243,8 +243,7 @@ def main():
else:
speedup = f"{legacy_us / us:.2f}x"
print(
- f"{s_q:>5} {kv_len:>7} {name:>10} {us:>9.1f} "
- f"{tbps:>9.3f} {speedup:>8}"
+ f"{s_q:>5} {kv_len:>7} {name:>10} {us:>9.1f} {tbps:>9.3f} {speedup:>8}"
)
print()
diff --git a/benchmark/kernels/elementwise/benchmark_concat_mla.py b/benchmark/kernels/elementwise/benchmark_concat_mla.py
index 7bc51d3da..8ae687aad 100644
--- a/benchmark/kernels/elementwise/benchmark_concat_mla.py
+++ b/benchmark/kernels/elementwise/benchmark_concat_mla.py
@@ -143,9 +143,7 @@ output_exp = execute_and_get_output(fn_cuda, data)
if not torch.all(output_ref == output_exp):
abs_delta = torch.abs(output_ref - output_exp)
raise AssertionError(
- f"{output_ref=} {output_exp=} "
- f"{abs_delta=} "
- f"{torch.argwhere(abs_delta != 0.0)=} "
+ f"{output_ref=} {output_exp=} {abs_delta=} {torch.argwhere(abs_delta != 0.0)=} "
)
diff --git a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py
index a74972c33..1c5880147 100644
--- a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py
+++ b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py
@@ -535,7 +535,6 @@ class BestConfigTrace:
class BenchmarkWorker:
-
def __init__(self, seed: int, server_args: ServerArgs) -> None:
torch.set_default_device("cuda")
torch.cuda.manual_seed_all(0)
@@ -729,8 +728,7 @@ class BenchmarkWorker:
down_use_tma_map[block_m] = time_cost_all[2] > time_cost_all[3]
print(
- f"Round 1 done. Down TMA decisions per BLOCK_SIZE_M: "
- f"{down_use_tma_map}"
+ f"Round 1 done. Down TMA decisions per BLOCK_SIZE_M: {down_use_tma_map}"
)
# === Round 2: Up with c_sorted from round 1 ===
diff --git a/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py b/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py
index aa684dfa8..84a2d8f5a 100755
--- a/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py
+++ b/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py
@@ -470,9 +470,9 @@ def _tune_shrink(
device: torch.device,
) -> tuple:
"""Tune shrink kernel for one layer type. Returns (best_configs, results)."""
- print(f"\n{'='*80}")
+ print(f"\n{'=' * 80}")
print(f"Tuning SHRINK โ {label} (K={K}, N={N}, slices={num_slices})")
- print(f"{'='*80}")
+ print(f"{'=' * 80}")
search = get_shrink_search_space()
print(f"Search space: {len(search)} configs")
@@ -508,7 +508,7 @@ def _tune_shrink(
best_config = config
if (i + 1) % 20 == 0:
print(
- f" chunk={chunk_size}: {i+1}/{len(search)} tested, best={best_time:.3f}ms"
+ f" chunk={chunk_size}: {i + 1}/{len(search)} tested, best={best_time:.3f}ms"
)
best_configs[chunk_size] = sort_config(best_config)
@@ -533,9 +533,9 @@ def _tune_expand(
device: torch.device,
) -> tuple:
"""Tune expand kernel for one layer type. Returns (best_configs, results)."""
- print(f"\n{'='*80}")
+ print(f"\n{'=' * 80}")
print(f"Tuning EXPAND โ {label} (output_dim={output_dim}, slices={num_slices})")
- print(f"{'='*80}")
+ print(f"{'=' * 80}")
search = get_expand_search_space()
print(f"Search space: {len(search)} configs")
@@ -584,7 +584,7 @@ def _tune_expand(
best_config = config
if (i + 1) % 50 == 0:
print(
- f" chunk={chunk_size}: {i+1}/{len(search)} tested, best={best_time:.3f}ms"
+ f" chunk={chunk_size}: {i + 1}/{len(search)} tested, best={best_time:.3f}ms"
)
best_configs[chunk_size] = sort_config(best_config)
@@ -673,9 +673,9 @@ def main(args: argparse.Namespace):
)
# --- Summary ---
- print(f"\n{'='*80}")
+ print(f"\n{'=' * 80}")
print(f"SUMMARY")
- print(f"{'='*80}")
+ print(f"{'=' * 80}")
print(
f"\n{'layer':<10} {'kernel':<8} {'K/dim':>6} {'chunk':>6}"
f" {'baseline':>10} {'tuned':>10} {'speedup':>8} config"
diff --git a/benchmark/lean_kernel_sweep.py b/benchmark/lean_kernel_sweep.py
index f6a6679ac..6e5ce6a8f 100755
--- a/benchmark/lean_kernel_sweep.py
+++ b/benchmark/lean_kernel_sweep.py
@@ -137,19 +137,21 @@ def main():
# b32 x 128K on the 8-KV-head config exceeds the microbench's single
# contiguous KV tensor (faults the GPU); real serving uses a paged pool.
if H_KV == 8 and B == 32 and S == 131072:
- print(f"{B:>5} {S//1024:>5}K {'skipped (contiguous-KV limit)':>30}")
+ print(
+ f"{B:>5} {S // 1024:>5}K {'skipped (contiguous-KV limit)':>30}"
+ )
rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,skip,")
continue
try:
std, lean, cos, gate = run(H_Q, H_KV, B, S)
except torch.cuda.OutOfMemoryError:
torch.cuda.empty_cache()
- print(f"{B:>5} {S//1024:>5}K {'OOM':>9}")
+ print(f"{B:>5} {S // 1024:>5}K {'OOM':>9}")
rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,OOM,")
continue
sp = std / lean
print(
- f"{B:>5} {S//1024:>5}K {std:>9.3f} {lean:>9.3f} {sp:>7.2f}x {('ON' if gate else 'OFF'):>5} {cos:>7.4f}"
+ f"{B:>5} {S // 1024:>5}K {std:>9.3f} {lean:>9.3f} {sp:>7.2f}x {('ON' if gate else 'OFF'):>5} {cos:>7.4f}"
)
rows.append(
f"{name},{H_Q},{H_KV},{B},{S},{std:.4f},{lean:.4f},{sp:.4f},{cos:.4f},{int(gate)}"
diff --git a/benchmark/mmlu/bench_sglang.py b/benchmark/mmlu/bench_sglang.py
index 9a2006e3d..b86aec46c 100644
--- a/benchmark/mmlu/bench_sglang.py
+++ b/benchmark/mmlu/bench_sglang.py
@@ -163,7 +163,7 @@ def main(args):
pt = 0
for subject, num_qs in zip(subjects[: args.nsub], num_questions):
print(
- f"subject: {subject}, #q:{num_qs}, acc: {np.mean(cors[pt: pt + num_qs]):.3f}"
+ f"subject: {subject}, #q:{num_qs}, acc: {np.mean(cors[pt : pt + num_qs]):.3f}"
)
pt += num_qs
assert pt == len(cors)
diff --git a/benchmark/ocr/bench_sglang.py b/benchmark/ocr/bench_sglang.py
index 651f61417..f210c6232 100644
--- a/benchmark/ocr/bench_sglang.py
+++ b/benchmark/ocr/bench_sglang.py
@@ -502,7 +502,7 @@ async def process_sample(
}
)
print(
- f"[INPUT ] [{i+1}] type={ttype!r:12s} expected: {expected[:120]}",
+ f"[INPUT ] [{i + 1}] type={ttype!r:12s} expected: {expected[:120]}",
flush=True,
)
# Print OCR output (truncate long outputs)
diff --git a/benchmark/ocr/generate_report.py b/benchmark/ocr/generate_report.py
index 770c827ee..95bba28d6 100644
--- a/benchmark/ocr/generate_report.py
+++ b/benchmark/ocr/generate_report.py
@@ -159,12 +159,12 @@ def _render_sample(sample: dict, failures_only: bool) -> str:
if failures_only and passed == total and not error:
return ""
- pct = f"{100*passed//total}%" if total else "โ"
+ pct = f"{100 * passed // total}%" if total else "โ"
header_cls = "fail" if (error or passed < total) else "pass"
parts = [f'
{html.escape(ti.get("text", ""))}')
+ parts.append(f"{html.escape(ti.get('text', ''))}")
elif ttype in ("order", "natural_reading_order"):
before = ti.get("before", "")
after = ti.get("after", "")
@@ -236,7 +236,7 @@ def _render_sample(sample: dict, failures_only: bool) -> str:
parts.append(f'โฆ and {len(matches)-6} more
' + f'โฆ and {len(matches) - 6} more
' ) else: parts.append( diff --git a/benchmark/prefill_only/util.py b/benchmark/prefill_only/util.py index 2451239d6..ad21dc1b3 100644 --- a/benchmark/prefill_only/util.py +++ b/benchmark/prefill_only/util.py @@ -383,14 +383,14 @@ async def send_warmup_requests( http_url, data=request_json, headers=headers ) as resp: if resp.status == 200: - print(f"Warmup request {i+1}/{num_warmup} completed successfully") + print(f"Warmup request {i + 1}/{num_warmup} completed successfully") else: print( - f"Warmup request {i+1}/{num_warmup} failed with status {resp.status}" + f"Warmup request {i + 1}/{num_warmup} failed with status {resp.status}" ) except Exception as e: - print(f"Warmup request {i+1}/{num_warmup} failed with error: {e}") + print(f"Warmup request {i + 1}/{num_warmup} failed with error: {e}") print("HTTP warmup requests completed") @@ -745,7 +745,6 @@ async def run_generic_benchmark( async with aiohttp.ClientSession( timeout=aiohttp.ClientTimeout(total=300) ) as session: - # Send START_PROFILE if profiling is enabled if config.profile: await send_profile_request("START_PROFILE", http_url, session=session) diff --git a/benchmark/scheduler/bench_token_storage.py b/benchmark/scheduler/bench_token_storage.py index ea8ef7418..79ee6aa6d 100644 --- a/benchmark/scheduler/bench_token_storage.py +++ b/benchmark/scheduler/bench_token_storage.py @@ -254,7 +254,7 @@ def print_breakdown(title: str, results: dict[str, dict[str, float]]) -> None: def microbench_torch_tensor_paths( - sizes: tuple[int, ...] = (1_000, 10_000, 100_000) + sizes: tuple[int, ...] = (1_000, 10_000, 100_000), ) -> None: """Compare three CPU-buffer -> pinned cuda tensor paths. @@ -293,9 +293,11 @@ def microbench_torch_tensor_paths( ), ( "(C) from_numpy(frombuf(array('q'))).pin() -> cuda", - lambda x: torch.from_numpy(np.frombuffer(x, dtype=np.int64)) - .pin_memory() - .to("cuda", non_blocking=True), + lambda x: ( + torch.from_numpy(np.frombuffer(x, dtype=np.int64)) + .pin_memory() + .to("cuda", non_blocking=True) + ), ), ]: cells = [] diff --git a/docs/demo/deepseek_v4_flash.ipynb b/docs/demo/deepseek_v4_flash.ipynb index 5f2f03495..8ffde4076 100644 --- a/docs/demo/deepseek_v4_flash.ipynb +++ b/docs/demo/deepseek_v4_flash.ipynb @@ -347,11 +347,11 @@ " try:\n", " paper[\"full_text\"] = download_and_extract(paper)\n", " print(\n", - " f\"[{i+1}/{N_FULL_PAPERS}] {paper['title'][:70]} โ {len(paper['full_text']):,} chars\"\n", + " f\"[{i + 1}/{N_FULL_PAPERS}] {paper['title'][:70]} โ {len(paper['full_text']):,} chars\"\n", " )\n", " except Exception as e:\n", " paper[\"full_text\"] = None\n", - " print(f\"[{i+1}/{N_FULL_PAPERS}] {paper['title'][:70]} โ failed: {e}\")" + " print(f\"[{i + 1}/{N_FULL_PAPERS}] {paper['title'][:70]} โ failed: {e}\")" ] }, { diff --git a/examples/frontend_language/usage/llava_video/srt_example_llava_v.py b/examples/frontend_language/usage/llava_video/srt_example_llava_v.py index ec5b334b0..d042a09da 100644 --- a/examples/frontend_language/usage/llava_video/srt_example_llava_v.py +++ b/examples/frontend_language/usage/llava_video/srt_example_llava_v.py @@ -124,7 +124,6 @@ def batch(video_dir, save_dir, cur_chunk, num_chunks, num_frames=16, batch_size= if __name__ == "__main__": - url = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4" cache_dir = os.path.expanduser("~/.cache") diff --git a/examples/frontend_language/usage/readme_examples.py b/examples/frontend_language/usage/readme_examples.py index 7269ef148..0b66e00fc 100644 --- a/examples/frontend_language/usage/readme_examples.py +++ b/examples/frontend_language/usage/readme_examples.py @@ -31,7 +31,7 @@ def tip_suggestion(s): forks = s.fork(2) for i, f in enumerate(forks): - f += f"Now, expand tip {i+1} into a paragraph:\n" + f += f"Now, expand tip {i + 1} into a paragraph:\n" f += sgl.gen(f"detailed_tip", max_tokens=256, stop="\n\n") s += "Tip 1:" + forks[0]["detailed_tip"] + "\n" diff --git a/examples/profiler/nsys_profile_tools/gputrc2graph.py b/examples/profiler/nsys_profile_tools/gputrc2graph.py index 4bfc4340e..68cde5045 100755 --- a/examples/profiler/nsys_profile_tools/gputrc2graph.py +++ b/examples/profiler/nsys_profile_tools/gputrc2graph.py @@ -86,7 +86,7 @@ class GPUTrace2Graph: # Update current_end for overlapping intervals for i in range(1, len(df)): if i % display_units == 0: - print(f"processing trace: {int(i/len(df) * 100)} %", end="\r") + print(f"processing trace: {int(i / len(df) * 100)} %", end="\r") if starts[i] <= current_end: if ends[i] > current_end: # Partial overlap @@ -182,9 +182,9 @@ class GPUTrace2Graph: def is_valid_file(self, base_file): """asserts if base_file is non-existent or is empty""" - assert ( - os.path.isfile(base_file) and os.path.getsize(base_file) > 0 - ), f"{base_file} doesn't exist or is empty" + assert os.path.isfile(base_file) and os.path.getsize(base_file) > 0, ( + f"{base_file} doesn't exist or is empty" + ) def should_gen_file(self, new_file, base_file): """figure out if new file should be generated from base_file""" diff --git a/examples/runtime/engine/fastapi_engine_inference.py b/examples/runtime/engine/fastapi_engine_inference.py index f5da9d715..66eb8ce12 100644 --- a/examples/runtime/engine/fastapi_engine_inference.py +++ b/examples/runtime/engine/fastapi_engine_inference.py @@ -130,7 +130,7 @@ def send_requests(server_url, prompts, max_new_tokens, temperature): """Sends generation requests to the running server for a list of prompts.""" # Iterate through prompts and send requests for i, prompt in enumerate(prompts): - print(f"\n[{i+1}/{len(prompts)}] Sending prompt: '{prompt}'") + print(f"\n[{i + 1}/{len(prompts)}] Sending prompt: '{prompt}'") payload = { "prompt": prompt, "max_new_tokens": max_new_tokens, diff --git a/examples/runtime/engine/offline_batch_inference_qwen_1m.py b/examples/runtime/engine/offline_batch_inference_qwen_1m.py index 664efa6d7..5505bf7b4 100644 --- a/examples/runtime/engine/offline_batch_inference_qwen_1m.py +++ b/examples/runtime/engine/offline_batch_inference_qwen_1m.py @@ -17,8 +17,7 @@ def load_prompt() -> str: # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/1m.txt with urlopen( - "https://qianwen-res.oss-cn-beijing.aliyuncs.com" - "/Qwen2.5-1M/test-data/64k.txt", + "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/64k.txt", timeout=5, ) as response: prompt = response.read().decode("utf-8") @@ -41,9 +40,7 @@ def process_requests(llm: sgl.Engine, prompts: list[str]) -> None: for output in outputs: prompt_token_ids = output["meta_info"]["prompt_tokens"] generated_text = output["text"] - print( - f"Prompt length: {prompt_token_ids}, " f"Generated text: {generated_text!r}" - ) + print(f"Prompt length: {prompt_token_ids}, Generated text: {generated_text!r}") # Create an LLM. diff --git a/examples/runtime/qwen3_vl_reranker.py b/examples/runtime/qwen3_vl_reranker.py index 09779996f..b94e27631 100644 --- a/examples/runtime/qwen3_vl_reranker.py +++ b/examples/runtime/qwen3_vl_reranker.py @@ -44,7 +44,7 @@ def rerank_text_only(): print("Results (sorted by relevance):") for i, result in enumerate(results): - print(f" {i+1}. Score: {result['score']:.4f} - {result['document'][:60]}...") + print(f" {i + 1}. Score: {result['score']:.4f} - {result['document'][:60]}...") print() @@ -99,7 +99,7 @@ def rerank_with_images(): print("Results (sorted by relevance):") for i, result in enumerate(results): - print(f" {i+1}. Index: {result['index']}, Score: {result['score']:.4f}") + print(f" {i + 1}. Index: {result['index']}, Score: {result['score']:.4f}") print() @@ -149,7 +149,7 @@ def rerank_multimodal_query(): print("Results (sorted by relevance):") for i, result in enumerate(results): - print(f" {i+1}. Index: {result['index']}, Score: {result['score']:.4f}") + print(f" {i + 1}. Index: {result['index']}, Score: {result['score']:.4f}") print() diff --git a/examples/usage/modelopt_quantize_and_export.py b/examples/usage/modelopt_quantize_and_export.py index 4394d917c..b3f26ede4 100755 --- a/examples/usage/modelopt_quantize_and_export.py +++ b/examples/usage/modelopt_quantize_and_export.py @@ -213,7 +213,7 @@ def deploy_exported_model( outputs = llm.generate(prompts, sampling_params) for i, output in enumerate(outputs): - print(f"Prompt {i+1}: {prompts[i]}") + print(f"Prompt {i + 1}: {prompts[i]}") print(f"Output: {output['text']}") print() diff --git a/examples/usage/reasoning_aware_compression/rac_collect_traces.py b/examples/usage/reasoning_aware_compression/rac_collect_traces.py index a9be8f69e..4c83b1618 100755 --- a/examples/usage/reasoning_aware_compression/rac_collect_traces.py +++ b/examples/usage/reasoning_aware_compression/rac_collect_traces.py @@ -335,7 +335,7 @@ def report(manifest: TraceManifest) -> None: print(f" decode tokens : {manifest.num_decode_tokens}") print(f" total tokens : {total}") print(f" decode share : {decode_share:.1%}") - print(f" wall clock : {manifest.elapsed_seconds/60:.1f} min") + print(f" wall clock : {manifest.elapsed_seconds / 60:.1f} min") if manifest.calibration_mode == "rac": print( "\nThe decode share is the activation mass that prompt-only " diff --git a/examples/usage/reasoning_aware_compression/rac_serve_and_eval.py b/examples/usage/reasoning_aware_compression/rac_serve_and_eval.py index 7448cf4db..27a70ab0a 100755 --- a/examples/usage/reasoning_aware_compression/rac_serve_and_eval.py +++ b/examples/usage/reasoning_aware_compression/rac_serve_and_eval.py @@ -209,7 +209,7 @@ def report(results: List[EvalResult]) -> None: f"{result.model_path:<{width}} " f"{result.accuracy:>7.3f} " f"{result.mean_completion_tokens:>16.0f} " - f"{result.elapsed_seconds/60:>10.1f}m" + f"{result.elapsed_seconds / 60:>10.1f}m" ) if len(results) > 1: diff --git a/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py b/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py index 5534b428d..e92bbdb22 100644 --- a/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py +++ b/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py @@ -194,9 +194,9 @@ def _direct_warm(worker_url: str, model_id: str, prefix: str) -> None: }, timeout=60.0, ) - assert ( - r.status_code == 200 - ), f"direct warm to {worker_url} failed: HTTP {r.status_code} {r.text!r}" + assert r.status_code == 200, ( + f"direct warm to {worker_url} failed: HTTP {r.status_code} {r.text!r}" + ) def _route_through(router_url: str, model_id: str, prompt: str) -> str: @@ -212,9 +212,9 @@ def _route_through(router_url: str, model_id: str, prompt: str) -> str: after = _success_counts_by_worker(router_url) deltas = {w: after.get(w, 0) - before.get(w, 0) for w in set(after) | set(before)} winners = [w for w, d in deltas.items() if d > 0] - assert ( - len(winners) == 1 - ), f"expected exactly one worker delta on {router_url}, got {deltas}" + assert len(winners) == 1, ( + f"expected exactly one worker delta on {router_url}, got {deltas}" + ) return winners[0] @@ -309,15 +309,15 @@ def test_routers_route_by_prefix_content( landed = _route_through( router.base_url, spec["model"], PREFIX_X ) - assert ( - landed == worker_x.url - ), f"router {label}: PREFIX_X must route to {worker_x.url}; landed on {landed}" + assert landed == worker_x.url, ( + f"router {label}: PREFIX_X must route to {worker_x.url}; landed on {landed}" + ) landed = _route_through( router.base_url, spec["model"], PREFIX_Y ) - assert ( - landed == worker_y.url - ), f"router {label}: PREFIX_Y must route to {worker_y.url}; landed on {landed}" + assert landed == worker_y.url, ( + f"router {label}: PREFIX_Y must route to {worker_y.url}; landed on {landed}" + ) except Exception: _dump_logs(logs) raise diff --git a/experimental/sgl-router/tests/e2e/chat_completions/test_validation.py b/experimental/sgl-router/tests/e2e/chat_completions/test_validation.py index 4b9ec369a..5216831fd 100644 --- a/experimental/sgl-router/tests/e2e/chat_completions/test_validation.py +++ b/experimental/sgl-router/tests/e2e/chat_completions/test_validation.py @@ -50,9 +50,9 @@ def test_chat_non_streaming_returns_assistant_message( body = resp.json() choice = body["choices"][0] assert choice["message"]["role"] == "assistant" - assert choice["message"][ - "content" - ], f"empty assistant content: {choice!r}" + assert choice["message"]["content"], ( + f"empty assistant content: {choice!r}" + ) assert choice.get("finish_reason"), choice finally: gpu_allocator.release(gpu) @@ -91,8 +91,8 @@ def test_chat_streaming_emits_sse_chunks_with_done( if line.startswith("data:"): chunks.append(line.strip()) assert len(chunks) >= 2, f"expected >=2 SSE chunks, got: {chunks}" - assert any( - "[DONE]" in c for c in chunks - ), f"no [DONE] terminator in stream: {chunks}" + assert any("[DONE]" in c for c in chunks), ( + f"no [DONE] terminator in stream: {chunks}" + ) finally: gpu_allocator.release(gpu) diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/test_discovery.py b/experimental/sgl-router/tests/e2e/k8s_integration/test_discovery.py index 83f0faab0..400d5cd41 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/test_discovery.py +++ b/experimental/sgl-router/tests/e2e/k8s_integration/test_discovery.py @@ -66,15 +66,17 @@ def test_router_discovers_multiple_workers(router_url): # Scale down to 1 โ router should still route after reconverging _scale_fake_worker(1) _poll_until( - lambda: httpx.post( - f"{router_url}/v1/chat/completions", - json={ - "model": "tiny", - "messages": [{"role": "user", "content": "post-scale-down"}], - }, - timeout=10.0, - ).status_code - == 200, + lambda: ( + httpx.post( + f"{router_url}/v1/chat/completions", + json={ + "model": "tiny", + "messages": [{"role": "user", "content": "post-scale-down"}], + }, + timeout=10.0, + ).status_code + == 200 + ), "router routes after scale-down to 1", timeout=60, interval=3, diff --git a/experimental/sgl-router/tests/e2e/test_chat_smoke.py b/experimental/sgl-router/tests/e2e/test_chat_smoke.py index 16d030bd3..82847e1f9 100644 --- a/experimental/sgl-router/tests/e2e/test_chat_smoke.py +++ b/experimental/sgl-router/tests/e2e/test_chat_smoke.py @@ -16,9 +16,9 @@ def test_models(router: str) -> None: assert resp.status_code == 200, resp.text data = resp.json() ids = [m["id"] for m in data.get("data", [])] - assert any( - MODEL in mid for mid in ids - ), f"Model {MODEL!r} not found in /v1/models response: {ids}" + assert any(MODEL in mid for mid in ids), ( + f"Model {MODEL!r} not found in /v1/models response: {ids}" + ) def test_chat_non_streaming(router: str) -> None: @@ -59,6 +59,6 @@ def test_chat_streaming(router: str) -> None: chunks.append(line) assert len(chunks) >= 2, f"Expected >=2 SSE chunks, got {len(chunks)}: {chunks}" - assert any( - "[DONE]" in c for c in chunks - ), f"No [DONE] chunk found in SSE stream: {chunks}" + assert any("[DONE]" in c for c in chunks), ( + f"No [DONE] chunk found in SSE stream: {chunks}" + ) diff --git a/experimental/sgl-router/tests/e2e/test_tokenize_smoke.py b/experimental/sgl-router/tests/e2e/test_tokenize_smoke.py index cee049f87..1e9babf6e 100644 --- a/experimental/sgl-router/tests/e2e/test_tokenize_smoke.py +++ b/experimental/sgl-router/tests/e2e/test_tokenize_smoke.py @@ -20,9 +20,9 @@ def test_tokenize_round_trip(router: str) -> None: ) assert tok_resp.status_code == 200, tok_resp.text tokens = tok_resp.json()["tokens"] - assert ( - isinstance(tokens, list) and len(tokens) > 0 - ), f"Expected non-empty token list, got: {tokens}" + assert isinstance(tokens, list) and len(tokens) > 0, ( + f"Expected non-empty token list, got: {tokens}" + ) # Detokenize detok_resp = httpx.post( @@ -32,6 +32,6 @@ def test_tokenize_round_trip(router: str) -> None: ) assert detok_resp.status_code == 200, detok_resp.text recovered = detok_resp.json()["text"] - assert ( - TEXT in recovered or recovered in TEXT - ), f"Round-trip mismatch: original={TEXT!r}, recovered={recovered!r}" + assert TEXT in recovered or recovered in TEXT, ( + f"Round-trip mismatch: original={TEXT!r}, recovered={recovered!r}" + ) diff --git a/experimental/sgl-router/tests/scripts/generate_parity_fixtures.py b/experimental/sgl-router/tests/scripts/generate_parity_fixtures.py index 5d33813f2..4e985823f 100644 --- a/experimental/sgl-router/tests/scripts/generate_parity_fixtures.py +++ b/experimental/sgl-router/tests/scripts/generate_parity_fixtures.py @@ -85,7 +85,7 @@ def load_tokenizer_with_fallback(primary, fallback, slug): raise continue raise RuntimeError( - f"No accessible tokenizer for slug={slug} " f"(tried: {primary}, {fallback})" + f"No accessible tokenizer for slug={slug} (tried: {primary}, {fallback})" ) diff --git a/python/sglang/benchmark/dspark_sps_profiler.py b/python/sglang/benchmark/dspark_sps_profiler.py index 1f914fa7c..88a3a13ca 100644 --- a/python/sglang/benchmark/dspark_sps_profiler.py +++ b/python/sglang/benchmark/dspark_sps_profiler.py @@ -687,8 +687,7 @@ def run_one_round( rank_rows = fetch_rank_rows(base_url=context.base_url) if len(rank_rows) != len(watermarks): raise RuntimeError( - f"DP rank count changed mid-profile: {len(watermarks)} -> " - f"{len(rank_rows)}." + f"DP rank count changed mid-profile: {len(watermarks)} -> {len(rank_rows)}." ) new_rank_rows = [ [row for row in rows if row.forward_ct > watermark] diff --git a/python/sglang/benchmark/offline_throughput.py b/python/sglang/benchmark/offline_throughput.py index 37acaad6f..0114c8a15 100644 --- a/python/sglang/benchmark/offline_throughput.py +++ b/python/sglang/benchmark/offline_throughput.py @@ -136,13 +136,13 @@ class BenchArgs: "--gsp-system-prompt-len", type=int, default=BenchArgs.gsp_system_prompt_len, - help="System prompt length, used" "only for generate-shared-prefix", + help="System prompt length, usedonly for generate-shared-prefix", ) parser.add_argument( "--gsp-question-len", type=int, default=BenchArgs.gsp_question_len, - help="Question length, used" "only for generate-shared-prefix", + help="Question length, usedonly for generate-shared-prefix", ) parser.add_argument( "--gsp-output-len", @@ -259,9 +259,9 @@ def throughput_test_once( ] if profile: - assert ( - "SGLANG_TORCH_PROFILER_DIR" in os.environ - ), "Please set SGLANG_TORCH_PROFILER_DIR." + assert "SGLANG_TORCH_PROFILER_DIR" in os.environ, ( + "Please set SGLANG_TORCH_PROFILER_DIR." + ) os.makedirs(os.environ["SGLANG_TORCH_PROFILER_DIR"], exist_ok=True) known_files = None backend.start_profile( diff --git a/python/sglang/benchmark/one_batch_server.py b/python/sglang/benchmark/one_batch_server.py index d904992c9..27a426336 100644 --- a/python/sglang/benchmark/one_batch_server.py +++ b/python/sglang/benchmark/one_batch_server.py @@ -486,7 +486,7 @@ def _warmup_cache( return print( - f"Warming up cache with {cache_hit_rate*100:.1f}% hit rate " + f"Warming up cache with {cache_hit_rate * 100:.1f}% hit rate " f"({cached_token_len} tokens per request)" ) # Create prefix input_ids for cache warming @@ -1024,7 +1024,7 @@ def get_report_summary( f"\nInput lens: {bench_args.input_len}. Output lens: {bench_args.output_len}." ) if bench_args.cache_hit_rate > 0.0: - summary += f" Cache hit rate: {bench_args.cache_hit_rate*100:.1f}%." + summary += f" Cache hit rate: {bench_args.cache_hit_rate * 100:.1f}%." summary += "\n" if is_blackwell(): @@ -1241,9 +1241,9 @@ def run_benchmark_internal( skip_max_running_requests_threshold = float("inf") skip_token_capacity_threshold = float("inf") else: - assert ( - max_running_requests_per_dp > 0 - ), f"effective_max_running_requests_per_dp is not set, {max_running_requests_per_dp=}" + assert max_running_requests_per_dp > 0, ( + f"effective_max_running_requests_per_dp is not set, {max_running_requests_per_dp=}" + ) skip_max_running_requests_threshold = max_running_requests_per_dp * dp_size print(f"{max_running_requests_per_dp=}") @@ -1288,9 +1288,9 @@ def run_benchmark_internal( "--lora-request-distribution=distinct/skewed requires more than " "one adapter via --lora-name." ) - assert ( - bench_args.lora_zipf_alpha > 1 - ), f"--lora-zipf-alpha must be > 1, got {bench_args.lora_zipf_alpha}" + assert bench_args.lora_zipf_alpha > 1, ( + f"--lora-zipf-alpha must be > 1, got {bench_args.lora_zipf_alpha}" + ) if bench_args.apply_chat_template and not ( bench_args.fixed_prompt_file or bench_args.dataset_name in REPLAY_TEXT_DATASETS diff --git a/python/sglang/benchmark/serving.py b/python/sglang/benchmark/serving.py index 9ba76b1cb..4362619d5 100644 --- a/python/sglang/benchmark/serving.py +++ b/python/sglang/benchmark/serving.py @@ -261,9 +261,9 @@ async def async_request_openai_completions( pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url - assert api_url.endswith( - "completions" - ), "OpenAI Completions API URL must end with 'completions'." + assert api_url.endswith("completions"), ( + "OpenAI Completions API URL must end with 'completions'." + ) prompt = request_func_input.prompt @@ -392,9 +392,9 @@ async def async_request_openai_chat_completions( latency, TTFT, ITL, and success status. """ api_url = request_func_input.api_url - assert api_url.endswith( - "chat/completions" - ), "OpenAI Chat Completions API URL must end with 'chat/completions'." + assert api_url.endswith("chat/completions"), ( + "OpenAI Chat Completions API URL must end with 'chat/completions'." + ) # TODO put it to other functions when `pbar` logic is refactored if getattr(args, "print_requests", False): @@ -1296,9 +1296,9 @@ def _normalize_round_messages(turn: Any) -> Optional[List[Dict[str, str]]]: def wrap_multi_turn_request_func(request_func: Callable, backend: str) -> Callable: - assert ( - backend in MULTI_TURN_BACKENDS - ), f"Multi-turn only supports chat backends: {MULTI_TURN_BACKENDS}, got {backend}" + assert backend in MULTI_TURN_BACKENDS, ( + f"Multi-turn only supports chat backends: {MULTI_TURN_BACKENDS}, got {backend}" + ) async def f( request_func_input: RequestFuncInput, @@ -1534,9 +1534,9 @@ async def benchmark( lora_name = lora_names[lora_idx] lora_idx = (lora_idx + 1) % len(lora_names) else: - assert ( - lora_request_distribution == "skewed" - ), f"Unexpected lora_request_distribution: {lora_request_distribution}. Expected 'skewed'." + assert lora_request_distribution == "skewed", ( + f"Unexpected lora_request_distribution: {lora_request_distribution}. Expected 'skewed'." + ) lora_name = np.random.choice(lora_names, p=lora_probs) else: @@ -2000,9 +2000,9 @@ def run_benchmark(args_: argparse.Namespace): extra_request_body["bootstrap_room"] = 0 if args.tokenize_prompt: - assert ( - args.backend == "sglang" - ), "`--tokenize-prompt` only compatible with `--backend sglang` currently" + assert args.backend == "sglang", ( + "`--tokenize-prompt` only compatible with `--backend sglang` currently" + ) # Set url if args.port is None: @@ -2079,18 +2079,18 @@ def run_benchmark(args_: argparse.Namespace): if args.dataset_name in ["image", "mmmu"]: args.apply_chat_template = True - assert ( - not args.tokenize_prompt - ), "`--tokenize-prompt` not compatible with image dataset" + assert not args.tokenize_prompt, ( + "`--tokenize-prompt` not compatible with image dataset" + ) if args.lora_request_distribution in ["distinct", "skewed"]: - assert ( - args.lora_name is not None and len(args.lora_name) > 1 - ), "More than 1 LoRA adapter must be specified via --lora-name to use 'distinct' or 'skewed' request distribution." + assert args.lora_name is not None and len(args.lora_name) > 1, ( + "More than 1 LoRA adapter must be specified via --lora-name to use 'distinct' or 'skewed' request distribution." + ) - assert ( - args.lora_zipf_alpha > 1 - ), f"Got invalid value for --lora-zipf-alpha of {args.lora_zipf_alpha}. It must be greater than 1." + assert args.lora_zipf_alpha > 1, ( + f"Got invalid value for --lora-zipf-alpha of {args.lora_zipf_alpha}. It must be greater than 1." + ) print(f"{args}\n") @@ -2364,13 +2364,13 @@ def cli_main(): "--image-format", type=str, default="jpeg", - help=("Format of images for image dataset. " "Supports jpeg and png."), + help=("Format of images for image dataset. Supports jpeg and png."), ) parser.add_argument( "--image-content", type=str, default="random", - help=("Content for images for image dataset. " "Supports random and blank."), + help=("Content for images for image dataset. Supports random and blank."), ) parser.add_argument( "--request-rate", diff --git a/python/sglang/cli/killall.py b/python/sglang/cli/killall.py index 1e672df2c..631988416 100755 --- a/python/sglang/cli/killall.py +++ b/python/sglang/cli/killall.py @@ -315,8 +315,7 @@ def _print_diagnostics(unkillable_pids): print(f" {line}") else: print( - "\n[killall] Diagnostic โ no sglang/python/gpu processes " - "in this container" + "\n[killall] Diagnostic โ no sglang/python/gpu processes in this container" ) diff --git a/python/sglang/cli/serve.py b/python/sglang/cli/serve.py index 3a934ada5..6e913622d 100644 --- a/python/sglang/cli/serve.py +++ b/python/sglang/cli/serve.py @@ -195,7 +195,7 @@ def serve(args, extra_argv): else: registered = registry.get(backend_name) logger.info( - "Dispatch override enabled: --model-type=%s " "(skip auto detection)", + "Dispatch override enabled: --model-type=%s (skip auto detection)", backend_name, ) diff --git a/python/sglang/cli/utils.py b/python/sglang/cli/utils.py index 26258be2b..f27ad6cf0 100644 --- a/python/sglang/cli/utils.py +++ b/python/sglang/cli/utils.py @@ -126,8 +126,7 @@ def get_model_path(extra_argv): ) else: raise Exception( - "Error: --model-path is required. " - "Please provide the path to the model." + "Error: --model-path is required. Please provide the path to the model." ) return model_path diff --git a/python/sglang/kernels/aot/benchmark/bench_amd_deterministic_allreduce.py b/python/sglang/kernels/aot/benchmark/bench_amd_deterministic_allreduce.py index 234d09342..60b658698 100644 --- a/python/sglang/kernels/aot/benchmark/bench_amd_deterministic_allreduce.py +++ b/python/sglang/kernels/aot/benchmark/bench_amd_deterministic_allreduce.py @@ -245,7 +245,7 @@ def worker(world_size, rank, port, results_queue): if input_size_bytes > custom_ar.max_size: if rank == 0: print( - f" Deterministic kernel skipped: input size ({input_size_bytes/(1024*1024):.1f} MB) > buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)" + f" Deterministic kernel skipped: input size ({input_size_bytes / (1024 * 1024):.1f} MB) > buffer size ({custom_ar.max_size / (1024 * 1024):.1f} MB)" ) deterministic_kernel_available = False else: @@ -412,17 +412,17 @@ def worker(world_size, rank, port, results_queue): } print( - f" All-Reduce: {lat_ar_median*1000:.3f}ms, Deterministic: {ar_deterministic}, Max variance: {ar_max_variance:.6f}" + f" All-Reduce: {lat_ar_median * 1000:.3f}ms, Deterministic: {ar_deterministic}, Max variance: {ar_max_variance:.6f}" ) print( - f" RS+All-Gather: {lat_rs_ag_median*1000:.3f}ms, Deterministic: {rs_ag_deterministic}, Max variance: {rs_ag_max_variance:.6f}" + f" RS+All-Gather: {lat_rs_ag_median * 1000:.3f}ms, Deterministic: {rs_ag_deterministic}, Max variance: {rs_ag_max_variance:.6f}" ) if custom_ar is not None and lat_custom_ar_median is not None: overhead_custom = ( (lat_custom_ar_median - lat_ar_median) / lat_ar_median ) * 100 print( - f" Custom AR: {lat_custom_ar_median*1000:.3f}ms, Deterministic: {custom_ar_deterministic}, Max variance: {custom_ar_max_variance:.6f}, Overhead: {overhead_custom:+.1f}%" + f" Custom AR: {lat_custom_ar_median * 1000:.3f}ms, Deterministic: {custom_ar_deterministic}, Max variance: {custom_ar_max_variance:.6f}, Overhead: {overhead_custom:+.1f}%" ) if lat_deterministic_kernel_median is not None: overhead_kernel = ( @@ -433,7 +433,7 @@ def worker(world_size, rank, port, results_queue): / lat_rs_ag_median ) * 100 print( - f" Deterministic Kernel: {lat_deterministic_kernel_median*1000:.3f}ms, Deterministic: {deterministic_kernel_deterministic}, Max variance: {deterministic_kernel_max_variance:.6f}, Overhead: {overhead_kernel:+.1f}%, Speedup vs RS+AG: {speedup_kernel_vs_rs_ag:+.1f}%" + f" Deterministic Kernel: {lat_deterministic_kernel_median * 1000:.3f}ms, Deterministic: {deterministic_kernel_deterministic}, Max variance: {deterministic_kernel_max_variance:.6f}, Overhead: {overhead_kernel:+.1f}%, Speedup vs RS+AG: {speedup_kernel_vs_rs_ag:+.1f}%" ) if lat_optimized_rs_ag_median is not None: overhead_opt = ( @@ -443,7 +443,7 @@ def worker(world_size, rank, port, results_queue): (lat_rs_ag_median - lat_optimized_rs_ag_median) / lat_rs_ag_median ) * 100 print( - f" Optimized RS+AG: {lat_optimized_rs_ag_median*1000:.3f}ms, Deterministic: {optimized_rs_ag_deterministic}, Max variance: {optimized_rs_ag_max_variance:.6f}, Overhead: {overhead_opt:+.1f}%, Speedup vs RS+AG: {speedup_vs_rs_ag:+.1f}%" + f" Optimized RS+AG: {lat_optimized_rs_ag_median * 1000:.3f}ms, Deterministic: {optimized_rs_ag_deterministic}, Max variance: {optimized_rs_ag_max_variance:.6f}, Overhead: {overhead_opt:+.1f}%, Speedup vs RS+AG: {speedup_vs_rs_ag:+.1f}%" ) print(f" RS+AG Overhead: {overhead_rs_ag:+.1f}%") @@ -515,8 +515,8 @@ def main(): ar_det_str = "โ" if r["all_reduce"]["deterministic"] else "โ" rs_ag_det_str = "โ" if r["rs_ag"]["deterministic"] else "โ" line = ( - f"{bs:<8} {r['all_reduce']['latency_median']*1000:<12.3f} {ar_det_str:<8} " - f"{r['rs_ag']['latency_median']*1000:<15.3f} {rs_ag_det_str:<10} " + f"{bs:<8} {r['all_reduce']['latency_median'] * 1000:<12.3f} {ar_det_str:<8} " + f"{r['rs_ag']['latency_median'] * 1000:<15.3f} {rs_ag_det_str:<10} " f"{r['overhead_rs_ag_pct']:<12.1f}" ) if r.get("custom_ar") is not None: @@ -526,7 +526,7 @@ def main(): (custom_ar["latency_median"] - r["all_reduce"]["latency_median"]) / r["all_reduce"]["latency_median"] ) * 100 - line += f" {custom_ar['latency_median']*1000:<18.3f} {custom_ar_det_str:<15} {custom_ar_overhead:<15.1f}" + line += f" {custom_ar['latency_median'] * 1000:<18.3f} {custom_ar_det_str:<15} {custom_ar_overhead:<15.1f}" if r.get("deterministic_kernel") is not None: det_kernel = r["deterministic_kernel"] det_kernel_det_str = "โ" if det_kernel["deterministic"] else "โ" @@ -538,7 +538,7 @@ def main(): (r["rs_ag"]["latency_median"] - det_kernel["latency_median"]) / r["rs_ag"]["latency_median"] ) * 100 - line += f" {det_kernel['latency_median']*1000:<18.3f} {det_kernel_det_str:<15} {det_kernel_overhead:<15.1f} {speedup_kernel:<10.1f}" + line += f" {det_kernel['latency_median'] * 1000:<18.3f} {det_kernel_det_str:<15} {det_kernel_overhead:<15.1f} {speedup_kernel:<10.1f}" if r.get("optimized_rs_ag") is not None: opt_rs_ag = r["optimized_rs_ag"] opt_rs_ag_det_str = "โ" if opt_rs_ag["deterministic"] else "โ" @@ -550,7 +550,7 @@ def main(): (r["rs_ag"]["latency_median"] - opt_rs_ag["latency_median"]) / r["rs_ag"]["latency_median"] ) * 100 - line += f" {opt_rs_ag['latency_median']*1000:<18.3f} {opt_rs_ag_det_str:<15} {opt_rs_ag_overhead:<15.1f} {speedup:<10.1f}" + line += f" {opt_rs_ag['latency_median'] * 1000:<18.3f} {opt_rs_ag_det_str:<15} {opt_rs_ag_overhead:<15.1f} {speedup:<10.1f}" print(line) print("=" * 80) diff --git a/python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py b/python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py index 461d8862b..9d24443f6 100644 --- a/python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py +++ b/python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py @@ -114,10 +114,8 @@ def benchmark(batch_size, seq_len, provider, block_size, num_kv_splits): q_size = qn.numel() * qn.element_size() + qr.numel() * qr.element_size() - gbps = ( - lambda ms: ( - q_size + q_size * dv / d + kv_cache.numel() * kv_cache.element_size() - ) + gbps = lambda ms: ( + (q_size + q_size * dv / d + kv_cache.numel() * kv_cache.element_size()) * 1e-9 / (ms * 1e-3) ) diff --git a/python/sglang/kernels/aot/benchmark/bench_fp4_gemm.py b/python/sglang/kernels/aot/benchmark/bench_fp4_gemm.py index df536f508..26d8aad30 100755 --- a/python/sglang/kernels/aot/benchmark/bench_fp4_gemm.py +++ b/python/sglang/kernels/aot/benchmark/bench_fp4_gemm.py @@ -368,9 +368,9 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): res_fi, backend="cudnn", ) - assert torch.allclose( - res_fi, res_cutlass, atol=1e-3, rtol=1e-3 - ), "cudnn fp4 doesn't match cutlass fp4" + assert torch.allclose(res_fi, res_cutlass, atol=1e-3, rtol=1e-3), ( + "cudnn fp4 doesn't match cutlass fp4" + ) mm_fp4( a_fp4, b_fp4_T, @@ -381,9 +381,9 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): res_fi, backend="trtllm", ) - assert torch.allclose( - res_fi, res_cutlass, atol=1e-3, rtol=1e-3 - ), "trtllm fp4 doesn't match cutlass fp4" + assert torch.allclose(res_fi, res_cutlass, atol=1e-3, rtol=1e-3), ( + "trtllm fp4 doesn't match cutlass fp4" + ) if csv_file: with open(csv_file, "a", newline="") as f: diff --git a/python/sglang/kernels/aot/benchmark/bench_int8_gemm.py b/python/sglang/kernels/aot/benchmark/bench_int8_gemm.py index 722d89d7c..0d0cb2a15 100644 --- a/python/sglang/kernels/aot/benchmark/bench_int8_gemm.py +++ b/python/sglang/kernels/aot/benchmark/bench_int8_gemm.py @@ -127,8 +127,8 @@ def benchmark(batch_size, provider, N, K): lambda: vllm_scaled_mm(a, b, scale_a, scale_b, torch.float16, bias), quantiles=quantiles, ) - gbps = ( - lambda ms: ( + gbps = lambda ms: ( + ( (2 * M * N * K - M * N) * a.element_size() + (3 * M * N) * scale_a.element_size() ) diff --git a/python/sglang/kernels/aot/python/sgl_kernel/attention.py b/python/sglang/kernels/aot/python/sgl_kernel/attention.py index faf23a4f0..3351d6d6f 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/attention.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/attention.py @@ -38,9 +38,9 @@ def cutlass_mla_decode( ) -> torch.Tensor: assert q_nope.ndim == 3, f"q_nope must be a 3D tensor, but got {q_nope.ndim}" assert q_pe.ndim == 3, f"q_pe must be a 3D tensor, but got {q_pe.ndim}" - assert ( - kv_c_and_k_pe_cache.ndim == 3 - ), f"kv_c_and_k_pe_cache must be a 3D tensor, but got {kv_c_and_k_pe_cache.ndim}" + assert kv_c_and_k_pe_cache.ndim == 3, ( + f"kv_c_and_k_pe_cache must be a 3D tensor, but got {kv_c_and_k_pe_cache.ndim}" + ) B_q, H, D_q_nope = q_nope.shape B_q_2, H_2, D_q_pe = q_pe.shape @@ -77,12 +77,12 @@ def cutlass_mla_decode( torch.bfloat16, ), f"q_nope.dtype needs to be fp16 or bf16 but got {q_nope.dtype}." assert q_nope.dtype == q_pe.dtype == kv_c_and_k_pe_cache.dtype - assert ( - seq_lens.dtype == torch.int32 - ), f"seq_lens.dtype needs to be int32 but got {seq_lens.dtype}." - assert ( - page_table.dtype == torch.int32 - ), f"page_table.dtype needs to be int32 but got {page_table.dtype}." + assert seq_lens.dtype == torch.int32, ( + f"seq_lens.dtype needs to be int32 but got {seq_lens.dtype}." + ) + assert page_table.dtype == torch.int32, ( + f"page_table.dtype needs to be int32 but got {page_table.dtype}." + ) out = q_nope.new_empty((B_q, MAX_HEADS, D_latent)) diff --git a/python/sglang/kernels/aot/python/sgl_kernel/elementwise.py b/python/sglang/kernels/aot/python/sgl_kernel/elementwise.py index aa325b277..d2315761c 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/elementwise.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/elementwise.py @@ -247,12 +247,12 @@ def gemma_fused_add_rmsnorm( def _check_shape(input: torch.Tensor, output: torch.Tensor) -> None: assert input.ndim == output.ndim, f"{input.ndim} != {output.ndim}" - assert ( - input.shape[:-1] == output.shape[:-1] - ), f"{input.shape[:-1]} != {output.shape[:-1]}" - assert ( - input.shape[-1] == 2 * output.shape[-1] - ), f"{input.shape[-1]} != {2 * output.shape[-1]}" + assert input.shape[:-1] == output.shape[:-1], ( + f"{input.shape[:-1]} != {output.shape[:-1]}" + ) + assert input.shape[-1] == 2 * output.shape[-1], ( + f"{input.shape[-1]} != {2 * output.shape[-1]}" + ) def silu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor: diff --git a/python/sglang/kernels/aot/python/sgl_kernel/flash_mla.py b/python/sglang/kernels/aot/python/sgl_kernel/flash_mla.py index bd8d558a6..2a15aad9c 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/flash_mla.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/flash_mla.py @@ -159,9 +159,9 @@ def flash_mla_with_kvcache( assert extra_topk_length is None if indices is not None: assert causal == False, "causal must be `false` if sparse attention is enabled." - assert (descale_q is None) == ( - descale_k is None - ), "descale_q and descale_k should be both None or both not None" + assert (descale_q is None) == (descale_k is None), ( + "descale_q and descale_k should be both None or both not None" + ) if indices is None and q.element_size() == 1: out, softmax_lse = torch.ops.sgl_kernel.fwd_kvcache_mla_fp8.default( @@ -257,9 +257,9 @@ def _flash_mla_with_kvcache_sched_meta( assert sched_meta.config.causal == causal, helper_msg assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, helper_msg assert sched_meta.config.topk == topk, helper_msg - assert ( - sched_meta.config.extra_page_block_size == extra_page_block_size - ), helper_msg + assert sched_meta.config.extra_page_block_size == extra_page_block_size, ( + helper_msg + ) assert sched_meta.config.extra_topk == extra_topk, helper_msg if topk is not None: diff --git a/python/sglang/kernels/aot/python/sgl_kernel/metal.py b/python/sglang/kernels/aot/python/sgl_kernel/metal.py index 8edec27c8..a98b976c8 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/metal.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/metal.py @@ -76,11 +76,11 @@ def rope_pool_fused( if q_shape != (q_shape[0], num_qo_heads, head_dim): raise ValueError( - "q shape must be [num_tokens, num_qo_heads, head_dim], " f"got {q.shape}" + f"q shape must be [num_tokens, num_qo_heads, head_dim], got {q.shape}" ) if k_shape != (q_shape[0], num_kv_heads, head_dim): raise ValueError( - "k shape must be [num_tokens, num_kv_heads, head_dim], " f"got {k.shape}" + f"k shape must be [num_tokens, num_kv_heads, head_dim], got {k.shape}" ) if v_shape != k_shape: raise ValueError(f"v shape must match k shape, got {v.shape} vs {k.shape}") diff --git a/python/sglang/kernels/aot/python/sgl_kernel/musa.py b/python/sglang/kernels/aot/python/sgl_kernel/musa.py index 49dd825ca..220bb4858 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/musa.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/musa.py @@ -86,9 +86,9 @@ def musa_fused_gemv( out_shape = x.shape[:-1] + ( qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2, ) - assert not ( - use_swigelu and use_rms_norm - ), "gemv only fused one activation (swigelu or rms_norm)!" + assert not (use_swigelu and use_rms_norm), ( + "gemv only fused one activation (swigelu or rms_norm)!" + ) if use_rms_norm: if gamma is None: @@ -113,9 +113,9 @@ def musa_fused_gemv( return output # w4a16 gemv elif qweight_scales is not None: - assert ( - x.dtype == torch.bfloat16 or x.dtype == torch.float16 - ), "W4A16 gemv only support bfloat16 or float16!" + assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, ( + "W4A16 gemv only support bfloat16 or float16!" + ) use_int4_w4a16 = True out_shape = x.shape[:-1] + ( qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2, diff --git a/python/sglang/kernels/aot/python/sgl_kernel/scalar_type.py b/python/sglang/kernels/aot/python/sgl_kernel/scalar_type.py index 571c1dca7..24dbaa93f 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/scalar_type.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/scalar_type.py @@ -70,9 +70,9 @@ class ScalarType: """ def _floating_point_max_int(self) -> int: - assert ( - self.mantissa <= 52 and self.exponent <= 11 - ), f"Cannot represent max/min as a double for type {self.__str__()}" + assert self.mantissa <= 52 and self.exponent <= 11, ( + f"Cannot represent max/min as a double for type {self.__str__()}" + ) max_mantissa = (1 << self.mantissa) - 1 if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN: @@ -80,9 +80,9 @@ class ScalarType: max_exponent = (1 << self.exponent) - 2 if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN or self.nan_repr == NanRepr.NONE: - assert ( - self.exponent < 11 - ), f"Cannot represent max/min as a double for type {self.__str__()}" + assert self.exponent < 11, ( + f"Cannot represent max/min as a double for type {self.__str__()}" + ) max_exponent = max_exponent + 1 # adjust the exponent to match that of a double @@ -109,25 +109,25 @@ class ScalarType: if self.is_floating_point(): return self._floating_point_max() else: - assert ( - self.size_bits < 64 or self.size_bits == 64 and self.is_signed() - ), "Cannot represent max as an int" + assert self.size_bits < 64 or self.size_bits == 64 and self.is_signed(), ( + "Cannot represent max as an int" + ) return (1 << self.mantissa) - 1 def _raw_min(self) -> Union[int, float]: if self.is_floating_point(): - assert ( - self.is_signed() - ), "We currently assume all floating point types are signed" + assert self.is_signed(), ( + "We currently assume all floating point types are signed" + ) sign_bit_double = 1 << 63 max_raw = self._floating_point_max_int() min_raw = max_raw | sign_bit_double return struct.unpack("!d", struct.pack("!Q", min_raw))[0] else: - assert ( - not self.is_signed() or self.size_bits <= 64 - ), "Cannot represent min as a int64_t" + assert not self.is_signed() or self.size_bits <= 64, ( + "Cannot represent min as a int64_t" + ) if self.is_signed(): return -(1 << (self.size_bits - 1)) diff --git a/python/sglang/kernels/aot/python/sgl_kernel/test_utils.py b/python/sglang/kernels/aot/python/sgl_kernel/test_utils.py index ede113fd0..4d7d1ca25 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/test_utils.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/test_utils.py @@ -94,9 +94,9 @@ def _compute_imbalanced_split( def assert_all_close_or_tiny_diff(a: torch.Tensor, b: torch.Tensor): - assert (a.shape == b.shape) and ( - a.dtype == b.dtype - ), f"{a.shape=} {b.shape=} {a.dtype=} {b.dtype=}" + assert (a.shape == b.shape) and (a.dtype == b.dtype), ( + f"{a.shape=} {b.shape=} {a.dtype=} {b.dtype=}" + ) numel = a.numel() if a.dtype == torch.float8_e4m3fn: diff --git a/python/sglang/kernels/aot/python/sgl_kernel/testing/rotary_embedding.py b/python/sglang/kernels/aot/python/sgl_kernel/testing/rotary_embedding.py index 1a4d90a9d..5a9eb6526 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/testing/rotary_embedding.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/testing/rotary_embedding.py @@ -112,9 +112,9 @@ class RotaryEmbedding(torch.nn.Module): fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """A PyTorch-native implementation of forward().""" - assert ( - fused_set_kv_buffer_arg is None - ), "fused_set_kv_buffer_arg is not supported for native implementation" + assert fused_set_kv_buffer_arg is None, ( + "fused_set_kv_buffer_arg is not supported for native implementation" + ) if offsets is not None: positions = positions + offsets @@ -182,9 +182,9 @@ class SglKernelRotaryEmbedding(RotaryEmbedding): offsets: Optional[torch.Tensor] = None, fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - assert ( - fused_set_kv_buffer_arg is None - ), "fused_set_kv_buffer_arg is not supported for sgl-kernel implementation" + assert fused_set_kv_buffer_arg is None, ( + "fused_set_kv_buffer_arg is not supported for sgl-kernel implementation" + ) if self.cos_sin_cache.dtype != query.dtype: self.cos_sin_cache = self.cos_sin_cache.to(query.dtype) torch.ops.sgl_kernel.rotary_embedding( diff --git a/python/sglang/kernels/aot/python/sgl_kernel/top_k.py b/python/sglang/kernels/aot/python/sgl_kernel/top_k.py index 4b842499a..c93dcd05c 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/top_k.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/top_k.py @@ -33,9 +33,9 @@ def fast_topk_v2( Returns: The topk indices tensor of shape (B, topk) """ - assert ( - topk == 2048 - ), "fast_topk_v2 is only optimized for deepseek v3.2 model, where topk=2048" + assert topk == 2048, ( + "fast_topk_v2 is only optimized for deepseek v3.2 model, where topk=2048" + ) assert score.dim() == 2 topk_indices = score.new_empty((score.size(0), topk), dtype=torch.int32) torch.ops.sgl_kernel.fast_topk(score, topk_indices, lengths, row_starts) @@ -68,9 +68,9 @@ def fast_topk_transform_fused( Returns: The topk indices tensor of shape (B, topk) """ - assert ( - topk == 2048 - ), "fast_topk_transform_fused is only optimized for deepseek v3.2 model, where topk=2048" + assert topk == 2048, ( + "fast_topk_transform_fused is only optimized for deepseek v3.2 model, where topk=2048" + ) assert score.dim() == 2 src_page_table = page_table_size_1 dst_page_table = score.new_empty((score.shape[0], topk), dtype=torch.int32) @@ -138,9 +138,9 @@ def fast_topk_transform_ragged_fused( Returns: The topk indices tensor of shape (B, topk) """ - assert ( - topk == 2048 - ), "fast_topk_transform_ragged_fused is only optimized for deepseek v3.2 model, where topk=2048" + assert topk == 2048, ( + "fast_topk_transform_ragged_fused is only optimized for deepseek v3.2 model, where topk=2048" + ) assert score.dim() == 2 topk_indices_ragged = score.new_empty((score.shape[0], topk), dtype=torch.int32) torch.ops.sgl_kernel.fast_topk_transform_ragged_fused( diff --git a/python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py b/python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py index 5a95f6e15..a828749d8 100644 --- a/python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py +++ b/python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py @@ -116,15 +116,15 @@ def test_tree_speculative_sampling_target_only( deterministic=True, ) - assert ( - predicts.tolist() == expected_predicts - ), f"Predicts mismatch for thresholds ({threshold_single}, {threshold_acc})" - assert ( - accept_index.tolist() == expected_accept_index - ), f"Accept index mismatch for thresholds ({threshold_single}, {threshold_acc})" - assert ( - accept_token_num.tolist() == expected_accept_token_num - ), f"Accept token num mismatch for thresholds ({threshold_single}, {threshold_acc})" + assert predicts.tolist() == expected_predicts, ( + f"Predicts mismatch for thresholds ({threshold_single}, {threshold_acc})" + ) + assert accept_index.tolist() == expected_accept_index, ( + f"Accept index mismatch for thresholds ({threshold_single}, {threshold_acc})" + ) + assert accept_token_num.tolist() == expected_accept_token_num, ( + f"Accept token num mismatch for thresholds ({threshold_single}, {threshold_acc})" + ) if __name__ == "__main__": diff --git a/python/sglang/kernels/aot/tests/test_custom_allreduce.py b/python/sglang/kernels/aot/tests/test_custom_allreduce.py index d729392c3..a60b4ae31 100644 --- a/python/sglang/kernels/aot/tests/test_custom_allreduce.py +++ b/python/sglang/kernels/aot/tests/test_custom_allreduce.py @@ -92,9 +92,9 @@ def multi_process_parallel( for i in range(world_size): procs[i].join() - assert ( - procs[i].exitcode == 0 - ), f"Process {i} failed with exit code {procs[i].exitcode}" + assert procs[i].exitcode == 0, ( + f"Process {i} failed with exit code {procs[i].exitcode}" + ) class TestCustomAllReduce(unittest.TestCase): diff --git a/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py b/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py index 1126894a4..0932ca18a 100644 --- a/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py +++ b/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py @@ -251,12 +251,14 @@ def test_sparse_attention( ref_out, ref_lse = ref_attn(q, k, v) - torch.testing.assert_close( - out, ref_out, atol=2e-2, rtol=1e-2 - ), f"{torch.max(torch.abs(out - ref_out))}" - torch.testing.assert_close( - lse, ref_lse, atol=2e-2, rtol=1e-2 - ), f"{torch.max(torch.abs(lse - ref_lse))}" + ( + torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=1e-2), + f"{torch.max(torch.abs(out - ref_out))}", + ) + ( + torch.testing.assert_close(lse, ref_lse, atol=2e-2, rtol=1e-2), + f"{torch.max(torch.abs(lse - ref_lse))}", + ) # sparse attention utils diff --git a/python/sglang/kernels/aot/tests/test_flashmla.py b/python/sglang/kernels/aot/tests/test_flashmla.py index 3afdd7866..f9b7574b6 100644 --- a/python/sglang/kernels/aot/tests/test_flashmla.py +++ b/python/sglang/kernels/aot/tests/test_flashmla.py @@ -198,9 +198,7 @@ def reference_torch_prefill( kvs = torch.index_select( kvs, 0, indices.masked_fill(invalid_indices_mask, 0).flatten() - ).view( - s_q, topk, 576 - ) # [s_q, topk, d_qk] + ).view(s_q, topk, 576) # [s_q, topk, d_qk] attn_score = qs @ kvs.transpose(1, 2) # [s_q, h_q, topk] attn_score.masked_fill_(invalid_indices_mask.unsqueeze(1), float("-inf")) attn_score *= sm_scale * math.log2(math.e) diff --git a/python/sglang/kernels/aot/tests/test_fused_qk_norm_rope.py b/python/sglang/kernels/aot/tests/test_fused_qk_norm_rope.py index 7f1daf5aa..57419b63b 100644 --- a/python/sglang/kernels/aot/tests/test_fused_qk_norm_rope.py +++ b/python/sglang/kernels/aot/tests/test_fused_qk_norm_rope.py @@ -76,9 +76,9 @@ def torch_ref_rms_norm_rope( v_size = num_heads_v * head_dim # Verify dimensions match - assert ( - hidden_size == q_size + k_size + v_size - ), f"Hidden size {hidden_size} doesn't match Q+K+V dimensions {q_size + k_size + v_size}" + assert hidden_size == q_size + k_size + v_size, ( + f"Hidden size {hidden_size} doesn't match Q+K+V dimensions {q_size + k_size + v_size}" + ) # Split the tensor into Q, K, V parts q = qkv[:, :q_size] diff --git a/python/sglang/kernels/aot/tests/test_moe_topk_sigmoid.py b/python/sglang/kernels/aot/tests/test_moe_topk_sigmoid.py index 6b2cba858..46a5cd5dc 100644 --- a/python/sglang/kernels/aot/tests/test_moe_topk_sigmoid.py +++ b/python/sglang/kernels/aot/tests/test_moe_topk_sigmoid.py @@ -44,13 +44,13 @@ def test_topk_sigmoid(num_tokens, num_experts, topk): topk_weights_ref, topk_indices_ref = torch.topk(sigmoid_output, topk, dim=-1) # Verify the top-k weights and indices match the torch native ones - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" + ) - assert torch.allclose( - topk_indices_ref.int(), topk_indices, atol=0, rtol=0 - ), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" + assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), ( + f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" + ) @pytest.mark.parametrize( @@ -87,13 +87,13 @@ def test_topk_sigmoid_dtype_regression(num_tokens, num_experts, topk, dtype): gating_output.float(), ) - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: SGLang old interface={topk_weights_ref} vs SGLang new interface={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: SGLang old interface={topk_weights_ref} vs SGLang new interface={topk_weights}" + ) - assert torch.allclose( - topk_indices_ref.int(), topk_indices, atol=0, rtol=0 - ), f"Indices mismatch: SGLang old interface={topk_indices_ref}, SGLang new interface={topk_indices}" + assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), ( + f"Indices mismatch: SGLang old interface={topk_indices_ref}, SGLang new interface={topk_indices}" + ) @pytest.mark.parametrize( @@ -136,13 +136,13 @@ def test_topk_sigmoid_renormalize(num_tokens, num_experts, topk): ) topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True) - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: SGLang w/o fused renormalize={topk_weights_ref} vs SGLang w/ fused renormalize={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: SGLang w/o fused renormalize={topk_weights_ref} vs SGLang w/ fused renormalize={topk_weights}" + ) - assert torch.allclose( - topk_indices_ref.int(), topk_indices, atol=0, rtol=0 - ), f"Indices mismatch: SGLang w/o fused renormalize={topk_indices_ref}, SGLang w/ fused renormalize={topk_indices}" + assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), ( + f"Indices mismatch: SGLang w/o fused renormalize={topk_indices_ref}, SGLang w/ fused renormalize={topk_indices}" + ) @pytest.mark.parametrize( @@ -180,13 +180,13 @@ def test_topk_sigmoid_renormalize_correction_bias(num_tokens, num_experts, topk) topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True) # Verify the top-k weights and indices match the torch native ones - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" + ) - assert torch.allclose( - topk_indices_ref.int(), topk_indices, atol=0, rtol=0 - ), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" + assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), ( + f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" + ) if __name__ == "__main__": diff --git a/python/sglang/kernels/aot/tests/test_moe_topk_softmax.py b/python/sglang/kernels/aot/tests/test_moe_topk_softmax.py index 77ffe8a46..80c04f983 100644 --- a/python/sglang/kernels/aot/tests/test_moe_topk_softmax.py +++ b/python/sglang/kernels/aot/tests/test_moe_topk_softmax.py @@ -41,13 +41,13 @@ def test_topkfast_softmax(num_tokens, num_experts, topk): topk_weights_ref, topk_indices_ref = torch.topk(softmax_output, topk, dim=-1) # Verify the top-k weights and indices match the torch native ones - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}" + ) - assert compare_topk_values( - gating_output, topk_indices_ref.int(), topk_indices - ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), ( + f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + ) @pytest.mark.parametrize( @@ -79,13 +79,13 @@ def test_topk_softmax(num_tokens, num_experts, topk): topk_weights_ref, topk_indices_ref = torch.topk(softmax_output, topk, dim=-1) # Verify the top-k weights and indices match the torch native ones - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}" + ) - assert compare_topk_values( - gating_output, topk_indices_ref.int(), topk_indices - ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), ( + f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + ) @pytest.mark.parametrize( @@ -122,13 +122,13 @@ def test_topk_softmax_dtype_regression(num_tokens, num_experts, topk, dtype): gating_output.float(), ) - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: SGLang old interface={topk_indices_ref} vs SGLang new interface={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: SGLang old interface={topk_indices_ref} vs SGLang new interface={topk_weights}" + ) - assert compare_topk_values( - gating_output, topk_indices_ref.int(), topk_indices - ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), ( + f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + ) @pytest.mark.parametrize( @@ -171,13 +171,13 @@ def test_topk_softmax_renormalize(num_tokens, num_experts, topk): ) topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True) - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: SGLang w/o fused renormalize={topk_indices_ref} vs SGLang w/ fused renormalize={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: SGLang w/o fused renormalize={topk_indices_ref} vs SGLang w/ fused renormalize={topk_weights}" + ) - assert compare_topk_values( - gating_output, topk_indices_ref.int(), topk_indices - ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), ( + f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + ) if __name__ == "__main__": diff --git a/python/sglang/kernels/jit/__main__.py b/python/sglang/kernels/jit/__main__.py index b69b5fef7..6507e71d4 100644 --- a/python/sglang/kernels/jit/__main__.py +++ b/python/sglang/kernels/jit/__main__.py @@ -72,9 +72,9 @@ def generate_clangd(): arch = make_jit_cuda_arch(int(major), int(minor)) else: arch = get_jit_cuda_arch() - assert ( - arch.major > 0 - ), "Cannot detect CUDA architecture, please specify --cuda-target explicitly." + assert arch.major > 0, ( + "Cannot detect CUDA architecture, please specify --cuda-target explicitly." + ) compile_flags = [ "-xcuda", diff --git a/python/sglang/kernels/jit/benchmark/marker.py b/python/sglang/kernels/jit/benchmark/marker.py index e34a00e31..71bf3fd9b 100644 --- a/python/sglang/kernels/jit/benchmark/marker.py +++ b/python/sglang/kernels/jit/benchmark/marker.py @@ -253,9 +253,9 @@ class Benchmark(Generic[F]): f"parametrize name {name!r} is not a parameter of " f"{self._fn.__name__}; available: {list(self._fn_params)}" ) - assert ( - name not in self._seen_args - ), f"parametrize name {name!r} is already used" + assert name not in self._seen_args, ( + f"parametrize name {name!r} is already used" + ) self._seen_args.add(name) self._configs.insert(0, (names, vals)) @@ -305,8 +305,7 @@ class Benchmark(Generic[F]): if p.default is inspect.Parameter.empty and p.kind in kinds } - (set(flat_names) | {self._line_arg}) assert not missing, ( - f"parameters not parametrized for {self._fn.__name__}: " - f"{sorted(missing)}" + f"parameters not parametrized for {self._fn.__name__}: {sorted(missing)}" ) results, bandwidths, should_log_bw = self._collect_results() @@ -360,13 +359,13 @@ def parametrize(names: str, vals: List[Any], ci_vals: Optional[List[Any]] = None return [(v,) for v in vs] out: List[Tuple[Any, ...]] = [] for v in vs: - assert isinstance( - v, (tuple, list) - ), f"parametrize: multi-name values must be tuples, got {v!r}" + assert isinstance(v, (tuple, list)), ( + f"parametrize: multi-name values must be tuples, got {v!r}" + ) t = tuple(v) - assert ( - len(t) == arity - ), f"parametrize: each value must have length {arity}, got {t!r}" + assert len(t) == arity, ( + f"parametrize: each value must have length {arity}, got {t!r}" + ) out.append(t) return out diff --git a/python/sglang/kernels/jit/utils/compile/loader.py b/python/sglang/kernels/jit/utils/compile/loader.py index 89fecce45..7b810a445 100644 --- a/python/sglang/kernels/jit/utils/compile/loader.py +++ b/python/sglang/kernels/jit/utils/compile/loader.py @@ -121,7 +121,7 @@ def load_jit( # Also the benign case where a concurrent GC unlinked the leaf # between the lookup and the load. logger.warning( - "Cached JIT module %s failed to load; rebuilding. " "Got error: %s", + "Cached JIT module %s failed to load; rebuilding. Got error: %s", spec.module_name, e, ) diff --git a/python/sglang/kernels/kda_kernels/causal_conv3d_cat_pad_jit.py b/python/sglang/kernels/kda_kernels/causal_conv3d_cat_pad_jit.py index 5d5bebd69..896873a4d 100644 --- a/python/sglang/kernels/kda_kernels/causal_conv3d_cat_pad_jit.py +++ b/python/sglang/kernels/kda_kernels/causal_conv3d_cat_pad_jit.py @@ -25,7 +25,7 @@ def _jit_causal_conv3d_cat_pad_module(dtype: torch.dtype) -> Module: cuda_wrappers=[ ( "causal_conv3d_cat_pad", - "causal_conv3d_cat_pad::" f"CausalConv3dCatPadKernel<{args}>::run", + f"causal_conv3d_cat_pad::CausalConv3dCatPadKernel<{args}>::run", ) ], ) diff --git a/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py b/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py index 607d1e634..45aaa9375 100644 --- a/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py +++ b/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py @@ -353,9 +353,9 @@ class _Qwen3xNvfp4Sm120Kernel: self.occupancy, ) - assert ( - self.epi_stage > 0 - ), "epi_stage <= 0, not enough shared memory. This configuration will be skipped." + assert self.epi_stage > 0, ( + "epi_stage <= 0, not enough shared memory. This configuration will be skipped." + ) ( self.a_smem_layout_staged, diff --git a/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py b/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py index fc17fa6c5..594b0cad8 100644 --- a/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py +++ b/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py @@ -34,11 +34,11 @@ def _jit_residual_gate_add_module(dtype: torch.dtype) -> Module: cuda_wrappers=[ ( "residual_gate_add", - "residual_gate_add::" f"ResidualGateAddKernel<{args}>::run", + f"residual_gate_add::ResidualGateAddKernel<{args}>::run", ), ( "residual_gate_add_transposed", - "residual_gate_add::" f"ResidualGateAddKernel<{args}>::run_transposed", + f"residual_gate_add::ResidualGateAddKernel<{args}>::run_transposed", ), ], ) diff --git a/python/sglang/kernels/ops/activation/activation.py b/python/sglang/kernels/ops/activation/activation.py index bf2f12cc8..f5765181e 100644 --- a/python/sglang/kernels/ops/activation/activation.py +++ b/python/sglang/kernels/ops/activation/activation.py @@ -157,9 +157,9 @@ def run_unary_activation( Unlike :func:`run_activation`, there is no gate/up split โ ``input`` and ``out`` share the same shape. """ - assert ( - op_name in SUPPORTED_UNARY_ACTIVATIONS - ), f"Unsupported unary activation: {op_name}" + assert op_name in SUPPORTED_UNARY_ACTIVATIONS, ( + f"Unsupported unary activation: {op_name}" + ) if out is None: out = torch.empty_like(input) _run_unary_activation_inplace(op_name, input, out) diff --git a/python/sglang/kernels/ops/activation/softcap.py b/python/sglang/kernels/ops/activation/softcap.py index 5d39de6f8..8b719c9b9 100644 --- a/python/sglang/kernels/ops/activation/softcap.py +++ b/python/sglang/kernels/ops/activation/softcap.py @@ -101,9 +101,9 @@ def softcap_inplace_logits(full_logits, final_logit_softcapping): row_stride = ncols else: assert full_logits.ndim == 2, "non-contiguous softcap requires 2D tensor" - assert ( - full_logits.stride(1) == 1 - ), "non-contiguous softcap requires contiguous columns" + assert full_logits.stride(1) == 1, ( + "non-contiguous softcap requires contiguous columns" + ) nrows, ncols = full_logits.shape row_stride = full_logits.stride(0) diff --git a/python/sglang/kernels/ops/attention/cutedsl_fp8_paged_mqa_logits.py b/python/sglang/kernels/ops/attention/cutedsl_fp8_paged_mqa_logits.py index b74401168..e0450405e 100644 --- a/python/sglang/kernels/ops/attention/cutedsl_fp8_paged_mqa_logits.py +++ b/python/sglang/kernels/ops/attention/cutedsl_fp8_paged_mqa_logits.py @@ -221,12 +221,12 @@ class FP8MQALogitsKernel: self.block_kv = block_kv self.phys_block_kv = phys_block_kv self.num_blocks_per_mma = block_kv // phys_block_kv - assert ( - block_kv % phys_block_kv == 0 - ), f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}" - assert ( - self.num_blocks_per_mma <= 4 - ), f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4" + assert block_kv % phys_block_kv == 0, ( + f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}" + ) + assert self.num_blocks_per_mma <= 4, ( + f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4" + ) self.remove_kv_wait_in_epilogue = remove_kv_wait_in_epilogue self.early_tmem_copy = early_tmem_copy self.smem_subpartition_opt = smem_subpartition_opt diff --git a/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py b/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py index cf50bb4cb..64c0817b0 100644 --- a/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py +++ b/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py @@ -3080,9 +3080,9 @@ def gated_delta_rule_mtp_wide_vec( assert K_val == 128 and V_val == 128 assert initial_state_source.dtype == torch.bfloat16 assert tile_v in (32, 64, 128), f"tile_v must be 32/64/128, got {tile_v}" - assert ( - V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0 - ), f"tile_v={tile_v} incompatible with 8 groups ร ILP=4 layout" + assert V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0, ( + f"tile_v={tile_v} incompatible with 8 groups ร ILP=4 layout" + ) if cache_ring: assert replayssm_rawv is not None and replayssm_rawk is not None @@ -3194,13 +3194,13 @@ def gated_delta_rule_mtp_wide_vec( ) # Validate recovery_steps for fused recovery+decode mode. - assert ( - 0 <= recovery_steps <= T_val - ), f"recovery_steps must be in [0, T={T_val}], got {recovery_steps}" + assert 0 <= recovery_steps <= T_val, ( + f"recovery_steps must be in [0, T={T_val}], got {recovery_steps}" + ) if recovery_steps > 0: - assert ( - not cache_intermediate_states - ), "recovery_steps > 0 is incompatible with intermediate state caching" + assert not cache_intermediate_states, ( + "recovery_steps > 0 is incompatible with intermediate state caching" + ) assert not disable_state_update, ( "recovery_steps > 0 requires state writeback " "(disable_state_update=False); the boundary writeback at i_t=K-1 " @@ -3220,12 +3220,12 @@ def gated_delta_rule_mtp_wide_vec( # accepted_steps[i] is the per-request phase boundary. per_request_accepted_steps = accepted_steps is not None if per_request_accepted_steps: - assert accepted_steps.shape == ( - B_val, - ), f"accepted_steps must have shape [B={B_val}], got {accepted_steps.shape}" - assert ( - accepted_steps.dtype == torch.int32 - ), f"accepted_steps must be int32, got {accepted_steps.dtype}" + assert accepted_steps.shape == (B_val,), ( + f"accepted_steps must have shape [B={B_val}], got {accepted_steps.shape}" + ) + assert accepted_steps.dtype == torch.int32, ( + f"accepted_steps must be int32, got {accepted_steps.dtype}" + ) assert accepted_steps.device == q.device # FLA-style per-token pool scatter (vLLM API compat). When the public @@ -3236,15 +3236,15 @@ def gated_delta_rule_mtp_wide_vec( # this entry point hit the same fail-fast errors. per_token_pool_scatter = ssm_state_indices is not None if per_token_pool_scatter: - assert ( - intermediate_states_buffer is None - ), "ssm_state_indices and intermediate_states_buffer are mutually exclusive" - assert ( - not disable_state_update - ), "ssm_state_indices requires state writes; disable_state_update must be False" - assert ( - recovery_steps == 0 - ), "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)" + assert intermediate_states_buffer is None, ( + "ssm_state_indices and intermediate_states_buffer are mutually exclusive" + ) + assert not disable_state_update, ( + "ssm_state_indices requires state writes; disable_state_update must be False" + ) + assert recovery_steps == 0, ( + "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)" + ) assert T_val >= 2, ( f"ssm_state_indices requires T >= 2 (got T={T_val}); " f"for T=1 use output_state_indices" @@ -3253,9 +3253,9 @@ def gated_delta_rule_mtp_wide_vec( f"ssm_state_indices must have shape [B={B_val}, T={T_val}], " f"got {tuple(ssm_state_indices.shape)}" ) - assert ( - ssm_state_indices.dtype == torch.int32 - ), f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}" + assert ssm_state_indices.dtype == torch.int32, ( + f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}" + ) assert ssm_state_indices.device == q.device phase_b_unroll = _select_wide_vec_phase_b_unroll( @@ -3517,9 +3517,9 @@ def gated_delta_rule_t1_wide_vec( assert K_val == 128 and V_val == 128 assert initial_state_source.dtype == torch.bfloat16 assert tile_v in (32, 64, 128), f"tile_v must be 32/64/128, got {tile_v}" - assert ( - V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0 - ), f"tile_v={tile_v} incompatible with 8 groups ร ILP=4 layout" + assert V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0, ( + f"tile_v={tile_v} incompatible with 8 groups ร ILP=4 layout" + ) if scale is None: scale = 1.0 / math.sqrt(K_val) @@ -3827,9 +3827,9 @@ def gated_delta_rule_mtp( f"intermediate_states_buffer dim 0 ({buffer_size}) must equal " f"batch size B={B}; the buffer is batch-scoped, not pool-scoped" ) - assert ( - cache_steps >= T - ), f"intermediate_states_buffer dim 1 ({cache_steps}) must be >= T={T}" + assert cache_steps >= T, ( + f"intermediate_states_buffer dim 1 ({cache_steps}) must be >= T={T}" + ) assert intermediate_states_buffer.dtype == torch.bfloat16 intermediate_states = intermediate_states_buffer.reshape( B * cache_steps * HV, V, K @@ -3860,28 +3860,28 @@ def gated_delta_rule_mtp( # results/2026-06-03/FLA_SCATTER_MODE_PLAN.md. per_token_pool_scatter = ssm_state_indices is not None if per_token_pool_scatter: - assert ( - intermediate_states_buffer is None - ), "ssm_state_indices and intermediate_states_buffer are mutually exclusive" - assert ( - not disable_state_update - ), "ssm_state_indices requires state writes; disable_state_update must be False" - assert ( - recovery_steps == 0 - ), "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)" - assert ( - T >= 2 - ), f"ssm_state_indices requires T >= 2 (got T={T}); for T=1 use output_state_indices" + assert intermediate_states_buffer is None, ( + "ssm_state_indices and intermediate_states_buffer are mutually exclusive" + ) + assert not disable_state_update, ( + "ssm_state_indices requires state writes; disable_state_update must be False" + ) + assert recovery_steps == 0, ( + "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)" + ) + assert T >= 2, ( + f"ssm_state_indices requires T >= 2 (got T={T}); for T=1 use output_state_indices" + ) assert ssm_state_indices.shape == (B, T), ( f"ssm_state_indices must have shape [B={B}, T={T}], " f"got {tuple(ssm_state_indices.shape)}" ) - assert ( - ssm_state_indices.dtype == torch.int32 - ), f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}" - assert ( - ssm_state_indices.device == q.device - ), f"ssm_state_indices device {ssm_state_indices.device} != q device {q.device}" + assert ssm_state_indices.dtype == torch.int32, ( + f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}" + ) + assert ssm_state_indices.device == q.device, ( + f"ssm_state_indices device {ssm_state_indices.device} != q device {q.device}" + ) # Dispatch to the wide_vec kernel when work_units (B*HV) amortizes its # lower per-CTA parallelism. ``_select_wide_vec_tile_v`` picks tile_v @@ -3960,12 +3960,12 @@ def gated_delta_rule_mtp( # Per-request K opt-in (see gated_delta_rule_mtp_wide_vec for full rationale). per_request_accepted_steps = accepted_steps is not None if per_request_accepted_steps: - assert accepted_steps.shape == ( - B, - ), f"accepted_steps must have shape [B={B}], got {accepted_steps.shape}" - assert ( - accepted_steps.dtype == torch.int32 - ), f"accepted_steps must be int32, got {accepted_steps.dtype}" + assert accepted_steps.shape == (B,), ( + f"accepted_steps must have shape [B={B}], got {accepted_steps.shape}" + ) + assert accepted_steps.dtype == torch.int32, ( + f"accepted_steps must be int32, got {accepted_steps.dtype}" + ) assert accepted_steps.device == q.device # Contiguous pool -> sentinel keys + slot dim marked dynamic (pool-size diff --git a/python/sglang/kernels/ops/attention/cutedsl_kda.py b/python/sglang/kernels/ops/attention/cutedsl_kda.py index 78c7b8495..312769e9d 100644 --- a/python/sglang/kernels/ops/attention/cutedsl_kda.py +++ b/python/sglang/kernels/ops/attention/cutedsl_kda.py @@ -1427,12 +1427,12 @@ def cutedsl_fused_sigmoid_gating_kda_update( N = initial_state_indices.shape[0] assert K == TILE_K, f"Current CuTe DSL KDA kernel requires K={TILE_K}, got {K}" - assert ( - V % TILE_V_SMALL == 0 - ), f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0, got V={V}" - assert ( - V % TILE_V == 0 - ), f"Current CuTe DSL KDA kernel requires V % {TILE_V} == 0, got V={V}" + assert V % TILE_V_SMALL == 0, ( + f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0, got V={V}" + ) + assert V % TILE_V == 0, ( + f"Current CuTe DSL KDA kernel requires V % {TILE_V} == 0, got V={V}" + ) assert (V // TILE_V_SMALL) % NUM_BLOCKS_PER_STATE_SMALL == 0, ( "Small-batch KDA kernel requires num_v_tiles_small divisible by " f"{NUM_BLOCKS_PER_STATE_SMALL}, got V={V}" diff --git a/python/sglang/kernels/ops/attention/decode_attention.py b/python/sglang/kernels/ops/attention/decode_attention.py index e09ce7fef..b535e19e2 100644 --- a/python/sglang/kernels/ops/attention/decode_attention.py +++ b/python/sglang/kernels/ops/attention/decode_attention.py @@ -1483,7 +1483,6 @@ def _lean_attention_decode_kernel( # Use a regular while loop instead of tl.static_range with a dynamic bound to avoid # Triton compiler crashes in the Coalesce pass (max_output_tile_cnt is runtime-computed). while iter < cta_end_tile_gid: - tile_row_idx = iter // tiles_per_khead tile_idx = tile_row_idx * batch_size tile_iter = tile_row_idx * tiles_per_khead diff --git a/python/sglang/kernels/ops/attention/deepseek_v4_rope.py b/python/sglang/kernels/ops/attention/deepseek_v4_rope.py index 1dd647465..fb0ee2c79 100644 --- a/python/sglang/kernels/ops/attention/deepseek_v4_rope.py +++ b/python/sglang/kernels/ops/attention/deepseek_v4_rope.py @@ -354,9 +354,9 @@ def apply_rotary_emb_triton( grid = (batch_size, n_heads if is_3d else 1, num_blocks_dim) if positions is not None: - assert positions.shape == ( - batch_size, - ), f"positions shape {positions.shape} != ({batch_size},)" + assert positions.shape == (batch_size,), ( + f"positions shape {positions.shape} != ({batch_size},)" + ) apply_rotary_emb_triton_kernel[grid]( x, @@ -374,9 +374,9 @@ def apply_rotary_emb_triton( BLOCK_SIZE=BLOCK_SIZE, ) else: - assert ( - freqs_real.shape[0] == batch_size - ), f"freqs_cis batch size {freqs_real.shape[0]} != x batch size {batch_size}" + assert freqs_real.shape[0] == batch_size, ( + f"freqs_cis batch size {freqs_real.shape[0]} != x batch size {batch_size}" + ) apply_rotary_emb_triton_kernel[grid]( x, @@ -621,9 +621,9 @@ def fused_norm_rope_inplace_triton( if weight is not None: assert weight.shape == (head_dim,) if positions is None: - assert ( - freqs_real.shape[0] == M - ), f"freqs_cis row count {freqs_real.shape[0]} != M={M}" + assert freqs_real.shape[0] == M, ( + f"freqs_cis row count {freqs_real.shape[0]} != M={M}" + ) else: assert positions.shape == (M,) and positions.dim() == 1 diff --git a/python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py b/python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py index af01f8a11..55562283f 100644 --- a/python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py +++ b/python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py @@ -181,9 +181,9 @@ def dequantize_k_cache_paged( output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache """ dim_quant = quant_k_cache.shape[-1] - assert ( - dim_quant == 656 - ), f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged" + assert dim_quant == 656, ( + f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged" + ) quant_k_cache = quant_k_cache.view((-1, dim_quant)) # num_tokens can exceed kv_cache_size due to prefix sharing (multiple seqs share same KV slots) diff --git a/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py b/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py index 2cd4dd539..6a7165b7d 100644 --- a/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py +++ b/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py @@ -308,9 +308,9 @@ def _set_k_and_s_triton( assert scale_dim == 1 if _is_hip: if _use_aiter_preshuffle: - assert ( - page_size % 16 == 0 - ), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}" + assert page_size % 16 == 0, ( + f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}" + ) else: assert page_size == 64 diff --git a/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py b/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py index d2dc85132..23d5e3025 100644 --- a/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py +++ b/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py @@ -155,9 +155,9 @@ def act_quant( - A tensor of scaling factors with dtype `torch.float32`. """ assert x.is_contiguous(), "Input tensor must be contiguous" - assert ( - x.size(-1) % block_size == 0 - ), f"Last dimension size must be divisible by block_size (block_size={block_size})" + assert x.size(-1) % block_size == 0, ( + f"Last dimension size must be divisible by block_size (block_size={block_size})" + ) N = x.size(-1) if _is_fp8_fnuz: y = torch.empty_like(x, dtype=torch.float8_e4m3fnuz) @@ -272,16 +272,16 @@ def sparse_attention_fwd_kernel_v1( num_stages=2, threads=256, ): - assert dim == tilelang.math.next_power_of_2( - dim - ), f"haven't check padding correctness yet, dim={dim}" - assert tail_dim == tilelang.math.next_power_of_2( - tail_dim - ), f"haven't check padding correctness yet, dim={tail_dim}" + assert dim == tilelang.math.next_power_of_2(dim), ( + f"haven't check padding correctness yet, dim={dim}" + ) + assert tail_dim == tilelang.math.next_power_of_2(tail_dim), ( + f"haven't check padding correctness yet, dim={tail_dim}" + ) assert is_causal == True, "non-casual is not supported" - assert ( - topk % block_I == 0 - ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + assert topk % block_I == 0, ( + "otherwise will load some index=0 thus causing wrong kv to be loaded" + ) if sm_scale is None: sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) else: @@ -361,7 +361,6 @@ def sparse_attention_fwd_kernel_v1( T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) for i_i in T.Pipelined(NI, num_stages=num_stages): - for bi_i in T.Parallel(BI): mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] >= 0 @@ -446,15 +445,15 @@ def sparse_attention_fwd_kernel_v2( sm_scale: Optional[float] = None, block_I: int = 64, ): - assert dim == tilelang.math.next_power_of_2( - dim - ), f"haven't check padding correctness yet, dim={dim}" - assert tail_dim == tilelang.math.next_power_of_2( - tail_dim - ), f"haven't check padding correctness yet, dim={tail_dim}" - assert ( - topk % block_I == 0 - ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + assert dim == tilelang.math.next_power_of_2(dim), ( + f"haven't check padding correctness yet, dim={dim}" + ) + assert tail_dim == tilelang.math.next_power_of_2(tail_dim), ( + f"haven't check padding correctness yet, dim={tail_dim}" + ) + assert topk % block_I == 0, ( + "otherwise will load some index=0 thus causing wrong kv to be loaded" + ) if sm_scale is None: sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) else: @@ -1078,9 +1077,9 @@ def sparse_mla_fwd_decode_partial_fp8( threads=256, ): assert d_v == 512, f"only support d_v=512" - assert ( - topk % block_I == 0 - ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + assert topk % block_I == 0, ( + "otherwise will load some index=0 thus causing wrong kv to be loaded" + ) # Softmax scores are in [0, 1]. We scale by fp8_max_val before FP8 cast # to better utilize FP8 dynamic range, then apply the inverse scale after GEMM. @@ -1104,9 +1103,9 @@ def sparse_mla_fwd_decode_partial_fp8( h_per_block = 16 # Match bf16 partial behavior: keep fixed 16-head tiles and use # sliced T.copy on H0:H1 for tail handling. - assert ( - num_heads <= h_per_block or num_heads % h_per_block == 0 - ), "num_heads must be <=16 or divisible by 16" + assert num_heads <= h_per_block or num_heads % h_per_block == 0, ( + "num_heads must be <=16 or divisible by 16" + ) head_blocks_per_seq = (num_heads + h_per_block - 1) // h_per_block batch = 1 @@ -1594,9 +1593,7 @@ def dpsk_v4_fp8_partial_kernel( sm_scale = sm_scale * log2e assert dim == 448 and tail_dim == 64 assert topk_1 % block_I == 0 - assert ( - topk_1 // block_I - ) % inner_iter_1 == 0, ( + assert (topk_1 // block_I) % inner_iter_1 == 0, ( f"NI_1={topk_1 // block_I} must be divisible by inner_iter_1={inner_iter_1}" ) assert block_size_kv_1 > 0 and (block_size_kv_1 & (block_size_kv_1 - 1)) == 0 @@ -1605,9 +1602,7 @@ def dpsk_v4_fp8_partial_kernel( if is_dual: assert inner_iter_2 > 0, "dual-cache call requires inner_iter_2 > 0" assert topk_2 % block_I == 0 - assert ( - topk_2 // block_I - ) % inner_iter_2 == 0, ( + assert (topk_2 // block_I) % inner_iter_2 == 0, ( f"NI_2={topk_2 // block_I} must be divisible by inner_iter_2={inner_iter_2}" ) assert block_size_kv_2 > 0 and (block_size_kv_2 & (block_size_kv_2 - 1)) == 0 @@ -2256,12 +2251,8 @@ def dpsk_v4_combine_kernel( @T.prim_func def main( - Partial_O: T.Tensor( - [batch, seq_len, n_groups, num_heads, DT], BF16 - ), # type: ignore - Partial_LSE: T.Tensor( - [batch, seq_len, n_groups, num_heads], accum_dtype - ), # type: ignore + Partial_O: T.Tensor([batch, seq_len, n_groups, num_heads, DT], BF16), # type: ignore + Partial_LSE: T.Tensor([batch, seq_len, n_groups, num_heads], accum_dtype), # type: ignore Topk_length_1: T.Tensor([batch], INT32), # type: ignore Topk_length_2: T.Tensor([batch], INT32), # type: ignore Attn_sink: T.Tensor([num_heads], FP32), # type: ignore @@ -2369,12 +2360,8 @@ def dpsk_v4_combine_kernel( @T.prim_func def main( - Partial_O: T.Tensor( - [batch, seq_len, n_groups, num_heads, DT], BF16 - ), # type: ignore - Partial_LSE: T.Tensor( - [batch, seq_len, n_groups, num_heads], accum_dtype - ), # type: ignore + Partial_O: T.Tensor([batch, seq_len, n_groups, num_heads, DT], BF16), # type: ignore + Partial_LSE: T.Tensor([batch, seq_len, n_groups, num_heads], accum_dtype), # type: ignore Attn_sink: T.Tensor([num_heads], FP32), # type: ignore Output: T.Tensor([batch, seq_len, num_heads, DT], BF16), # type: ignore LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype), # type: ignore diff --git a/python/sglang/kernels/ops/attention/dsa/triton_kernel.py b/python/sglang/kernels/ops/attention/dsa/triton_kernel.py index 0d2969804..4ce7a4c1e 100644 --- a/python/sglang/kernels/ops/attention/dsa/triton_kernel.py +++ b/python/sglang/kernels/ops/attention/dsa/triton_kernel.py @@ -99,9 +99,9 @@ def act_quant( - A tensor of scaling factors with dtype `torch.float32`. """ assert x.is_contiguous(), "Input tensor must be contiguous" - assert ( - x.size(-1) % block_size == 0 - ), f"Last dimension size must be divisible by block_size (block_size={block_size})" + assert x.size(-1) % block_size == 0, ( + f"Last dimension size must be divisible by block_size (block_size={block_size})" + ) # Flatten all dims except last N = x.size(-1) diff --git a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py index 2c6ea0c9e..0a3c06310 100644 --- a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py +++ b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py @@ -69,9 +69,7 @@ def _sparse_mla_fwd_kernel( ) # [H, D_V] q_tail = tl.load( q_rope_ptr + s_i * H * D_TAIL + h[:, None] * D_TAIL + dt[None, :] - ).to( - q_nope_ptr.dtype.element_ty - ) # [H, D_TAIL] + ).to(q_nope_ptr.dtype.element_ty) # [H, D_TAIL] m_i = tl.full([H], -float("inf"), tl.float32) l_i = tl.zeros([H], tl.float32) @@ -89,9 +87,7 @@ def _sparse_mla_fwd_kernel( ) # [BLOCK_N, D_V] -- reused as V kv_tail = tl.load( kbase + (D_V + dt)[None, :], mask=valid[:, None], other=0.0 - ).to( - q_nope_ptr.dtype.element_ty - ) # [BLOCK_N, D_TAIL] + ).to(q_nope_ptr.dtype.element_ty) # [BLOCK_N, D_TAIL] qk = tl.dot(q_main, tl.trans(kv_main)).to(tl.float32) qk += tl.dot(q_tail, tl.trans(kv_tail)).to(tl.float32) diff --git a/python/sglang/kernels/ops/attention/dsv4/index_buf_accessor.py b/python/sglang/kernels/ops/attention/dsv4/index_buf_accessor.py index 8536e18a9..50738d8f8 100644 --- a/python/sglang/kernels/ops/attention/dsv4/index_buf_accessor.py +++ b/python/sglang/kernels/ops/attention/dsv4/index_buf_accessor.py @@ -209,7 +209,9 @@ def _set_k_and_s_torch( == num_tokens_to_write_nope == num_tokens_to_write_rope == num_tokens_to_write_scale - ), f"{num_tokens_to_write=} {num_tokens_to_write_nope=} {num_tokens_to_write_rope=} {num_tokens_to_write_scale=}" + ), ( + f"{num_tokens_to_write=} {num_tokens_to_write_nope=} {num_tokens_to_write_rope=} {num_tokens_to_write_scale=}" + ) assert buf.dtype == torch.uint8 assert loc.dtype in [ diff --git a/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py b/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py index efe579f1a..d477a0238 100644 --- a/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py +++ b/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py @@ -110,9 +110,9 @@ def _init_compressed_attn_metadata_triton( # no cache-write locations. Keep the write buffers unpadded and mask those # rows in the kernel. num_write_tokens = raw_out_loc.shape[0] - assert ( - num_write_tokens <= bs - ), f"raw_out_loc has {num_write_tokens} rows, expected at most {bs} metadata rows" + assert num_write_tokens <= bs, ( + f"raw_out_loc has {num_write_tokens} rows, expected at most {bs} metadata rows" + ) device = seq_lens.device c4_out_loc = torch.empty(num_write_tokens, dtype=torch.int64, device=device) @@ -126,12 +126,12 @@ def _init_compressed_attn_metadata_triton( c128_seq_lens_clamp1 = torch.empty(bs, dtype=torch.int32, device=device) if compute_page_indices: - assert ( - page_table is not None - ), "page_table required when compute_page_indices=True" - assert ( - page_size >= 128 and page_size % 128 == 0 - ), "page_size must be a multiple of 128 when compute_page_indices=True" + assert page_table is not None, ( + "page_table required when compute_page_indices=True" + ) + assert page_size >= 128 and page_size % 128 == 0, ( + "page_size must be a multiple of 128 when compute_page_indices=True" + ) max_pages = page_table.shape[1] c128_page_size = page_size // 128 c128_cur_max_seq_len = c128_page_size * max_pages diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py b/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py index 4188b95c4..5245639f5 100644 --- a/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py +++ b/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py @@ -1090,9 +1090,9 @@ class FlashAttentionForwardSm120(FlashAttentionForwardBase): ): assert blocksparse_tensors is None, "Block sparsity is not supported on SM120" assert (mBias is not None) == self.has_bias - assert ( - mPageTable is None or self.paged_kv - ), "SM120 paged KV requires the dedicated DMA-warp specialization" + assert mPageTable is None or self.paged_kv, ( + "SM120 paged KV requires the dedicated DMA-warp specialization" + ) self._check_type( *( t.element_type if t is not None else None @@ -1251,7 +1251,9 @@ class FlashAttentionForwardSm120(FlashAttentionForwardBase): TileScheduler = ( Sm120UniformBatchScheduler if is_varlen and self.direct_uniform_batch - else SingleTileVarlenScheduler if is_varlen else SingleTileScheduler + else SingleTileVarlenScheduler + if is_varlen + else SingleTileScheduler ) tile_sched_args = TileSchedulerArguments( num_block=cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m), diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py b/python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py index 9cc7cb365..5fe3df516 100644 --- a/python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py +++ b/python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py @@ -75,9 +75,9 @@ class Sm120UniformBatchScheduler: loc=None, ip=None, ) -> Params: - assert ( - scheduling_mode == SchedulingMode.STATIC - ), f"SM120 uniform-batch scheduler only supports STATIC, got {scheduling_mode!r}" + assert scheduling_mode == SchedulingMode.STATIC, ( + f"SM120 uniform-batch scheduler only supports STATIC, got {scheduling_mode!r}" + ) return Sm120UniformBatchScheduler.Params.create(args, loc=loc, ip=ip) @staticmethod diff --git a/python/sglang/kernels/ops/attention/fla/chunk.py b/python/sglang/kernels/ops/attention/fla/chunk.py index 8acb9d14e..80cfb4e1f 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk.py +++ b/python/sglang/kernels/ops/attention/fla/chunk.py @@ -85,7 +85,6 @@ def chunk_gated_delta_rule_fwd( class ChunkGatedDeltaRuleFunction(torch.autograd.Function): - @staticmethod @input_guard @autocast_custom_fwd @@ -207,12 +206,12 @@ def chunk_gated_delta_rule( ) """ assert q.dtype == k.dtype == v.dtype - assert ( - q.dtype != torch.float32 - ), "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16." - assert ( - len(beta.shape) == 3 - ), "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise." + assert q.dtype != torch.float32, ( + "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16." + ) + assert len(beta.shape) == 3, ( + "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise." + ) if head_first: raise DeprecationWarning( diff --git a/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py b/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py index 2fe5e623d..d80aa25d1 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py @@ -82,9 +82,10 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( i_v, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -326,9 +327,9 @@ def chunk_gated_delta_rule_fwd_h( chunk_indices: Optional[torch.LongTensor] = None, use_exp2: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - assert not ( - use_exp2 and g is not None - ), "use_exp2 covers only the per-channel gk path; scalar g stays natural-exp" + assert not (use_exp2 and g is not None), ( + "use_exp2 covers only the per-channel gk path; scalar g stays natural-exp" + ) B, T, Hg, K, V = *k.shape, u.shape[-1] H = u.shape[-2] BT = CHUNK_SIZE diff --git a/python/sglang/kernels/ops/attention/fla/chunk_fwd.py b/python/sglang/kernels/ops/attention/fla/chunk_fwd.py index 828ddc4c9..4e7b06810 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_fwd.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_fwd.py @@ -71,12 +71,14 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T diff --git a/python/sglang/kernels/ops/attention/fla/chunk_intra.py b/python/sglang/kernels/ops/attention/fla/chunk_intra.py index 454eba022..908384d9d 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_intra.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_intra.py @@ -88,12 +88,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -819,12 +821,14 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T diff --git a/python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py b/python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py index 7481b9923..feea3c53f 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py @@ -62,9 +62,10 @@ def chunk_kda_fwd_kernel_intra_token_parallel( left = mid + 1 i_n = left - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos i_t = i_tg - bos else: diff --git a/python/sglang/kernels/ops/attention/fla/chunk_o.py b/python/sglang/kernels/ops/attention/fla/chunk_o.py index c2c04312a..f644f15fc 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_o.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_o.py @@ -53,12 +53,14 @@ def chunk_fwd_kernel_o( if IS_VARLEN: i_tg = i_t - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos NT = tl.cdiv(T, BT) else: diff --git a/python/sglang/kernels/ops/attention/fla/cumsum.py b/python/sglang/kernels/ops/attention/fla/cumsum.py index 182211e55..b31331c72 100644 --- a/python/sglang/kernels/ops/attention/fla/cumsum.py +++ b/python/sglang/kernels/ops/attention/fla/cumsum.py @@ -37,12 +37,14 @@ def chunk_local_cumsum_scalar_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -97,12 +99,14 @@ def chunk_local_cumsum_vector_kernel( i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -169,9 +173,9 @@ def chunk_local_cumsum_scalar( B, H, T = g.shape else: B, T, H = g.shape - assert chunk_size == 2 ** ( - chunk_size.bit_length() - 1 - ), "chunk_size must be a power of 2" + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( + "chunk_size must be a power of 2" + ) BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) @@ -216,9 +220,9 @@ def chunk_local_cumsum_vector( if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - assert chunk_size == 2 ** ( - chunk_size.bit_length() - 1 - ), "chunk_size must be a power of 2" + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( + "chunk_size must be a power of 2" + ) g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) @@ -260,9 +264,9 @@ def chunk_local_cumsum( **kwargs, ) -> torch.Tensor: if cu_seqlens is not None: - assert ( - g.shape[0] == 1 - ), "Only batch size 1 is supported when cu_seqlens are provided" + assert g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) if len(g.shape) == 3: return chunk_local_cumsum_scalar( g=g, diff --git a/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py b/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py index 6514db459..cdc3df1cc 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py +++ b/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py @@ -390,9 +390,9 @@ class FusedRMSNormGated(nn.Module): residual_in_fp32: bool = False, ) -> torch.Tensor: if _use_cpu: - assert ( - self.activation == "silu" - ), "CPU rmsnorm_gated currently only supports activation silu" + assert self.activation == "silu", ( + "CPU rmsnorm_gated currently only supports activation silu" + ) return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu( x, self.weight, g, self.eps ) diff --git a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py b/python/sglang/kernels/ops/attention/fla/fused_recurrent.py index 457535586..1a59ad064 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py +++ b/python/sglang/kernels/ops/attention/fla/fused_recurrent.py @@ -43,9 +43,10 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( i_n, i_hv = i_nh // HV, i_nh % HV i_h = i_hv // (HV // H) if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int64) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) all = T T = eos - bos else: @@ -708,7 +709,6 @@ def fused_recurrent_kda_packed_decode( class FusedRecurrentFunction(torch.autograd.Function): - @staticmethod @input_guard def forward( @@ -907,9 +907,10 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( i_n, i_hv = i_nh // HV, i_nh % HV i_h = i_hv // (HV // H) if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int64) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) all = T T = eos - bos else: @@ -1144,7 +1145,6 @@ def fused_recurrent_gated_delta_rule_update_fwd( class FusedRecurrentUpdateFunction(torch.autograd.Function): - @staticmethod @input_guard def forward( diff --git a/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py b/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py index 0b4ef1a69..aa1e53122 100644 --- a/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py +++ b/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py @@ -678,9 +678,9 @@ def _launch_gdn_spec( num_slots, HV, V, K = checkpoint_state.shape H = k.shape[1] B = query_start_loc.shape[0] - 1 - assert ( - max_cache_len & (max_cache_len - 1) == 0 - ), "circular cache requires power-of-two max_cache_len" + assert max_cache_len & (max_cache_len - 1) == 0, ( + "circular cache requires power-of-two max_cache_len" + ) assert d_cache.shape[2] == max_cache_len BK = triton.next_power_of_2(K) diff --git a/python/sglang/kernels/ops/attention/fla/kda.py b/python/sglang/kernels/ops/attention/fla/kda.py index ad9720c53..605583a31 100644 --- a/python/sglang/kernels/ops/attention/fla/kda.py +++ b/python/sglang/kernels/ops/attention/fla/kda.py @@ -1046,18 +1046,18 @@ def kda_gate_chunk_cumsum( Cumulative-summed gated tensor of shape [B, T, H, K]. """ if cu_seqlens is not None: - assert ( - g.shape[0] == 1 - ), "Only batch size 1 is supported when cu_seqlens are provided" + assert g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) assert len(g.shape) == 4 B, T, H, S = g.shape BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - assert chunk_size == 2 ** ( - chunk_size.bit_length() - 1 - ), "chunk_size must be a power of 2" + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( + "chunk_size must be a power of 2" + ) g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) diff --git a/python/sglang/kernels/ops/attention/fla/l2norm.py b/python/sglang/kernels/ops/attention/fla/l2norm.py index 0527324a8..a55e322c3 100644 --- a/python/sglang/kernels/ops/attention/fla/l2norm.py +++ b/python/sglang/kernels/ops/attention/fla/l2norm.py @@ -120,7 +120,6 @@ def l2norm_fwd( class L2NormFunction(torch.autograd.Function): - @staticmethod @input_guard def forward(ctx, x, eps=1e-6, output_dtype=None): @@ -137,7 +136,6 @@ l2_norm = l2norm class L2Norm(nn.Module): - def __init__(self, eps: float = 1e-6, output_dtype: Optional[torch.dtype] = None): super().__init__() self.eps = eps diff --git a/python/sglang/kernels/ops/attention/fla/layernorm_gated.py b/python/sglang/kernels/ops/attention/fla/layernorm_gated.py index a599e761f..05631f006 100644 --- a/python/sglang/kernels/ops/attention/fla/layernorm_gated.py +++ b/python/sglang/kernels/ops/attention/fla/layernorm_gated.py @@ -345,7 +345,6 @@ def rms_norm_gated( class LayerNormFn(torch.autograd.Function): - @staticmethod def forward( ctx, @@ -389,7 +388,6 @@ def layernorm_fn( class LayerNorm(torch.nn.Module): - def __init__( self, hidden_size, @@ -431,7 +429,6 @@ class LayerNorm(torch.nn.Module): class RMSNorm(torch.nn.Module): - def __init__( self, hidden_size, @@ -465,7 +462,9 @@ class RMSNorm(torch.nn.Module): self.norm_before_gate and self.group_size is None and self.activation == "swish" - ), "CPU rmsnorm_gated currently only supports norm before gate without group size or activation other than swish" + ), ( + "CPU rmsnorm_gated currently only supports norm before gate without group size or activation other than swish" + ) return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu( x, self.weight, z, self.eps ) diff --git a/python/sglang/kernels/ops/attention/fla/utils.py b/python/sglang/kernels/ops/attention/fla/utils.py index 4154a3c52..7c0b89006 100644 --- a/python/sglang/kernels/ops/attention/fla/utils.py +++ b/python/sglang/kernels/ops/attention/fla/utils.py @@ -326,9 +326,9 @@ if torch_release >= (2, 4): return device_torch_lib.device(index) else: - assert ( - device == "cuda" - ), "Only cuda device is supported for PyTorch version < 2.4.0." + assert device == "cuda", ( + "Only cuda device is supported for PyTorch version < 2.4.0." + ) autocast_custom_fwd = device_torch_lib.amp.custom_fwd autocast_custom_bwd = device_torch_lib.amp.custom_bwd diff --git a/python/sglang/kernels/ops/attention/fla/wy_fast.py b/python/sglang/kernels/ops/attention/fla/wy_fast.py index 9956b6a92..762e09030 100644 --- a/python/sglang/kernels/ops/attention/fla/wy_fast.py +++ b/python/sglang/kernels/ops/attention/fla/wy_fast.py @@ -43,12 +43,14 @@ def recompute_w_u_fwd_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T diff --git a/python/sglang/kernels/ops/attention/flash_mla_sm120.py b/python/sglang/kernels/ops/attention/flash_mla_sm120.py index bf8dc234d..906d4aa22 100644 --- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py +++ b/python/sglang/kernels/ops/attention/flash_mla_sm120.py @@ -77,9 +77,7 @@ def _gather_and_dequant(k_cache, indices, page_size): raw_pages = k_cache.as_strided( (num_pages, page_bytes), (page_bytes, 1), - ).view( - torch.uint8 - ) # (num_pages, page_bytes) uint8 + ).view(torch.uint8) # (num_pages, page_bytes) uint8 # Note: float8_e4m3fn and uint8 are both 1 byte, view is safe # Compute byte offsets within each page diff --git a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/chunk_fwd.py b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/chunk_fwd.py index de399d33e..480b3b621 100644 --- a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/chunk_fwd.py +++ b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/chunk_fwd.py @@ -922,8 +922,7 @@ def chunk_kda_fwd( if needs_eqlen_pad: if B != 1 and T % BT != 0: raise NotImplementedError( - f"eqlen with B>1 and T % {BT} != 0 not supported " - f"(got B={B}, T={T})." + f"eqlen with B>1 and T % {BT} != 0 not supported (got B={B}, T={T})." ) T_padded = ((T + CPB_BT - 1) // CPB_BT) * CPB_BT # Pre-allocated padded scratch buffers (per (B,T_padded,H,K,dtype) cache diff --git a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_k4_only_persistent.py b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_k4_only_persistent.py index 289049e7e..526cd6170 100644 --- a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_k4_only_persistent.py +++ b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_k4_only_persistent.py @@ -407,7 +407,6 @@ def k4_persistent_kernel( # ==== WG1: State readout + decay ==== if warpgroup_idx == STATE_WG: - cId_128 = cute.make_identity_tensor((M6, N6)) tCtState_mn = transform_partitioned_tensor_layout(tCtState) @@ -698,7 +697,6 @@ def k4_persistent_kernel( # ==== TMA warp (warp 2) ==== elif warp_idx == TMA_WARP: - cta_layout = cute.make_layout(1) scheduler = GDNTileScheduler.create( @@ -995,7 +993,6 @@ def k4_persistent_kernel( # ==== WG2: W/NV/O readout (SS-mode, no Phase 2) ==== elif warpgroup_idx == READOUT_WG: - tCtW_mn = transform_partitioned_tensor_layout(tCtW) tCtNV_mn = transform_partitioned_tensor_layout(tCtNV) tCtO_mn = transform_partitioned_tensor_layout(tCtO) diff --git a/python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py b/python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py index 1a24cc0a7..bf87c27d6 100644 --- a/python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py +++ b/python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py @@ -107,9 +107,9 @@ def chunk_kda_fwd( Returns the fla-shaped 12-tuple: (o [B,T,H,128] bf16, final_state [N,H,128,128] fp32 or None, then Nones, ..., h, initial_state). """ - assert ( - chunk_size == CHUNK - ), f"kda_prefill supports chunk_size={CHUNK} only, got {chunk_size}" + assert chunk_size == CHUNK, ( + f"kda_prefill supports chunk_size={CHUNK} only, got {chunk_size}" + ) if cp_context is not None or disable_recompute: raise NotImplementedError( "kda_prefill is the inference forward path: cp_context, " @@ -124,9 +124,9 @@ def chunk_kda_fwd( if state_v_first and initial_state is not None: # [V,K]-layout state: pure transpose (K==V==128), exact, ~us/call initial_state = initial_state.transpose(-1, -2).contiguous() - assert ( - q.dim() == 4 and q.shape[-1] == K and v.shape[-1] == K - ), f"expected [B,T,H,{K}] q/k/v, got q={tuple(q.shape)} v={tuple(v.shape)}" + assert q.dim() == 4 and q.shape[-1] == K and v.shape[-1] == K, ( + f"expected [B,T,H,{K}] q/k/v, got q={tuple(q.shape)} v={tuple(v.shape)}" + ) B, T, H, _ = q.shape cu_cpu = None @@ -145,9 +145,9 @@ def chunk_kda_fwd( betaf = beta.reshape(Tt, H).contiguous() if use_gate_in_kernel: - assert ( - A_log is not None and dt_bias is not None - ), "use_gate_in_kernel=True requires A_log and dt_bias" + assert A_log is not None and dt_bias is not None, ( + "use_gate_in_kernel=True requires A_log and dt_bias" + ) assert g.dtype == torch.bfloat16, f"raw gate input must be bf16, got {g.dtype}" gf = g.reshape(Tt, H, K).contiguous() sg = lower_bound is not None # fla: lb presence selects safe-gate diff --git a/python/sglang/kernels/ops/attention/linear/lightning_attn.py b/python/sglang/kernels/ops/attention/linear/lightning_attn.py index d415eefc9..1a968dd88 100644 --- a/python/sglang/kernels/ops/attention/linear/lightning_attn.py +++ b/python/sglang/kernels/ops/attention/linear/lightning_attn.py @@ -392,7 +392,6 @@ def _fwd_none_diag_kernel( class _attention(torch.autograd.Function): - @staticmethod def forward(ctx, q, k, v, s, kv_history): # Forward pass of the lightning attention algorithm diff --git a/python/sglang/kernels/ops/attention/linear/seg_la.py b/python/sglang/kernels/ops/attention/linear/seg_la.py index abda8b529..9f5f6a44c 100644 --- a/python/sglang/kernels/ops/attention/linear/seg_la.py +++ b/python/sglang/kernels/ops/attention/linear/seg_la.py @@ -158,7 +158,6 @@ def seg_la_kernel( state = state * block_decay + tl.dot(k, v) else: - qk = tl.dot(q, k) * softmax_scale decays = tl.exp(decay_scale * (offs_b[:, None] - offs_b[None, :])) decays = tl.where(offs_b[None, :] <= offs_b[:, None], decays, 0.0) diff --git a/python/sglang/kernels/ops/attention/metadata.py b/python/sglang/kernels/ops/attention/metadata.py index 85426d045..88230f5e3 100644 --- a/python/sglang/kernels/ops/attention/metadata.py +++ b/python/sglang/kernels/ops/attention/metadata.py @@ -496,9 +496,9 @@ def draft_extend_set_metadata( row tails keep stale values that attention kernels never read past cache_seqlens, matching the eager replay path's bounded writes. """ - assert ( - page_size > 0 and (page_size & (page_size - 1)) == 0 - ), f"page_size must be a power of two, got {page_size}" + assert page_size > 0 and (page_size & (page_size - 1)) == 0, ( + f"page_size must be a power of two, got {page_size}" + ) batch_size = cache_seqlens_int32.shape[0] max_seq_pages = page_table.shape[1] @@ -588,9 +588,9 @@ def normal_decode_set_metadata( page_table / swa_page_table row is (re)written; the tail keeps stale values across CUDA-graph replays, so consumers must bound reads by cache_seqlens. """ - assert ( - page_size > 0 and (page_size & (page_size - 1)) == 0 - ), f"page_size must be a power of two, got {page_size}" + assert page_size > 0 and (page_size & (page_size - 1)) == 0, ( + f"page_size must be a power of two, got {page_size}" + ) batch_size = cache_seqlens_int32.shape[0] device = seq_lens.device diff --git a/python/sglang/kernels/ops/attention/minimax_qknorm_rope.py b/python/sglang/kernels/ops/attention/minimax_qknorm_rope.py index a67d25928..201186779 100644 --- a/python/sglang/kernels/ops/attention/minimax_qknorm_rope.py +++ b/python/sglang/kernels/ops/attention/minimax_qknorm_rope.py @@ -113,9 +113,9 @@ def minimax_qknorm_rope_grouped( """ groups = [(w, off, cnt) for (w, off, cnt) in groups if cnt > 0] num_groups = len(groups) - assert ( - 1 <= num_groups <= _MAX_GROUPS - ), f"need 1..{_MAX_GROUPS} groups, got {num_groups}" + assert 1 <= num_groups <= _MAX_GROUPS, ( + f"need 1..{_MAX_GROUPS} groups, got {num_groups}" + ) weights: List[torch.Tensor] = [g[0] for g in groups] offsets: List[int] = [int(g[1]) for g in groups] diff --git a/python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py b/python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py index c981585a7..c27cebe74 100644 --- a/python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py +++ b/python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py @@ -331,9 +331,9 @@ def flash_decode_with_gqa_share_sparse( max_slots, num_kv_heads, _ = k_cache.shape assert slot_ids.shape[0] == batch_size and seq_lens.shape[0] == batch_size assert topk_idx.shape[0] == num_kv_heads - assert ( - triton.next_power_of_2(block_size) == block_size - ), f"block_size must be a power of 2, but got {block_size}" + assert triton.next_power_of_2(block_size) == block_size, ( + f"block_size must be a power of 2, but got {block_size}" + ) # assert slot_ids.max() < max_slots, f"get slot_ids {slot_ids}, but kv_cache shape is {kv_cache.shape}" max_kv_len = req_to_token.shape[1] # gqa diff --git a/python/sglang/kernels/ops/attention/minimax_sparse/prefill/flash_with_topk_idx.py b/python/sglang/kernels/ops/attention/minimax_sparse/prefill/flash_with_topk_idx.py index 98e370d22..932a08e0e 100644 --- a/python/sglang/kernels/ops/attention/minimax_sparse/prefill/flash_with_topk_idx.py +++ b/python/sglang/kernels/ops/attention/minimax_sparse/prefill/flash_with_topk_idx.py @@ -495,9 +495,9 @@ def flash_prefill_with_topk_index( assert qk_head_dim <= 256 and v_head_dim <= 256, "head_dim must be less than 256" if sink is not None: assert sink.shape[0] == num_heads and sink.shape[1] == qk_head_dim - assert ( - init_blocks + local_blocks <= topk - ), "init_blocks + local_blocks must be less than topk" + assert init_blocks + local_blocks <= topk, ( + "init_blocks + local_blocks must be less than topk" + ) if sm_scale is None: sm_scale = qk_head_dim**-0.5 # q_scale multiplies every Q-side logit (QK dot and sink), so it folds into diff --git a/python/sglang/kernels/ops/attention/mla_kv_pack_quantize_fp8.py b/python/sglang/kernels/ops/attention/mla_kv_pack_quantize_fp8.py index 2d221af4d..e78729537 100644 --- a/python/sglang/kernels/ops/attention/mla_kv_pack_quantize_fp8.py +++ b/python/sglang/kernels/ops/attention/mla_kv_pack_quantize_fp8.py @@ -189,21 +189,21 @@ def mla_kv_pack_quantize_fp8( torch.bfloat16, torch.float16, ), f"k_nope must be bf16/fp16, got {k_nope.dtype}" - assert ( - k_pe.dtype == k_nope.dtype and v.dtype == k_nope.dtype - ), "k_nope, k_pe, v must share dtype" + assert k_pe.dtype == k_nope.dtype and v.dtype == k_nope.dtype, ( + "k_nope, k_pe, v must share dtype" + ) assert fp8_dtype in (torch.float8_e4m3fn, torch.float8_e5m2) s, num_heads, qk_nope = k_nope.shape qk_rope = k_pe.shape[-1] v_head = v.shape[-1] - assert ( - v.shape[0] == s and v.shape[1] == num_heads - ), f"v shape {tuple(v.shape)} mismatches k_nope {tuple(k_nope.shape)}" - assert ( - k_pe.shape[0] == s - ), f"k_pe first dim {k_pe.shape[0]} mismatches k_nope first dim {s}" + assert v.shape[0] == s and v.shape[1] == num_heads, ( + f"v shape {tuple(v.shape)} mismatches k_nope {tuple(k_nope.shape)}" + ) + assert k_pe.shape[0] == s, ( + f"k_pe first dim {k_pe.shape[0]} mismatches k_nope first dim {s}" + ) assert k_nope.stride(-1) == 1, "k_nope must have stride-1 inner dim" assert v.stride(-1) == 1, "v must have stride-1 inner dim" assert k_pe.stride(-1) == 1, "k_pe must have stride-1 inner dim" diff --git a/python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py b/python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py index 837dacee5..51786fc66 100644 --- a/python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py +++ b/python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py @@ -326,9 +326,9 @@ def _decode_grouped_att_m_fwd_rope( is_neox_style=True, ): if use_rope: - assert ( - k_pe_tokens_out is not None - ), "We must output the k_pe tokens with rope applied if rope fusion enabled." + assert k_pe_tokens_out is not None, ( + "We must output the k_pe tokens with rope applied if rope fusion enabled." + ) BLOCK = 32 diff --git a/python/sglang/kernels/ops/attention/score_mod.py b/python/sglang/kernels/ops/attention/score_mod.py index d39c20fbb..98540169d 100644 --- a/python/sglang/kernels/ops/attention/score_mod.py +++ b/python/sglang/kernels/ops/attention/score_mod.py @@ -30,9 +30,9 @@ import triton.language as tl def unpack_aux_tensors(score_mod, aux_tensors): if score_mod is None: return None, 0, 0, 0 - assert ( - aux_tensors is not None and len(aux_tensors) == 1 - ), "Triton score_mod currently requires exactly one aux tensor" + assert aux_tensors is not None and len(aux_tensors) == 1, ( + "Triton score_mod currently requires exactly one aux tensor" + ) aux0 = aux_tensors[0] assert aux0.dim() == 3 and aux0.stride(2) == 1, ( f"aux_tensors[0] must be 3D with a contiguous last dim, " diff --git a/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py b/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py index bd014c0f8..312febb67 100644 --- a/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py +++ b/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py @@ -303,7 +303,7 @@ def sparse_mla_q8kv8_prefill_fwd( ) if indices.ndim != 3: raise ValueError( - "indices must have shape (s_q, h_kv, topk), " f"got {tuple(indices.shape)}" + f"indices must have shape (s_q, h_kv, topk), got {tuple(indices.shape)}" ) s_q, h_q, d_qk = q.shape @@ -362,8 +362,7 @@ def sparse_mla_q8kv8_prefill_fwd( if indices.shape[:2] != (s_q, h_kv): raise ValueError( - "indices must have shape " - f"({s_q}, {h_kv}, topk), got {tuple(indices.shape)}" + f"indices must have shape ({s_q}, {h_kv}, topk), got {tuple(indices.shape)}" ) if indices.dtype != torch.int32: @@ -385,14 +384,13 @@ def sparse_mla_q8kv8_prefill_fwd( raise ValueError("topk_length must be a CUDA tensor") if topk_length.device != device: raise ValueError( - "topk_length must be on q's device " - f"{device}, got {topk_length.device}" + f"topk_length must be on q's device {device}, got {topk_length.device}" ) if not topk_length.is_contiguous(): raise ValueError("topk_length must be contiguous") if torch.any(topk_length < 0).item() or torch.any(topk_length > topk).item(): raise ValueError( - "topk_length values must satisfy " f"0 <= topk_length <= topk ({topk})" + f"topk_length values must satisfy 0 <= topk_length <= topk ({topk})" ) if d_v != 512: diff --git a/python/sglang/kernels/ops/communication/inkling_all_reduce.py b/python/sglang/kernels/ops/communication/inkling_all_reduce.py index 050d57ef4..e7b7af2e7 100644 --- a/python/sglang/kernels/ops/communication/inkling_all_reduce.py +++ b/python/sglang/kernels/ops/communication/inkling_all_reduce.py @@ -121,9 +121,9 @@ _AR_TUNED_TP8 = { } _AR_TUNED = {4: _AR_TUNED_TP4, 8: _AR_TUNED_TP8} _AR_TUNED_TOKENS = sorted(_AR_TUNED_TP4) # same token grid for every table -assert all( - set(t) == set(_AR_TUNED_TP4) for t in _AR_TUNED.values() -), "all tuned tables must share the same token grid" +assert all(set(t) == set(_AR_TUNED_TP4) for t in _AR_TUNED.values()), ( + "all tuned tables must share the same token grid" +) def select_ar_config(num_tokens: int, world_size: int = 4): diff --git a/python/sglang/kernels/ops/communication/mp.py b/python/sglang/kernels/ops/communication/mp.py index 957ac0647..817d158f1 100644 --- a/python/sglang/kernels/ops/communication/mp.py +++ b/python/sglang/kernels/ops/communication/mp.py @@ -161,8 +161,7 @@ def multigpu_launch( for N in num_gpus: if N <= 1 or N > num_devices: raise ValueError( - f"Invalid number of GPUs requested: {N} " - f"(available: {num_devices})" + f"Invalid number of GPUs requested: {N} (available: {num_devices})" ) os.environ[env_key] = "1" os.environ[pid_key] = str(os.getpid()) diff --git a/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_chunkwise_triton.py b/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_chunkwise_triton.py index cf697d07e..1272fae4a 100644 --- a/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_chunkwise_triton.py +++ b/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_chunkwise_triton.py @@ -1666,14 +1666,14 @@ def cam_scan_bidi_chunkwise( q, k, v: camera-prepared ``(B, H, D, N)`` fp32; beta: ``(B, H, F, S)`` fp32; decay: ``(B, H, F)`` fp32. Returns ``(B, H, D, N)`` fp32. """ - assert ( - q.shape == k.shape == v.shape - ), f"q/k/v shape mismatch: {q.shape} {k.shape} {v.shape}" + assert q.shape == k.shape == v.shape, ( + f"q/k/v shape mismatch: {q.shape} {k.shape} {v.shape}" + ) assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() assert beta.is_contiguous() and decay.is_contiguous() - assert ( - q.dtype == torch.float32 - ), f"cam_scan_bidi_chunkwise requires fp32 q/k/v, got {q.dtype}" + assert q.dtype == torch.float32, ( + f"cam_scan_bidi_chunkwise requires fp32 q/k/v, got {q.dtype}" + ) B, H, D, N = q.shape F = beta.shape[2] diff --git a/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_triton.py b/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_triton.py index 07760f61c..e7873d852 100644 --- a/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_triton.py +++ b/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_triton.py @@ -175,9 +175,9 @@ def fused_qk_inv_rms( qkv: (B, N, 3, H, D) contiguous. Returns (q_inv_rms, k_inv_rms), each (B, N) float32. """ assert qkv.is_contiguous(), "qkv must be contiguous (B, N, 3, H, D)" - assert ( - qkv.dim() == 5 and qkv.shape[2] == 3 - ), f"expected (B, N, 3, H, D), got {tuple(qkv.shape)}" + assert qkv.dim() == 5 and qkv.shape[2] == 3, ( + f"expected (B, N, 3, H, D), got {tuple(qkv.shape)}" + ) B, N, _, H, D = qkv.shape C = H * D q_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device) diff --git a/python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py b/python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py index 38321b4f1..816706a9b 100644 --- a/python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py +++ b/python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py @@ -26,7 +26,7 @@ def _jit_usp_relayout_module(dtype: torch.dtype) -> Module: cuda_wrappers=[ ( "usp_merge_heads", - "usp_relayout::" f"UspMergeHeadsKernel<{args}>::run", + f"usp_relayout::UspMergeHeadsKernel<{args}>::run", ), ], ) diff --git a/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py b/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py index 6741c1a2e..7a9a3b61f 100644 --- a/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py +++ b/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py @@ -413,9 +413,9 @@ def fuse_scale_shift_kernel( num_warps = 2 if block_n == 64 else 4 grid = (rows, triton.cdiv(C, block_n)) num_frames = scale.shape[1] - assert ( - L % num_frames == 0 - ), "seq_len must be divisible by num_frames for 4D scale/shift" + assert L % num_frames == 0, ( + "seq_len must be divisible by num_frames for 4D scale/shift" + ) frame_seqlen = L // num_frames # Compact scale [B, F, 1, C] -> [B*F, C] (per-frame) diff --git a/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py b/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py index cf5e45428..ddf5ffc28 100644 --- a/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py +++ b/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py @@ -51,9 +51,9 @@ def _build_fused_norm_module(D: int, is_rms: bool, has_gate: bool, has_weight: b VEC = _VEC NUM_WAVES = _NUM_WAVES BLOCK = NUM_WAVES * WARP_SIZE - assert ( - D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0 - ), f"FlyDSL fused_residual_norm requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" + assert D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0, ( + f"FlyDSL fused_residual_norm requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" + ) NUM_ITERS = D // (BLOCK * VEC) @flyc.kernel(known_block_size=[BLOCK, 1, 1]) @@ -543,9 +543,9 @@ def _build_norm_scale_shift_module(D: int, is_rms: bool, has_weight: bool): VEC = _VEC NUM_WAVES = _NUM_WAVES BLOCK = NUM_WAVES * WARP_SIZE - assert ( - D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0 - ), f"FlyDSL norm_scale_shift requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" + assert D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0, ( + f"FlyDSL norm_scale_shift requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" + ) NUM_ITERS = D // (BLOCK * VEC) @flyc.kernel(known_block_size=[BLOCK, 1, 1]) diff --git a/python/sglang/kernels/ops/elementwise/elementwise.py b/python/sglang/kernels/ops/elementwise/elementwise.py index 326c89ccf..85ba16c84 100644 --- a/python/sglang/kernels/ops/elementwise/elementwise.py +++ b/python/sglang/kernels/ops/elementwise/elementwise.py @@ -103,9 +103,9 @@ fused_dual_residual_rmsnorm_kernel_autotune = rmsnorm_autotune( def fused_dual_residual_rmsnorm(x, residual, weight1, weight2, eps, autotune=False): assert len(x.shape) == 2 - assert ( - x.shape == residual.shape and x.dtype == residual.dtype - ), f"{x.shape=} {residual.shape=} {x.dtype=} {residual.dtype=}" + assert x.shape == residual.shape and x.dtype == residual.dtype, ( + f"{x.shape=} {residual.shape=} {x.dtype=} {residual.dtype=}" + ) output, mid = torch.empty_like(x), torch.empty_like(x) bs, hidden_dim = x.shape if autotune: @@ -434,9 +434,9 @@ def fused_sigmoid_mul( gate_stride_head = gate.stride(1) else: # Flat path: both tensors have the same shape - assert ( - attn_output.shape == gate.shape - ), "attn_output and gate must have the same shape" + assert attn_output.shape == gate.shape, ( + "attn_output and gate must have the same shape" + ) hidden_dim = attn_output.shape[-1] num_tokens = attn_output.numel() // hidden_dim head_dim = hidden_dim diff --git a/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py b/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py index 0df7e9470..9cd83c793 100644 --- a/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py +++ b/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py @@ -338,9 +338,9 @@ def _dsv3_fused_a_gemm_run(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Te assert mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16 assert K % 1024 == 0, f"K must be a multiple of 1024, got {K}" assert N % TILE_M == 0, f"N must be a multiple of {TILE_M}, got {N}" - assert ( - tuple(mat_b.shape) == (K, N) and mat_b.stride(0) == 1 - ), "mat_b must be [K, N] column-major" + assert tuple(mat_b.shape) == (K, N) and mat_b.stride(0) == 1, ( + "mat_b must be [K, N] column-major" + ) assert 1 <= M <= 16, "num_tokens must be in [1, 16]" assert mat_a.stride(1) == 1, "mat_a must be row-major [M, K]" diff --git a/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py b/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py index 9597f836d..92ec82470 100644 --- a/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py +++ b/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py @@ -229,7 +229,9 @@ def default_tactic(m: int, n: int, k: int) -> SplitKTactic: ab_stages=( _MIN_AB_STAGES if k <= 2 * _CTA_K and m > 8 - else min(max_stages, 6) if k <= 4 * _CTA_K else max_stages + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages ), ) validate_tactic(tactic, m, n, k) diff --git a/python/sglang/kernels/ops/grammar/bitmask_ops.py b/python/sglang/kernels/ops/grammar/bitmask_ops.py index 9a195c006..e439effa1 100644 --- a/python/sglang/kernels/ops/grammar/bitmask_ops.py +++ b/python/sglang/kernels/ops/grammar/bitmask_ops.py @@ -114,9 +114,9 @@ def apply_token_bitmask_inplace_triton( indices = torch.tensor(indices, dtype=torch.int32, device=logits.device) num_rows = indices.shape[0] else: - assert ( - logits_shape[0] == bitmask_shape[0] - ), f"batch size mismatch: logits {logits_shape[0]} vs bitmask {bitmask_shape[0]}" + assert logits_shape[0] == bitmask_shape[0], ( + f"batch size mismatch: logits {logits_shape[0]} vs bitmask {bitmask_shape[0]}" + ) num_rows = logits_shape[0] if NUM_SMS > 0: diff --git a/python/sglang/kernels/ops/kimi_k3/sp_collective.py b/python/sglang/kernels/ops/kimi_k3/sp_collective.py index abb93b474..2d87f84e5 100644 --- a/python/sglang/kernels/ops/kimi_k3/sp_collective.py +++ b/python/sglang/kernels/ops/kimi_k3/sp_collective.py @@ -51,10 +51,7 @@ def _device_name(device: torch.device) -> str: def _table(world_size: int, hidden_size: int, device: torch.device) -> Optional[dict]: path = os.path.join( _CONFIG_DIR, - ( - f"world={world_size},H={hidden_size}," - f"device_name={_device_name(device)}.json" - ), + (f"world={world_size},H={hidden_size},device_name={_device_name(device)}.json"), ) if path not in _TABLES: if os.path.exists(path): diff --git a/python/sglang/kernels/ops/kvcache/cache_ops.py b/python/sglang/kernels/ops/kvcache/cache_ops.py index 6bc1450d5..bc4f71034 100644 --- a/python/sglang/kernels/ops/kvcache/cache_ops.py +++ b/python/sglang/kernels/ops/kvcache/cache_ops.py @@ -100,18 +100,18 @@ def concat_and_cast_mha_k_triton( k_rope: torch.Tensor, ): # The source data type will be implicitly converted to the target data type. - assert ( - len(k.shape) == 3 and len(k_nope.shape) == 3 and len(k_rope.shape) == 3 - ), f"shape should be 3d, but got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" - assert ( - k.shape[0] == k_nope.shape[0] and k.shape[0] == k_rope.shape[0] - ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" - assert ( - k.shape[1] == k_nope.shape[1] and 1 == k_rope.shape[1] - ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" - assert ( - k.shape[-1] == k_nope.shape[-1] + k_rope.shape[-1] - ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + assert len(k.shape) == 3 and len(k_nope.shape) == 3 and len(k_rope.shape) == 3, ( + f"shape should be 3d, but got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + ) + assert k.shape[0] == k_nope.shape[0] and k.shape[0] == k_rope.shape[0], ( + f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + ) + assert k.shape[1] == k_nope.shape[1] and 1 == k_rope.shape[1], ( + f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + ) + assert k.shape[-1] == k_nope.shape[-1] + k_rope.shape[-1], ( + f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + ) nope_dim = k_nope.shape[-1] rope_dim = k_rope.shape[-1] @@ -638,9 +638,9 @@ def absorbed_bmm_concat_cast_q_fp8( assert q_fp8_pad.shape[0] >= num_tokens and q_fp8_pad.shape[1] >= num_heads assert q_fp8_pad.shape[2] == n_dim + rope_dim # tl.arange / tl.dot constraints - assert ( - k_dim % 16 == 0 and 16 <= k_dim <= 256 - ), "K must be a multiple of 16 in [16, 256]" + assert k_dim % 16 == 0 and 16 <= k_dim <= 256, ( + "K must be a multiple of 16 in [16, 256]" + ) assert (rope_dim & (rope_dim - 1)) == 0, "ROPE must be a power of two" assert n_dim % block_n == 0, "N must be a multiple of block_n" assert q_nope.stride(2) == 1 and q_rope.stride(2) == 1 @@ -680,22 +680,22 @@ def absorbed_bmm_concat_cast_q_fp8( # Largest power-of-2 divisor of K, capped at 128 (K % 16 == 0 # makes this >= 16), unless the caller pinned block_k. blk_k = block_k or min(k_dim & -k_dim, 128) - assert ( - k_dim % blk_k == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16 - ), "loop needs BLOCK_K a power-of-2 divisor of K >= 16" + assert k_dim % blk_k == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16, ( + "loop needs BLOCK_K a power-of-2 divisor of K >= 16" + ) k_mode = 1 elif v == "two_dot": blk_k = 1 << (k_dim.bit_length() - 1) # largest power of 2 < K k1 = k_dim - blk_k - assert ( - k1 & (k1 - 1) == 0 and k1 >= 16 - ), "two_dot needs K = pow2 + pow2 with both halves >= 16" + assert k1 & (k1 - 1) == 0 and k1 >= 16, ( + "two_dot needs K = pow2 + pow2 with both halves >= 16" + ) k_mode = 2 elif v == "three_dot": blk_k = k_dim // 3 - assert ( - k_dim % 3 == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16 - ), "three_dot needs K = 3 * pow2 with pow2 >= 16" + assert k_dim % 3 == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16, ( + "three_dot needs K = 3 * pow2 with pow2 >= 16" + ) k_mode = 3 elif v == "pad": blk_k = 1 << k_dim.bit_length() # next power of 2 above K diff --git a/python/sglang/kernels/ops/kvcache/hisparse.py b/python/sglang/kernels/ops/kvcache/hisparse.py index cfebdcbf3..2e2281915 100644 --- a/python/sglang/kernels/ops/kvcache/hisparse.py +++ b/python/sglang/kernels/ops/kvcache/hisparse.py @@ -292,9 +292,9 @@ def _load_cache_to_device_buffer_mla( miss_count: torch.Tensor | None, skip_io: bool, ) -> None: - assert ( - hot_buffer_size >= num_top_k - ), f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" + assert hot_buffer_size >= num_top_k, ( + f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" + ) record_miss_plan = miss_src is not None module = _jit_sparse_module( diff --git a/python/sglang/kernels/ops/kvcache/kv_read_table.py b/python/sglang/kernels/ops/kvcache/kv_read_table.py index e19436ea6..9c3141a5e 100644 --- a/python/sglang/kernels/ops/kvcache/kv_read_table.py +++ b/python/sglang/kernels/ops/kvcache/kv_read_table.py @@ -192,9 +192,9 @@ def build_kv_read_table( region's live prefix is written -- never rebound, never tail-cleared. """ bs = int(req_pool_indices.numel()) - assert ( - out.dtype == torch.int32 - ), f"build_kv_read_table: out must be int32, got {out.dtype}" + assert out.dtype == torch.int32, ( + f"build_kv_read_table: out must be int32, got {out.dtype}" + ) assert out.dim() == 2 and out.shape[0] >= bs and out.shape[1] >= max_pages, ( f"build_kv_read_table: out {tuple(out.shape)} cannot hold " f"(bs={bs}, max_pages={max_pages})" @@ -262,7 +262,7 @@ def build_kv_read_table_packed( """ bs = int(req_pool_indices.numel()) assert out.dtype in (torch.int32, torch.int64), ( - f"build_kv_read_table_packed: out must be int32 or int64, got " f"{out.dtype}" + f"build_kv_read_table_packed: out must be int32 or int64, got {out.dtype}" ) assert out.dim() == 1 and out.numel() >= max_tokens, ( f"build_kv_read_table_packed: out {tuple(out.shape)} cannot hold " diff --git a/python/sglang/kernels/ops/kvcache/rope_cache.py b/python/sglang/kernels/ops/kvcache/rope_cache.py index 49b4b8ca9..49615558b 100644 --- a/python/sglang/kernels/ops/kvcache/rope_cache.py +++ b/python/sglang/kernels/ops/kvcache/rope_cache.py @@ -613,37 +613,37 @@ def fused_qk_rope_reshape_and_cache( value_shuffle_layout = False (t_slot,) = slot_mapping.shape - assert ( - t == tk == tv and t_slot <= tk - ), f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}" - assert ( - block_size == block_size_v - ), f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}" - assert ( - kh == vh == kh_cache == vh_cache - ), "KV head should be identical for k, v, key_cache, and value_cache" - assert ( - t_cache == t_cache_v - ), "Number of tokens should be identical for key_cache, and value_cache" + assert t == tk == tv and t_slot <= tk, ( + f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}" + ) + assert block_size == block_size_v, ( + f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}" + ) + assert kh == vh == kh_cache == vh_cache, ( + "KV head should be identical for k, v, key_cache, and value_cache" + ) + assert t_cache == t_cache_v, ( + "Number of tokens should be identical for key_cache, and value_cache" + ) if flash_layout: - assert ( - d == dk == dv == dk_cache == dv_cache - ), "D dimension should be identical for q, k, and v" + assert d == dk == dv == dk_cache == dv_cache, ( + "D dimension should be identical for q, k, and v" + ) else: - assert ( - d == dk == dv == dkx_cache * x_cache == dv_cache - ), "D dimension should be identical for q, k, and v" + assert d == dk == dv == dkx_cache * x_cache == dv_cache, ( + "D dimension should be identical for q, k, and v" + ) assert x_cache == triton.next_power_of_2(x_cache), "x_size should be power of 2" assert d == triton.next_power_of_2(d), "D dimension should be power of 2" - assert block_size == triton.next_power_of_2( - block_size - ), "block_size should be power of 2" + assert block_size == triton.next_power_of_2(block_size), ( + "block_size should be power of 2" + ) assert qh % kh == 0, "Q heads must be multiple of H heads" d_freq = cos_sin.shape[-1] // 2 - assert (d_freq == d // 2) or ( - d_freq == d - ), "cos/sin last dim should be the same or half of the qk last dim" + assert (d_freq == d // 2) or (d_freq == d), ( + "cos/sin last dim should be the same or half of the qk last dim" + ) reuse_freqs_front_part = d_freq == d // 2 if q_out is None: @@ -654,9 +654,9 @@ def fused_qk_rope_reshape_and_cache( if zeros_out is not None: tz, qhz, dz = zeros_out.shape - assert ( - t == tz and qh == qhz and d == dz - ), f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}" + assert t == tz and qh == qhz and d == dz, ( + f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}" + ) output_zeros = True elif output_zeros: zeros_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device) diff --git a/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py b/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py index 67d23a576..ec54f6fed 100644 --- a/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py +++ b/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py @@ -113,12 +113,12 @@ def build_trtllm_mha_page_table( ``full_to_swa`` is provided, which then also requires ``swa_page_table``. """ has_swa = full_to_swa is not None - assert has_swa == ( - swa_page_table is not None - ), "full_to_swa and swa_page_table must be provided together" - assert ( - _MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0 - ), f"page_size={page_size} must divide _MHA_KV_INDEX_BLOCK_TOKENS={_MHA_KV_INDEX_BLOCK_TOKENS}" + assert has_swa == (swa_page_table is not None), ( + "full_to_swa and swa_page_table must be provided together" + ) + assert _MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0, ( + f"page_size={page_size} must divide _MHA_KV_INDEX_BLOCK_TOKENS={_MHA_KV_INDEX_BLOCK_TOKENS}" + ) bs, num_pages = page_table.shape full_to_swa_numel = full_to_swa.numel() if has_swa else 0 create_trtllm_mha_kv_indices_triton[ diff --git a/python/sglang/kernels/ops/layernorm/__init__.py b/python/sglang/kernels/ops/layernorm/__init__.py index fa17572fa..00f857771 100644 --- a/python/sglang/kernels/ops/layernorm/__init__.py +++ b/python/sglang/kernels/ops/layernorm/__init__.py @@ -385,8 +385,7 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp): "(rocm-triton, sglang.kernels.jit)." ), KernelBackend.TORCH: ( - "Gemma-style fused residual-add + RMS normalization " - "(pure-torch reference)." + "Gemma-style fused residual-add + RMS normalization (pure-torch reference)." ), } diff --git a/python/sglang/kernels/ops/layernorm/mhc.py b/python/sglang/kernels/ops/layernorm/mhc.py index d34cca1b4..4784ebf0d 100644 --- a/python/sglang/kernels/ops/layernorm/mhc.py +++ b/python/sglang/kernels/ops/layernorm/mhc.py @@ -1103,9 +1103,9 @@ def mhc_pre( gemm_out_sqrsum = torch.empty( n_splits, num_tokens, dtype=torch.float32, device=residual.device ) - assert ( - n_splits == 1 - ), "The simple TileLang version gemm_sqrsum doesn't support split-k" + assert n_splits == 1, ( + "The simple TileLang version gemm_sqrsum doesn't support split-k" + ) _mhc_pre_gemm_sqrsum_dispatch()( residual_flat.view(num_tokens, hc_mult * hidden_size), fn_flat, @@ -1119,9 +1119,9 @@ def mhc_pre( if norm_weight is not None: assert norm_eps is not None, "norm_eps required when norm_weight is provided" - assert norm_weight.shape == ( - hidden_size, - ), f"norm_weight shape {tuple(norm_weight.shape)} != (hidden_size={hidden_size},)" + assert norm_weight.shape == (hidden_size,), ( + f"norm_weight shape {tuple(norm_weight.shape)} != (hidden_size={hidden_size},)" + ) norm_weight_bf = ( norm_weight.bfloat16() if norm_weight.dtype != torch.bfloat16 diff --git a/python/sglang/kernels/ops/lplb/cuda_solver.py b/python/sglang/kernels/ops/lplb/cuda_solver.py index 61ab5cfe6..3a658cdf2 100644 --- a/python/sglang/kernels/ops/lplb/cuda_solver.py +++ b/python/sglang/kernels/ops/lplb/cuda_solver.py @@ -250,9 +250,9 @@ def dispatch_probability( if random_vals is None: random_vals = torch.rand(n, dtype=torch.float32, device=topk_ids.device) else: - assert random_vals.shape == ( - n, - ), f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}" + assert random_vals.shape == (n,), ( + f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}" + ) module = _dispatch_module(max_copies, DISPATCH_BLOCK_DIM) module.dispatch_probability(out, flat_ids, log2phy_prob, map32, random_vals) return out.view(original_shape).to(topk_ids.dtype) @@ -299,9 +299,9 @@ def dispatch_probability_torch_reference( n = flat_ids.shape[0] num_logical, max_copies = log2phy_prob.shape assert log2phy_map.shape == (num_logical, max_copies) - assert random_vals.shape == ( - n, - ), f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}" + assert random_vals.shape == (n,), ( + f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}" + ) # Gather per-row probabilities and physical maps. probs = log2phy_prob[flat_ids] # (N, max_copies), float32 diff --git a/python/sglang/kernels/ops/lplb/shmem_budget.py b/python/sglang/kernels/ops/lplb/shmem_budget.py index c91f04e19..99b3d4946 100644 --- a/python/sglang/kernels/ops/lplb/shmem_budget.py +++ b/python/sglang/kernels/ops/lplb/shmem_budget.py @@ -111,8 +111,8 @@ def assert_fits(nc: int, nv: int, gpu: str = "h100") -> None: cap = gpu_budget_bytes(gpu) if used > cap: raise ValueError( - f"fused IPM kernel needs {used/1024:.1f} KiB of shared memory for " - f"NC={nc}, NV={nv}, but {gpu} allows {cap/1024:.1f} KiB/block. " + f"fused IPM kernel needs {used / 1024:.1f} KiB of shared memory for " + f"NC={nc}, NV={nv}, but {gpu} allows {cap / 1024:.1f} KiB/block. " f"Either reduce problem size or switch to a tiled design." ) @@ -145,8 +145,8 @@ def report(nc: int, nv: int, gpu: str = "h100") -> str: status = "FITS" if bd.total_bytes <= cap else "OVER BUDGET" return ( f"[shmem] NC={nc} NV={nv} gpu={gpu} | " - f"A={bd.a_bytes/1024:.1f}K " - f"ata={bd.ata_bytes/1024:.1f}K " - f"rest={(bd.c_bytes+bd.x_bytes+bd.rhs_bytes+bd.d_bytes)/1024:.1f}K | " - f"total={bd.total_bytes/1024:.1f}K / {cap/1024:.1f}K {status}" + f"A={bd.a_bytes / 1024:.1f}K " + f"ata={bd.ata_bytes / 1024:.1f}K " + f"rest={(bd.c_bytes + bd.x_bytes + bd.rhs_bytes + bd.d_bytes) / 1024:.1f}K | " + f"total={bd.total_bytes / 1024:.1f}K / {cap / 1024:.1f}K {status}" ) diff --git a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py index be5cca881..321a15d0d 100644 --- a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py +++ b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py @@ -332,7 +332,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching matrix_w = w_col0 matrix_x = col0 for j in tl.static_range(KERNEL_WIDTH): - if KERNEL_WIDTH == 2: if j == 1: # KERNEL_WIDTH-1: matrix_w = w_col1 @@ -502,9 +501,9 @@ def causal_conv1d_fn( assert padded_batch == cache_indices.size(0) if has_initial_state is not None: assert has_initial_state.size() == (padded_batch,) - assert ( - conv_states is not None - ), "ERROR: `has_initial_state` is used, which needs also `conv_states`" + assert conv_states is not None, ( + "ERROR: `has_initial_state` is used, which needs also `conv_states`" + ) assert weight.stride(1) == 1 assert (dim, width) == weight.shape assert is_channel_last, "Need to run in channel-last layout" @@ -1053,9 +1052,9 @@ def causal_conv1d_update( if validate_data: assert dim == weight.size(0) - assert ( - conv_state.stride(-2) == 1 - ), f"ERROR: expect contiguous along feat-dim of conv_state (currently stride={conv_state.stride()})" + assert conv_state.stride(-2) == 1, ( + f"ERROR: expect contiguous along feat-dim of conv_state (currently stride={conv_state.stride()})" + ) assert state_len >= width - 1 # when above happens, we don't shift-left to keep any records in conv_state assert dim == conv_state.size(1) diff --git a/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py b/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py index 3ed64d131..d58c8798f 100644 --- a/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py @@ -52,8 +52,9 @@ cvt.rs.f16x2.f32 $0, $2, $1, $3; @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) @triton.heuristics( { - "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] - is not None + "HAS_STATE_BATCH_INDICES": lambda args: ( + args["state_batch_indices_ptr"] is not None + ) } ) @triton.heuristics( @@ -61,24 +62,23 @@ cvt.rs.f16x2.f32 $0, $2, $1, $3; ) @triton.heuristics( { - "CACHE_INTERMEDIATE_STATES": lambda args: args["intermediate_states_buffer"] - is not None + "CACHE_INTERMEDIATE_STATES": lambda args: ( + args["intermediate_states_buffer"] is not None + ) } ) @triton.heuristics( { - "HAS_EAGLE_TREE_CUSTOM_ATTN_MASK": lambda args: args[ - "retrieve_parent_token_ptr" - ] - is not None + "HAS_EAGLE_TREE_CUSTOM_ATTN_MASK": lambda args: ( + args["retrieve_parent_token_ptr"] is not None + ) } ) @triton.heuristics( { - "HAS_INTERMEDIATE_STATE_INDICES": lambda args: args[ - "intermediate_state_indices_ptr" - ] - is not None + "HAS_INTERMEDIATE_STATE_INDICES": lambda args: ( + args["intermediate_state_indices_ptr"] is not None + ) } ) @triton.jit(do_not_specialize=["T"]) diff --git a/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_scan.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_scan.py index 52b197139..5fd4fb138 100644 --- a/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_scan.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_scan.py @@ -176,7 +176,6 @@ def _chunk_scan_fwd_kernel( ) # if a seq is changed exactly on boundary or (c_off > 0) # implies a new example (pseudo chunk) ): - # - replace prev_states_ptr with init_states prev_states_ptr = ( initstates_ptr @@ -193,7 +192,6 @@ def _chunk_scan_fwd_kernel( # - handle chunk state limit if HAS_INITSTATES: - # have to split this if otherwise compilation will have problems dA_cs_m_boundary = 0.0 @@ -214,7 +212,6 @@ def _chunk_scan_fwd_kernel( # (logical) chunk indices. if (c_idx == c_idx_n) or c_off > 0: - # get the next offset c_off_n = tl.load( chunk_offsets_ptr + (pid_c + 1), @@ -265,7 +262,6 @@ def _chunk_scan_fwd_kernel( + offs_k_dstate[:, None] * prev_states_dstate ) if HAS_SEQ_IDX: - if not HAS_INITSTATES: # - this is for continuous batching where there is no init states scale_m = tl.where(seq_idx_m == seq_idx_prev, tl.exp(dA_cs_m), 0.0) @@ -455,9 +451,9 @@ def _chunk_scan_fwd( # with initial states, we need to take care of how # seq_idx crosses the boundaries assert batch == 1, "chunk scan only supports initial states with batch 1" - assert ( - chunk_indices is not None and chunk_offsets is not None - ), "chunk_indices and chunk_offsets should have been set" + assert chunk_indices is not None and chunk_offsets is not None, ( + "chunk_indices and chunk_offsets should have been set" + ) else: chunk_indices, chunk_offsets = None, None else: diff --git a/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_state.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_state.py index 162d859d4..d33808a51 100644 --- a/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_state.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_state.py @@ -390,7 +390,6 @@ def _chunk_state_varlen_kernel( # - if start_idx < pid_c * chunk_size, then we need to take the past_states_ptrs # - if state_idx >= pid * chunk_size, then we need to insert initstates if (start_idx < pid_c * chunk_size) or (HAS_INITSTATES): # first chunk - dA_cs_boundary = 0.0 # default if not HAS_INITSTATES: @@ -399,7 +398,6 @@ def _chunk_state_varlen_kernel( + offs_n[None, :] * stride_chunk_states_dstate ) else: - # - this seems repetitive, buts its to help the compiler if start_idx < pid_c * chunk_size: past_states_ptrs = chunk_states_ptr + ( diff --git a/python/sglang/kernels/ops/mamba/triton_ops/ssd_combined.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_combined.py index c7f16e70e..b79ff2365 100644 --- a/python/sglang/kernels/ops/mamba/triton_ops/ssd_combined.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/ssd_combined.py @@ -163,9 +163,9 @@ def _mamba_chunk_scan_combined_fwd( if cu_seqlens is None: return out_x, dt, dA_cumsum, states, final_states else: - assert ( - batch == 1 - ), "passing cu_seqlens to get the varlen states is only supported if batch dimension is 1" + assert batch == 1, ( + "passing cu_seqlens to get the varlen states is only supported if batch dimension is 1" + ) varlen_states = chunk_state_varlen( B.squeeze(0), x.squeeze(0), @@ -223,9 +223,9 @@ def mamba_chunk_scan_combined( if not return_varlen_states: cu_seqlens = None else: - assert ( - cu_seqlens is not None - ), "cu_seqlens must be provided if return_varlen_states is True" + assert cu_seqlens is not None, ( + "cu_seqlens must be provided if return_varlen_states is True" + ) out_x, dt_out, dA_cumsum, states, final_states, *rest = ( _mamba_chunk_scan_combined_fwd( x, diff --git a/python/sglang/kernels/ops/mamba/triton_ops/ssd_state_passing.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_state_passing.py index d448a1d5c..c25ea23a1 100644 --- a/python/sglang/kernels/ops/mamba/triton_ops/ssd_state_passing.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/ssd_state_passing.py @@ -189,15 +189,15 @@ def _state_passing_fwd( # - if cu_seqlens is provided, then the initial states # are used for continuous batching. In which case we # require seq_idx to be provided - assert ( - seq_idx is not None - ), "seq_idx must be provided for continuous batching" + assert seq_idx is not None, ( + "seq_idx must be provided for continuous batching" + ) # - we also need chunk_offsets to be provided, to account # for computation of dA_cumsum from the start of the # sequence - assert ( - chunk_offsets is not None - ), "chunk_offsets must be provided for continuous batching" + assert chunk_offsets is not None, ( + "chunk_offsets must be provided for continuous batching" + ) else: # - this is the regular batching case, where initial # states are used are for each example of the batch. diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 20263ae3a..98c6c78be 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/python/sglang/kernels/ops/moe/ep_moe_kernels.py @@ -418,16 +418,16 @@ def silu_and_mul_masked_post_quant_fwd( if output_scale.dtype == torch.int32: assert scale_ue8m0, "packed int32 scales are UE8M0 by definition" - assert ( - num_real_tokens is not None and topk is not None - ), "the packed schedule sizes its grid from num_real_tokens * topk" + assert num_real_tokens is not None and topk is not None, ( + "the packed schedule sizes its grid from num_real_tokens * topk" + ) E, m_max, _ = input.shape G = size_n // quant_group_size assert G % 4 == 0, "packed UE8M0 path requires num_groups % 4 == 0" BLOCK_N = quant_group_size * 4 - assert ( - size_n % BLOCK_N == 0 - ), "packed UE8M0 path requires size_n % (4*group) == 0" + assert size_n % BLOCK_N == 0, ( + "packed UE8M0 path requires size_n % (4*group) == 0" + ) hidden_dim_split = size_n // BLOCK_N assert tuple(output_scale.shape) == (E, hidden_dim_split, m_max) @@ -1175,9 +1175,9 @@ def ep_scatter( is_fp8 = recv_x_scale is not None and recv_x.dtype != torch.bfloat16 if is_fp8: - assert ( - recv_x_scale.dtype == output_tensor_scale.dtype - ), f"recv_x_scale.dtype: {recv_x_scale.dtype}, output_tensor_scale.dtype: {output_tensor_scale.dtype}" + assert recv_x_scale.dtype == output_tensor_scale.dtype, ( + f"recv_x_scale.dtype: {recv_x_scale.dtype}, output_tensor_scale.dtype: {output_tensor_scale.dtype}" + ) assert ( recv_x_scale.shape[1] == output_tensor_scale.shape[1] == scale_hidden_size ) diff --git a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py index 7fa26c7b3..e90e0fb08 100644 --- a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py +++ b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py @@ -866,9 +866,9 @@ def invoke_fused_moe_kernel( assert B_scale is not None if block_shape is None: # activation channel-wise int8 quantization - assert ( - per_channel_quant - ), "int8 quantization only supports channel-wise quantization except for block-wise quantization" + assert per_channel_quant, ( + "int8 quantization only supports channel-wise quantization except for block-wise quantization" + ) A, A_scale = per_token_quant_int8(A) else: # activation block-wise int8 quantization @@ -902,23 +902,23 @@ def invoke_fused_moe_kernel( if fuse_sum_all_reduce: assert not c_sorted, "fuse_sum_all_reduce only supports c_sorted=False" if fuse_add_to_output: - assert ( - not fuse_sum_all_reduce - ), "fuse_add_to_output and fuse_sum_all_reduce are mutually exclusive" - assert ( - add_output_mask is not None - ), "add_output_mask required when fuse_add_to_output=True" + assert not fuse_sum_all_reduce, ( + "fuse_add_to_output and fuse_sum_all_reduce are mutually exclusive" + ) + assert add_output_mask is not None, ( + "add_output_mask required when fuse_add_to_output=True" + ) # ===== TO BE REFACTORED ==== if mask_output: - assert ( - not fuse_add_to_output - ), "mask_output and fuse_add_to_output are mutually exclusive" - assert ( - not fuse_sum_all_reduce - ), "mask_output and fuse_sum_all_reduce are mutually exclusive" - assert ( - add_output_mask is not None - ), "add_output_mask required when mask_output=True" + assert not fuse_add_to_output, ( + "mask_output and fuse_add_to_output are mutually exclusive" + ) + assert not fuse_sum_all_reduce, ( + "mask_output and fuse_sum_all_reduce are mutually exclusive" + ) + assert add_output_mask is not None, ( + "add_output_mask required when mask_output=True" + ) # ===== END TO BE REFACTORED ==== if ( @@ -926,9 +926,9 @@ def invoke_fused_moe_kernel( and block_shape is not None and block_shape[1] > 0 ): - assert ( - not fuse_sum_all_reduce - ), "fuse_sum_all_reduce is not supported for GPTQ/AWQ kernels" + assert not fuse_sum_all_reduce, ( + "fuse_sum_all_reduce is not supported for GPTQ/AWQ kernels" + ) assert B_scale is not None and B_scale.ndim == 3 assert B_zp is None or B_zp.ndim == 3 assert bias is None @@ -1559,9 +1559,9 @@ def fused_append_shared_experts_with_weights( ``apply_sigmoid`` (the sigmoid is intrinsic), so the two are mutually exclusive. """ - assert not ( - fuse_gate and apply_sigmoid - ), "fuse_gate already applies sigmoid in-kernel; do not also set apply_sigmoid" + assert not (fuse_gate and apply_sigmoid), ( + "fuse_gate already applies sigmoid in-kernel; do not also set apply_sigmoid" + ) assert N is not None, "N (shared expert base id) must be provided" m, k = topk_ids.shape s = int(num_fused_shared_experts) @@ -1569,9 +1569,9 @@ def fused_append_shared_experts_with_weights( return topk_ids, topk_weights if fuse_gate: - assert ( - hidden_states is not None and gate_weight is not None - ), "fuse_gate=True requires hidden_states and gate_weight" + assert hidden_states is not None and gate_weight is not None, ( + "fuse_gate=True requires hidden_states and gate_weight" + ) hidden_arg = hidden_states.contiguous() wgate_arg = gate_weight.reshape(-1).contiguous() hidden_dim = hidden_arg.shape[1] diff --git a/python/sglang/kernels/ops/moe/inkling_gate_topk_renorm.py b/python/sglang/kernels/ops/moe/inkling_gate_topk_renorm.py index 5601a2859..452275c85 100644 --- a/python/sglang/kernels/ops/moe/inkling_gate_topk_renorm.py +++ b/python/sglang/kernels/ops/moe/inkling_gate_topk_renorm.py @@ -244,9 +244,9 @@ def _get_fused_scratch(device: torch.device) -> tuple[torch.Tensor, torch.Tensor # buffers in the capture pool, where other graphs' replays can reuse # (clobber) them. Call ensure_gate_gemv_fused_scratch() eagerly first # (InklingGate.__init__ does). - assert ( - not torch.cuda.is_current_stream_capturing() - ), "fused gate scratch must be allocated before CUDA graph capture" + assert not torch.cuda.is_current_stream_capturing(), ( + "fused gate scratch must be allocated before CUDA graph capture" + ) workspace = torch.empty( (_FUSED_MAX_TOKENS, _LOGITS_PAD), dtype=torch.float32, device=device ) diff --git a/python/sglang/kernels/ops/moe/inkling_moe.py b/python/sglang/kernels/ops/moe/inkling_moe.py index 04496548a..8769c917d 100644 --- a/python/sglang/kernels/ops/moe/inkling_moe.py +++ b/python/sglang/kernels/ops/moe/inkling_moe.py @@ -331,14 +331,14 @@ def silu_and_mul_triton( Adapted from ``inkling_kernels.activation.silu_and_mul_fwd`` (without MXFP). """ - assert ( - gateup_output.is_contiguous() - ), f"{gateup_output.shape=} {gateup_output.stride()=}" + assert gateup_output.is_contiguous(), ( + f"{gateup_output.shape=} {gateup_output.stride()=}" + ) assert gateup_output.ndim == 2, f"{gateup_output.shape=}" if topk_weights is not None: - assert ( - topk_weights.is_contiguous() - ), f"{topk_weights.shape=} {topk_weights.stride()=}" + assert topk_weights.is_contiguous(), ( + f"{topk_weights.shape=} {topk_weights.stride()=}" + ) assert topk_weights.ndim == 1, f"{topk_weights.shape=}" M = gateup_output.shape[0] @@ -502,9 +502,9 @@ def compute_expert_block_metadata( block_size_m: int = BLOCK_SIZE_M, ): assert num_tokens_per_expert.ndim == 1, f"{num_tokens_per_expert.shape=}" - assert ( - num_tokens_per_expert.is_contiguous() - ), f"{num_tokens_per_expert.shape=} {num_tokens_per_expert.stride()=}" + assert num_tokens_per_expert.is_contiguous(), ( + f"{num_tokens_per_expert.shape=} {num_tokens_per_expert.stride()=}" + ) num_experts = num_tokens_per_expert.numel() max_num_blocks = _get_max_num_blocks(num_routed_tokens, [block_size_m], num_experts) diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index 640a6610b..b60ecce2e 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/moe_fused_gate.py @@ -52,9 +52,9 @@ def moe_fused_gate_jit( apply_routed_scaling_factor_on_output: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: scoring_func_int = _SCORING_FUNC_MAP.get(scoring_func.lower()) - assert ( - scoring_func_int is not None - ), f"Unknown scoring_func '{scoring_func}', must be one of {list(_SCORING_FUNC_MAP.keys())}" + assert scoring_func_int is not None, ( + f"Unknown scoring_func '{scoring_func}', must be one of {list(_SCORING_FUNC_MAP.keys())}" + ) assert input.dtype == torch.float32, "input must be float32" assert bias.dtype == torch.float32, "bias must be float32" @@ -214,9 +214,7 @@ def _router_triton_kernel( win_lane = tl.min(lane_id, axis=1)[:, None].to(tl.int32) # [BLOCK_M, 1] win_activated = tl.sum( tl.where(offs_n[None, :] == win_lane, activated, 0.0), axis=1 - )[ - :, None - ] # [BLOCK_M, 1] + )[:, None] # [BLOCK_M, 1] slot = offs_k[None, :] == k # [1, BLOCK_K] selected_vals = tl.where(slot, win_activated, selected_vals) selected_idx = tl.where(slot, win_lane, selected_idx) @@ -278,9 +276,9 @@ def moe_fused_gate( the existing call sites. """ scoring_func_int = _SCORING_FUNC_MAP.get(scoring_func.lower()) - assert ( - scoring_func_int is not None - ), f"Unknown scoring_func '{scoring_func}', must be one of {list(_SCORING_FUNC_MAP.keys())}" + assert scoring_func_int is not None, ( + f"Unknown scoring_func '{scoring_func}', must be one of {list(_SCORING_FUNC_MAP.keys())}" + ) assert scores.dtype in ( torch.float32, torch.float16, @@ -288,9 +286,9 @@ def moe_fused_gate( ), "scores must be float32/float16/bfloat16" assert scores.ndim == 2, "scores must be 2D" if bias is None: - assert ( - scoring_func.lower() == "softmax" - ), "bias is required for non-softmax routing" + assert scoring_func.lower() == "softmax", ( + "bias is required for non-softmax routing" + ) else: # The kernel loads the bias and upcasts it to fp32 in-register (see # _router_triton_kernel), so a non-fp32 bias (DeepSeek-V4 stores the @@ -301,9 +299,9 @@ def moe_fused_gate( torch.bfloat16, ), "bias must be float32/float16/bfloat16" assert bias.ndim == 1, "bias must be 1D" - assert scores.size(1) == bias.size( - 0 - ), "scores and bias must have same num_experts" + assert scores.size(1) == bias.size(0), ( + "scores and bias must have same num_experts" + ) assert topk > num_fused_shared_experts, "topk must be > num_fused_shared_experts" if routed_scaling_factor is None: routed_scaling_factor = 1.0 diff --git a/python/sglang/kernels/ops/moe/pack_topk_ids.py b/python/sglang/kernels/ops/moe/pack_topk_ids.py index 3c548c952..437aef580 100644 --- a/python/sglang/kernels/ops/moe/pack_topk_ids.py +++ b/python/sglang/kernels/ops/moe/pack_topk_ids.py @@ -15,7 +15,6 @@ from sglang.kernels.jit.utils import is_arch_support_pdl class PackTopkIds: - @classmethod def execute( cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor @@ -33,18 +32,18 @@ class PackTopkIds: @classmethod def triton(cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor) -> torch.Tensor: - assert ( - topk_ids.shape == topk_weights.shape - ), f"shape mismatch: {topk_ids.shape=} vs {topk_weights.shape=}" + assert topk_ids.shape == topk_weights.shape, ( + f"shape mismatch: {topk_ids.shape=} vs {topk_weights.shape=}" + ) assert topk_ids.ndim >= 1, f"expected >=1D, got {topk_ids.shape=}" assert topk_ids.dtype in ( torch.int32, torch.int64, ), f"topk_ids must be int32 or int64, got {topk_ids.dtype}" - assert ( - topk_weights.dtype == torch.float32 - ), f"topk_weights must be float32, got {topk_weights.dtype}" + assert topk_weights.dtype == torch.float32, ( + f"topk_weights must be float32, got {topk_weights.dtype}" + ) assert topk_ids.is_contiguous(), "topk_ids must be contiguous" assert topk_weights.is_contiguous(), "topk_weights must be contiguous" diff --git a/python/sglang/kernels/ops/moe/rocm_moe_utils.py b/python/sglang/kernels/ops/moe/rocm_moe_utils.py index 43382610f..92f60e1cc 100644 --- a/python/sglang/kernels/ops/moe/rocm_moe_utils.py +++ b/python/sglang/kernels/ops/moe/rocm_moe_utils.py @@ -92,11 +92,11 @@ def rocm_fused_experts_tkw1( # AITER tkw1 kernel for FP8 models with `apply_router_weight_on_input` # This applies topk_weights on the GEMM output of the first FC layer # rather than the second FC. - assert ( - topk_weights.dim() == 2 - ), "`topk_weights` should be in shape (num_tokens, topk)" + assert topk_weights.dim() == 2, ( + "`topk_weights` should be in shape (num_tokens, topk)" + ) assert topk_weights.shape[-1] == 1, ( - "Only support topk=1 when" " `apply_router_weight_on_input` is True" + "Only support topk=1 when `apply_router_weight_on_input` is True" ) return rocm_aiter_asm_moe_tkw1( @@ -307,7 +307,9 @@ def upscale_mxfp4(hidden_state, hidden_state_scale, recv_token_num, output_dtype OUT_TL = ( tl.float16 if output_dtype == torch.float16 - else tl.bfloat16 if output_dtype == torch.bfloat16 else tl.float32 + else tl.bfloat16 + if output_dtype == torch.bfloat16 + else tl.float32 ) upscale_fp4x2_block32_kernel[grid]( diff --git a/python/sglang/kernels/ops/moe/shuffle_rows_with_scales.py b/python/sglang/kernels/ops/moe/shuffle_rows_with_scales.py index 5bc616699..fb5abc7e0 100644 --- a/python/sglang/kernels/ops/moe/shuffle_rows_with_scales.py +++ b/python/sglang/kernels/ops/moe/shuffle_rows_with_scales.py @@ -95,12 +95,12 @@ def shuffle_rows_with_scales( assert q.dim() == 2 and scale.dim() == 2, "q and scale must be 2D" assert q.is_contiguous() and scale.is_contiguous(), "q and scale must be contiguous" assert q.element_size() == 1, f"q must be a 1-byte dtype, got {q.dtype}" - assert ( - q.shape[0] == scale.shape[0] - ), f"row count mismatch: q {q.shape[0]} vs scale {scale.shape[0]}" - assert ( - dst2src_map.numel() >= num_dst_rows - ), f"map holds {dst2src_map.numel()} rows, need {num_dst_rows}" + assert q.shape[0] == scale.shape[0], ( + f"row count mismatch: q {q.shape[0]} vs scale {scale.shape[0]}" + ) + assert dst2src_map.numel() >= num_dst_rows, ( + f"map holds {dst2src_map.numel()} rows, need {num_dst_rows}" + ) # The kernel reads the map as whatever dtype it carries and casts to int64, # so a float map would truncate into a plausible-looking row id instead of # failing. diff --git a/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py index 9221086e9..0f55f5047 100644 --- a/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py +++ b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py @@ -164,16 +164,16 @@ def sigmoid_gate_topk_renorm( # Only column-stride-1 is required (the kernel reads rows via stride_lm). In # InklingGate the gate logits are a [t,258] slice of a padded [t,264] tensor, so # they are NOT contiguous but are column-contiguous -- no copy needed. - assert ( - logits.ndim == 2 and logits.stride(1) == 1 - ), f"{logits.shape=} {logits.stride()=}" - assert ( - logits.shape[0] * logits.stride(0) <= 2**31 - ), f"assumes int32 indexing: {logits.stride()=}" + assert logits.ndim == 2 and logits.stride(1) == 1, ( + f"{logits.shape=} {logits.stride()=}" + ) + assert logits.shape[0] * logits.stride(0) <= 2**31, ( + f"assumes int32 indexing: {logits.stride()=}" + ) assert k <= 32, f"topk kernels only support k <= 32: {k=}" - assert ( - n_shared_experts >= 0 - ), f"expected non-negative shared experts: {n_shared_experts=}" + assert n_shared_experts >= 0, ( + f"expected non-negative shared experts: {n_shared_experts=}" + ) M, G = logits.shape N = G - n_shared_experts A = k + n_shared_experts diff --git a/python/sglang/kernels/ops/moe/trtllm_lora_temp/kimi_k2_moe_fused_gate.py b/python/sglang/kernels/ops/moe/trtllm_lora_temp/kimi_k2_moe_fused_gate.py index a78bdd18d..7ad2476f4 100644 --- a/python/sglang/kernels/ops/moe/trtllm_lora_temp/kimi_k2_moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/trtllm_lora_temp/kimi_k2_moe_fused_gate.py @@ -39,12 +39,12 @@ def kimi_k2_moe_fused_gate( (output_weights, expert_indices). """ _supported = (torch.float32, torch.bfloat16, torch.float16) - assert ( - input.dtype in _supported - ), f"input must be float32/bfloat16/float16, got {input.dtype}" - assert ( - bias.dtype in _supported - ), f"bias must be float32/bfloat16/float16, got {bias.dtype}" + assert input.dtype in _supported, ( + f"input must be float32/bfloat16/float16, got {input.dtype}" + ) + assert bias.dtype in _supported, ( + f"bias must be float32/bfloat16/float16, got {bias.dtype}" + ) assert input.ndim == 2, "input must be 2D" assert bias.ndim == 1, "bias must be 1D" assert input.size(1) == bias.size(0), "input and bias must have same num_experts" diff --git a/python/sglang/kernels/ops/quantization/fp8_kernel.py b/python/sglang/kernels/ops/quantization/fp8_kernel.py index 6c8c805f0..5cfe48e09 100644 --- a/python/sglang/kernels/ops/quantization/fp8_kernel.py +++ b/python/sglang/kernels/ops/quantization/fp8_kernel.py @@ -269,9 +269,9 @@ def _per_token_group_quant_8bit_raw( Returns: Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization. """ - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` cannot be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" if _is_hip: @@ -585,9 +585,9 @@ def _run_per_token_group_quant_8bit_kernel( ) return - assert ( - eps == 1e-10 - ), f"per_token_group_quant bakes the absmax floor in at 1e-10, got {eps}" + assert eps == 1e-10, ( + f"per_token_group_quant bakes the absmax floor in at 1e-10, got {eps}" + ) expected_range = (-448.0, 448.0) if x_q.dtype == fp8_dtype else (-128.0, 127.0) assert (fp8_min, fp8_max) == expected_range, ( f"per_token_group_quant bakes the {x_q.dtype} quant range in at {expected_range}, " @@ -614,9 +614,9 @@ def sglang_per_token_group_quant_fp8( fuse_silu_and_mul: bool = False, masked_m: Optional[torch.Tensor] = None, ): - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` cannot be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" if ( @@ -676,9 +676,9 @@ def sglang_per_token_group_quant_fp8_row_padded( bit-exact; the caller still slices the GEMM output back to m. """ assert x.dim() == 2, "row-padded quant expects a 2D input" - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` must be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` must be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" supported_group_sizes = ( @@ -726,9 +726,9 @@ def sglang_per_token_group_quant_fp8_ue8m0( group_size: int, eps: float = 1e-10, ) -> Tuple[torch.Tensor, torch.Tensor]: - assert ( - x.shape[-1] % group_size == 0 - ), f"hidden ({x.shape[-1]}) must be divisible by group_size ({group_size})" + assert x.shape[-1] % group_size == 0, ( + f"hidden ({x.shape[-1]}) must be divisible by group_size ({group_size})" + ) assert x.is_contiguous(), "x must be contiguous" *x_batch, x_q_mn, x_q_k = x.shape @@ -1466,9 +1466,9 @@ def prepare_block_fp8_matmul_inputs( if As.dtype == torch.float: assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] elif As.dtype == torch.int: - assert ( - triton.cdiv(triton.cdiv(A.shape[-1], block_k), 4) == As.shape[-1] - ), f"{A.shape=} {As.shape=} {block_size=}" + assert triton.cdiv(triton.cdiv(A.shape[-1], block_k), 4) == As.shape[-1], ( + f"{A.shape=} {As.shape=} {block_size=}" + ) else: raise NotImplementedError @@ -1484,9 +1484,9 @@ def prepare_block_fp8_matmul_inputs( assert triton.cdiv(K, block_k) == Bs.shape[1] elif Bs.dtype == torch.int: assert N == Bs.shape[0], f"{B.shape=} {Bs.shape=} {block_size=}" - assert ( - triton.cdiv(triton.cdiv(K, block_k), 4) == Bs.shape[1] - ), f"{B.shape=} {Bs.shape=} {block_size=}" + assert triton.cdiv(triton.cdiv(K, block_k), 4) == Bs.shape[1], ( + f"{B.shape=} {Bs.shape=} {block_size=}" + ) else: raise NotImplementedError @@ -1932,9 +1932,9 @@ if _is_hip: _native_dynamic_per_tensor_quant_fp8(output, input, scale) else: # Static scaling - assert ( - scale.numel() == 1 - ), f"Expected scalar scale, got numel={scale.numel()}" + assert scale.numel() == 1, ( + f"Expected scalar scale, got numel={scale.numel()}" + ) if _use_aiter: static_per_tensor_quant(output, input, scale) elif _has_vllm: @@ -1973,9 +1973,9 @@ else: ) # False for dynamic else: # Static scaling - assert ( - scale.numel() == 1 - ), f"Expected scalar scale, got numel={scale.numel()}" + assert scale.numel() == 1, ( + f"Expected scalar scale, got numel={scale.numel()}" + ) sgl_per_tensor_quant_fp8( input, output, scale, is_static=True ) # True for static @@ -2054,9 +2054,9 @@ def per_token_group_quant_fp8_hopper_moe_mn_major( ) -> Tuple[torch.Tensor, torch.Tensor]: assert A.dim() == 2 assert A.is_contiguous(), "`A` is not contiguous" - assert ( - A.shape[-1] % group_size == 0 - ), "the last dimension of `A` cannot be divisible by `group_size`" + assert A.shape[-1] % group_size == 0, ( + "the last dimension of `A` cannot be divisible by `group_size`" + ) a_q = torch.empty_like(A, device=A.device, dtype=fp8_dtype) M, K = A.shape[0], A.shape[1] diff --git a/python/sglang/kernels/ops/quantization/fp8_quantize.py b/python/sglang/kernels/ops/quantization/fp8_quantize.py index 6eadb78a0..53c4c6636 100644 --- a/python/sglang/kernels/ops/quantization/fp8_quantize.py +++ b/python/sglang/kernels/ops/quantization/fp8_quantize.py @@ -81,7 +81,7 @@ def _flatten_to_2d(x: torch.Tensor): if x.stride(d) != expected: raise ValueError( f"cannot flatten dim {d}: stride={x.stride(d)} but expected " - f"shape[{d+1}]*stride[{d+1}]={expected}. Tensor shape={tuple(x.shape)}, " + f"shape[{d + 1}]*stride[{d + 1}]={expected}. Tensor shape={tuple(x.shape)}, " f"stride={tuple(x.stride())}." ) return M, N, row_stride diff --git a/python/sglang/kernels/ops/quantization/int8_kernel.py b/python/sglang/kernels/ops/quantization/int8_kernel.py index 3c73744f4..2748db0d8 100644 --- a/python/sglang/kernels/ops/quantization/int8_kernel.py +++ b/python/sglang/kernels/ops/quantization/int8_kernel.py @@ -152,9 +152,9 @@ def per_token_group_quant_int8( Returns: Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization. """ - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` cannot be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" iinfo = torch.iinfo(dtype) @@ -197,15 +197,15 @@ def sglang_per_token_group_quant_int8( eps: float = 1e-10, dtype: torch.dtype = torch.int8, ): - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` cannot be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" assert dtype == torch.int8 # per_token_group_quant bakes the int8 constants in ([-128, 127], eps 1e-10). - assert ( - eps == 1e-10 - ), f"per_token_group_quant bakes the absmax floor in at 1e-10, got {eps}" + assert eps == 1e-10, ( + f"per_token_group_quant bakes the absmax floor in at 1e-10, got {eps}" + ) return per_token_group_quant(x, group_size=group_size, out_dtype=dtype) diff --git a/python/sglang/kernels/ops/quantization/mxfp8_amd_gfx95.py b/python/sglang/kernels/ops/quantization/mxfp8_amd_gfx95.py index 9c90ce8ed..773ab1511 100644 --- a/python/sglang/kernels/ops/quantization/mxfp8_amd_gfx95.py +++ b/python/sglang/kernels/ops/quantization/mxfp8_amd_gfx95.py @@ -309,9 +309,9 @@ def dot_scaled_mxfp8_blockscaled_linear( kernel_out_dtype = input_2d.dtype else: # Activations already MXFP8-quantized by a fused upstream op. - assert ( - input_2d.dtype == MXFP8_VALUE_DTYPE - ), "pre-quantized input must be FP8 E4M3 when input_scale is given." + assert input_2d.dtype == MXFP8_VALUE_DTYPE, ( + "pre-quantized input must be FP8 E4M3 when input_scale is given." + ) assert input_scale.dtype == torch.uint8 and input_scale.shape == ( m, k // 32, diff --git a/python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py b/python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py index ede6ce7a2..8a9ce307d 100644 --- a/python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py +++ b/python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py @@ -69,9 +69,9 @@ def store_sf_interleaved( page_size: int = 128, ): """Scatter-write per-token scale factors into interleaved page layout.""" - assert ( - page_size == 128 - ), f"Interleaved SF layout requires page_size=128, got {page_size}" + assert page_size == 128, ( + f"Interleaved SF layout requires page_size=128, got {page_size}" + ) num_tokens, nheads, sf_dim = sf_in.shape assert sf_dim == 4, f"Expected sf_dim=4 (hdim=128, sf_vec_size=32), got {sf_dim}" diff --git a/python/sglang/kernels/ops/sampling/murmur_hash.py b/python/sglang/kernels/ops/sampling/murmur_hash.py index 2a090a4ca..3028136b8 100644 --- a/python/sglang/kernels/ops/sampling/murmur_hash.py +++ b/python/sglang/kernels/ops/sampling/murmur_hash.py @@ -102,12 +102,12 @@ def murmur_hash32_kernel( def murmur_hash32(seed, positions, col_indices): - assert ( - seed.shape == positions.shape - ), "Seed and positions must have the same shape (n,)" - assert ( - len(seed.shape) == 1 and len(col_indices.shape) == 1 - ), f"Inputs must be 1D tensors {seed.shape=} {col_indices.shape=}" + assert seed.shape == positions.shape, ( + "Seed and positions must have the same shape (n,)" + ) + assert len(seed.shape) == 1 and len(col_indices.shape) == 1, ( + f"Inputs must be 1D tensors {seed.shape=} {col_indices.shape=}" + ) n = seed.shape[0] m = col_indices.shape[0] device = seed.device diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py index 948a9d60a..a63748c0d 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py @@ -306,9 +306,9 @@ def softmax_temp( ) -> torch.Tensor: num_rows = logits.shape[0] bs = num_rows // rows_per_request - assert ( - bs * rows_per_request == num_rows - ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + assert bs * rows_per_request == num_rows, ( + f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + ) temp_per_row = torch.repeat_interleave( temperatures.reshape(bs).to(torch.float32), rows_per_request, dim=0 ) @@ -366,9 +366,9 @@ def softmax_temp_triton( ) -> torch.Tensor: num_rows, vocab = logits.shape[0], logits.shape[-1] bs = num_rows // rows_per_request - assert ( - bs * rows_per_request == num_rows - ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + assert bs * rows_per_request == num_rows, ( + f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + ) temperatures = temperatures.reshape(bs).to(torch.float32).contiguous() out = torch.empty((num_rows, vocab), dtype=torch.float32, device=logits.device) BLOCK_V = 4096 @@ -397,9 +397,9 @@ def softmax_temp_flashinfer( ) num_rows, vocab = logits.shape[0], logits.shape[-1] bs = num_rows // rows_per_request - assert ( - bs * rows_per_request == num_rows - ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + assert bs * rows_per_request == num_rows, ( + f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + ) temp_per_row = torch.repeat_interleave( temperatures.reshape(bs).to(torch.float32), rows_per_request, dim=0 ).contiguous() diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py b/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py index 0de315543..f733ba2bd 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py @@ -390,9 +390,9 @@ def compact_row_index_triton( verify_lens = verify_lens.to(device=device, dtype=torch.int64).contiguous() bs = verify_lens.shape[0] # The search converges only for bs <= 2**(NBITS-1); beyond it silently mismaps. - assert bs <= 1 << ( - _SEARCH_NBITS - 1 - ), f"bs={bs} exceeds row-index search capacity {1 << (_SEARCH_NBITS - 1)}" + assert bs <= 1 << (_SEARCH_NBITS - 1), ( + f"bs={bs} exceeds row-index search capacity {1 << (_SEARCH_NBITS - 1)}" + ) incl = torch.cumsum(verify_lens, dim=0).contiguous() req = torch.empty(padded_total, dtype=torch.int64, device=device) within = torch.empty(padded_total, dtype=torch.int64, device=device) diff --git a/python/sglang/lang/backend/openai.py b/python/sglang/lang/backend/openai.py index a2d006bb7..3ee2f4fb1 100644 --- a/python/sglang/lang/backend/openai.py +++ b/python/sglang/lang/backend/openai.py @@ -129,9 +129,9 @@ class OpenAI(BaseBackend): if key not in self.spec_kwargs: self.spec_kwargs[key] = value else: - assert ( - value == self.spec_kwargs[key] - ), "sampling parameters should be consistent if turn on api speculative execution." + assert value == self.spec_kwargs[key], ( + "sampling parameters should be consistent if turn on api speculative execution." + ) self.spec_format.append( {"text": "", "stop": params["stop"], "name": spec_var_name} ) @@ -180,9 +180,9 @@ class OpenAI(BaseBackend): ) # Keep the returned list (or string) as is. elif sampling_params.dtype in [str, "str", "string"]: - assert ( - not self.is_chat_model - ), "constrained type not supported on chat model" + assert not self.is_chat_model, ( + "constrained type not supported on chat model" + ) kwargs = sampling_params.to_openai_kwargs() kwargs.pop("stop") comp = openai_completion( @@ -200,9 +200,9 @@ class OpenAI(BaseBackend): else: comp = '"' + comp + '"' elif sampling_params.dtype in [int, "int"]: - assert ( - not self.is_chat_model - ), "constrained type not supported on chat model" + assert not self.is_chat_model, ( + "constrained type not supported on chat model" + ) kwargs = sampling_params.to_openai_kwargs() kwargs.pop("stop") comp = openai_completion( diff --git a/python/sglang/lang/backend/runtime_endpoint.py b/python/sglang/lang/backend/runtime_endpoint.py index db61e431f..84849d4ea 100644 --- a/python/sglang/lang/backend/runtime_endpoint.py +++ b/python/sglang/lang/backend/runtime_endpoint.py @@ -133,18 +133,14 @@ class RuntimeEndpoint(BaseBackend): dtype_regex = None if sampling_params.dtype in ["int", int]: - dtype_regex = REGEX_INT sampling_params.stop.extend([" ", "\n"]) elif sampling_params.dtype in ["float", float]: - dtype_regex = REGEX_FLOAT sampling_params.stop.extend([" ", "\n"]) elif sampling_params.dtype in ["str", str]: - dtype_regex = REGEX_STR elif sampling_params.dtype in ["bool", bool]: - dtype_regex = REGEX_BOOL else: raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}") diff --git a/python/sglang/lang/choices.py b/python/sglang/lang/choices.py index e52c6b362..ceac01b69 100644 --- a/python/sglang/lang/choices.py +++ b/python/sglang/lang/choices.py @@ -12,7 +12,6 @@ class ChoicesDecision: class ChoicesSamplingMethod(ABC): - @property def requires_unconditional_logprobs(self) -> bool: return False @@ -30,7 +29,6 @@ class ChoicesSamplingMethod(ABC): class TokenLengthNormalized(ChoicesSamplingMethod): - def __call__( self, *, @@ -54,7 +52,6 @@ token_length_normalized = TokenLengthNormalized() class GreedyTokenSelection(ChoicesSamplingMethod): - def __call__( self, *, @@ -108,7 +105,6 @@ greedy_token_selection = GreedyTokenSelection() class UnconditionalLikelihoodNormalized(ChoicesSamplingMethod): - @property def requires_unconditional_logprobs(self) -> bool: return True diff --git a/python/sglang/lang/interpreter.py b/python/sglang/lang/interpreter.py index 90dd41857..c07a84bd2 100644 --- a/python/sglang/lang/interpreter.py +++ b/python/sglang/lang/interpreter.py @@ -624,9 +624,9 @@ class StreamExecutor: self.meta_info[name] = meta_info self.variable_event[name].set() else: - assert ( - self.num_api_spec_tokens is None - ), "stream is not supported with api speculative execution" + assert self.num_api_spec_tokens is None, ( + "stream is not supported with api speculative execution" + ) generator = self.backend.generate_stream( self, sampling_params=sampling_params ) diff --git a/python/sglang/lang/ir.py b/python/sglang/lang/ir.py index 45cb6c859..771e929a4 100644 --- a/python/sglang/lang/ir.py +++ b/python/sglang/lang/ir.py @@ -531,7 +531,6 @@ class SglRoleEnd(SglExpr): class SglSelect(SglExpr): - def __init__( self, name: str, diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py index 6133d5f97..36527369a 100755 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py @@ -2356,8 +2356,7 @@ def main(): args.breakable_cuda_graph or args.quality_bcg_matrix ): parser.error( - "--bcg-text-buckets requires --breakable-cuda-graph or " - "--quality-bcg-matrix" + "--bcg-text-buckets requires --breakable-cuda-graph or --quality-bcg-matrix" ) if args.cleanup_model_cache and not args.model_cache_root: parser.error("--cleanup-model-cache requires --model-cache-root") diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_flux_pipeline.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_flux_pipeline.py index 63c099682..59147a6c3 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_flux_pipeline.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_flux_pipeline.py @@ -135,12 +135,12 @@ def test_comfyui_flux_pipeline_direct() -> None: assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" - assert ( - noise_pred.device.type == "cuda" - ), f"noise_pred should be on cuda, got {noise_pred.device}" - assert ( - noise_pred.dtype == torch.bfloat16 - ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" + assert noise_pred.device.type == "cuda", ( + f"noise_pred should be on cuda, got {noise_pred.device}" + ) + assert noise_pred.dtype == torch.bfloat16, ( + f"noise_pred should be bfloat16, got {noise_pred.dtype}" + ) print("โ Successfully retrieved noise_pred from OutputBatch!") print(f" noise_pred shape: {noise_pred.shape}") diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_h3_request.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_h3_request.py index b46cf3b5e..6b5b8e231 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_h3_request.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_h3_request.py @@ -145,15 +145,19 @@ def _run_node(**node_kwargs): ) node = SGLDiffusionGenerateH3() - with mock.patch( - f"{PKG}.core.server_api.requests.post", - side_effect=fake_post, - ), mock.patch( - f"{PKG}.core.server_api.requests.get", - side_effect=fake_get, - ), mock.patch( - f"{PKG}.nodes.get_image_path", - side_effect=lambda image: "/tmp/frame.png", + with ( + mock.patch( + f"{PKG}.core.server_api.requests.post", + side_effect=fake_post, + ), + mock.patch( + f"{PKG}.core.server_api.requests.get", + side_effect=fake_get, + ), + mock.patch( + f"{PKG}.nodes.get_image_path", + side_effect=lambda image: "/tmp/frame.png", + ), ): result = node.generate(sgld_client=client, **node_kwargs) return captured, result @@ -243,13 +247,16 @@ def test_extra_fields_win_over_generic_defaults(): captured.update(json) return _Response({"id": "job-1"}) - with mock.patch( - f"{PKG}.core.server_api.requests.post", - side_effect=fake_post, - ), mock.patch( - f"{PKG}.core.server_api.requests.get", - side_effect=lambda *a, **k: _Response( - {"id": "job-1", "status": "completed", "size": RESOLVED_SIZE} + with ( + mock.patch( + f"{PKG}.core.server_api.requests.post", + side_effect=fake_post, + ), + mock.patch( + f"{PKG}.core.server_api.requests.get", + side_effect=lambda *a, **k: _Response( + {"id": "job-1", "status": "completed", "size": RESOLVED_SIZE} + ), ), ): client.generate_video( diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_edit_pipeline.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_edit_pipeline.py index 609f8df04..422441d56 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_edit_pipeline.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_edit_pipeline.py @@ -116,12 +116,12 @@ def test_comfyui_qwen_image_edit_pipeline_direct() -> None: assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" - assert ( - noise_pred.device.type == "cuda" - ), f"noise_pred should be on cuda, got {noise_pred.device}" - assert ( - noise_pred.dtype == torch.bfloat16 - ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" + assert noise_pred.device.type == "cuda", ( + f"noise_pred should be on cuda, got {noise_pred.device}" + ) + assert noise_pred.dtype == torch.bfloat16, ( + f"noise_pred should be bfloat16, got {noise_pred.dtype}" + ) print("โ Successfully retrieved noise_pred from OutputBatch (Edit Mode)!") print(f" noise_pred shape: {noise_pred.shape}") diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_pipeline.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_pipeline.py index 88d73f663..bb7b070f0 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_pipeline.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_pipeline.py @@ -100,12 +100,12 @@ def test_comfyui_qwen_image_pipeline_direct() -> None: assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" - assert ( - noise_pred.device.type == "cuda" - ), f"noise_pred should be on cuda, got {noise_pred.device}" - assert ( - noise_pred.dtype == torch.bfloat16 - ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" + assert noise_pred.device.type == "cuda", ( + f"noise_pred should be on cuda, got {noise_pred.device}" + ) + assert noise_pred.dtype == torch.bfloat16, ( + f"noise_pred should be bfloat16, got {noise_pred.dtype}" + ) print("โ Successfully retrieved noise_pred from OutputBatch!") print(f" noise_pred shape: {noise_pred.shape}") diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py index 77e6811bf..4b053f1db 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py @@ -102,12 +102,12 @@ def test_comfyui_zimage_pipeline_direct() -> None: assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" - assert ( - noise_pred.device.type == "cuda" - ), f"noise_pred should be on cuda, got {noise_pred.device}" - assert ( - noise_pred.dtype == torch.bfloat16 - ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" + assert noise_pred.device.type == "cuda", ( + f"noise_pred should be on cuda, got {noise_pred.device}" + ) + assert noise_pred.dtype == torch.bfloat16, ( + f"noise_pred should be bfloat16, got {noise_pred.dtype}" + ) print("โ Successfully retrieved noise_pred from OutputBatch!") print(f" noise_pred shape: {noise_pred.shape}") diff --git a/python/sglang/multimodal_gen/benchmarks/bench_serving.py b/python/sglang/multimodal_gen/benchmarks/bench_serving.py index dca9a4c28..b5a1dffee 100644 --- a/python/sglang/multimodal_gen/benchmarks/bench_serving.py +++ b/python/sglang/multimodal_gen/benchmarks/bench_serving.py @@ -631,7 +631,7 @@ async def benchmark(args): warm_out = await limited_request_func(warm_req, session, None) warmup_pairs.append((warm_req, warm_out)) logger.info( - f"Warmup {i+1}/{args.warmup_requests}: " + f"Warmup {i + 1}/{args.warmup_requests}: " f"latency={warm_out.latency:.2f}s, success={warm_out.success}" ) diff --git a/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py b/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py index 6caefb45e..03b18fe08 100644 --- a/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py +++ b/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py @@ -67,7 +67,6 @@ class LTX2ConnectorArchConfig(AdapterArchConfig): @dataclass class LTX2ConnectorConfig(AdapterConfig): - arch_config: AdapterArchConfig = field(default_factory=LTX2ConnectorArchConfig) prefix: str = "LTX2" diff --git a/python/sglang/multimodal_gen/configs/models/base.py b/python/sglang/multimodal_gen/configs/models/base.py index 3bca5cc67..a789b25af 100644 --- a/python/sglang/multimodal_gen/configs/models/base.py +++ b/python/sglang/multimodal_gen/configs/models/base.py @@ -83,9 +83,9 @@ class ModelConfig: arch_config.__post_init__() def update_model_config(self, source_model_dict: dict[str, Any]) -> None: - assert ( - "arch_config" not in source_model_dict - ), "Source model config shouldn't contain arch_config." + assert "arch_config" not in source_model_dict, ( + "Source model config shouldn't contain arch_config." + ) valid_fields = {f.name for f in fields(self)} diff --git a/python/sglang/multimodal_gen/configs/models/dits/flux.py b/python/sglang/multimodal_gen/configs/models/dits/flux.py index 97adc01e8..d1a8a6885 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/flux.py +++ b/python/sglang/multimodal_gen/configs/models/dits/flux.py @@ -110,7 +110,6 @@ class FluxArchConfig(DiTArchConfig): @dataclass class FluxConfig(DiTConfig): - arch_config: DiTArchConfig = field(default_factory=FluxArchConfig) prefix: str = "Flux" diff --git a/python/sglang/multimodal_gen/configs/models/dits/krea2.py b/python/sglang/multimodal_gen/configs/models/dits/krea2.py index 99a83f97e..b330b5364 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/krea2.py +++ b/python/sglang/multimodal_gen/configs/models/dits/krea2.py @@ -57,9 +57,9 @@ class Krea2ArchConfig(DiTArchConfig): self.num_attention_heads = self.heads self.num_channels_latents = self.channels assert self.features % self.heads == 0 - assert ( - sum(self.axes_dims) == self.features // self.heads - ), f"sum(axes_dims)={sum(self.axes_dims)} != head_dim={self.features // self.heads}" + assert sum(self.axes_dims) == self.features // self.heads, ( + f"sum(axes_dims)={sum(self.axes_dims)} != head_dim={self.features // self.heads}" + ) @property def head_dim(self) -> int: diff --git a/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py b/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py index ad2ee5d93..637025a4f 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py +++ b/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py @@ -50,9 +50,9 @@ class MOVAAudioArchConfig(DiTArchConfig): self.hidden_size = self.dim self.num_attention_heads = self.num_heads self.num_channels_latents = self.out_dim - assert ( - not self.has_image_input - ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." + assert not self.has_image_input, ( + "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." + ) @dataclass diff --git a/python/sglang/multimodal_gen/configs/models/dits/mova_video.py b/python/sglang/multimodal_gen/configs/models/dits/mova_video.py index b0c2e882f..e3f010e88 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/mova_video.py +++ b/python/sglang/multimodal_gen/configs/models/dits/mova_video.py @@ -49,9 +49,9 @@ class MOVAVideoArchConfig(DiTArchConfig): self.hidden_size = self.dim self.num_attention_heads = self.num_heads self.num_channels_latents = self.out_dim - assert ( - not self.has_image_input - ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." + assert not self.has_image_input, ( + "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." + ) @dataclass diff --git a/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py b/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py index 85a3acda8..daad209ee 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py @@ -100,9 +100,7 @@ class WanVideoArchConfig(DiTArchConfig): local_attn_size: int = ( -1 ) # Window size for temporal local attention (-1 indicates global attention) - sink_size: int = ( - 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache - ) + sink_size: int = 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache num_frames_per_block: int = 3 sliding_window_num_frames: int = 21 attention_type: str = "original" diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index a15e1a5cf..e86ad7bde 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -1180,9 +1180,9 @@ class PipelineConfig: elif isinstance(current_value, tuple) and all( isinstance(v, ModelConfig) for v in current_value ): - assert len(current_value) == len( - new_value - ), "Users shouldn't delete or add text encoder config objects in your json" + assert len(current_value) == len(new_value), ( + "Users shouldn't delete or add text encoder config objects in your json" + ) for target_config, source_config in zip( current_value, new_value, strict=True ): diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py index 209a484ab..feccd04a2 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py @@ -17,7 +17,6 @@ logger = init_logger(__name__) @dataclass class LongLive2T2VConfig(Wan2_2_TI2V_5B_Config): - is_causal: bool = True task_type: ModelTaskType = ModelTaskType.TI2V vae_precision: str = "bf16" diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py index e564947d5..c3967b40e 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py @@ -41,7 +41,6 @@ def sana_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Ten @dataclass class SanaPipelineConfig(SpatialImagePipelineConfig): - task_type: ModelTaskType = ModelTaskType.T2I # should_use_guidance=False disables *embedded* guidance (timestep-conditioned diff --git a/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py index fa5323f47..a3c1bd62e 100644 --- a/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py +++ b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py @@ -12,7 +12,6 @@ from sglang.multimodal_gen.runtime.post_training.sp_utils import ( class QwenImageRolloutPipelineMixin: - def gather_denoising_env_static_for_sp(self, batch, cond_kwargs: dict | None): if cond_kwargs is None: return None diff --git a/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py index 5c4e60ad5..f49e63ac3 100644 --- a/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py +++ b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py @@ -12,7 +12,6 @@ from sglang.multimodal_gen.runtime.post_training.sp_utils import ( class ZImageRolloutPipelineMixin: - def gather_denoising_env_static_for_sp(self, batch, cond_kwargs: dict | None): if cond_kwargs is None: return None diff --git a/python/sglang/multimodal_gen/configs/sample/cosmos3.py b/python/sglang/multimodal_gen/configs/sample/cosmos3.py index e75e9e71b..75512d53f 100644 --- a/python/sglang/multimodal_gen/configs/sample/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/sample/cosmos3.py @@ -501,8 +501,7 @@ class Cosmos3SamplingParams(SamplingParams): raise ValueError("num_conditional_frames must be non-negative") if self.num_conditional_frames >= self.num_video_frames_per_chunk: raise ValueError( - "num_conditional_frames must be smaller than " - "num_video_frames_per_chunk" + "num_conditional_frames must be smaller than num_video_frames_per_chunk" ) if self.num_first_chunk_conditional_frames < 0: raise ValueError("num_first_chunk_conditional_frames must be non-negative") diff --git a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py index 85934ec90..4c4d4546b 100644 --- a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py +++ b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py @@ -267,7 +267,9 @@ class MiniMaxH3SamplingParams(SamplingParams): seed=( _seed_override if _seed_override is not None - else self.seed if isinstance(self.seed, int) else None + else self.seed + if isinstance(self.seed, int) + else None ), ) ) diff --git a/python/sglang/multimodal_gen/configs/sample/qwenimage.py b/python/sglang/multimodal_gen/configs/sample/qwenimage.py index 127dfa5fc..8e0b09033 100644 --- a/python/sglang/multimodal_gen/configs/sample/qwenimage.py +++ b/python/sglang/multimodal_gen/configs/sample/qwenimage.py @@ -17,9 +17,7 @@ class QwenImageSamplingParams(SamplingParams): @dataclass class QwenImage2512SamplingParams(QwenImageSamplingParams): - negative_prompt: str = ( - "ไฝๅ่พจ็๏ผไฝ็ป่ดจ๏ผ่ขไฝ็ธๅฝข๏ผๆๆ็ธๅฝข๏ผ็ป้ข่ฟ้ฅฑๅ๏ผ่กๅๆ๏ผไบบ่ธๆ ็ป่๏ผ่ฟๅบฆๅ ๆป๏ผ็ป้ขๅ ทๆAIๆใๆๅพๆททไนฑใๆๅญๆจก็ณ๏ผๆญๆฒใ" - ) + negative_prompt: str = "ไฝๅ่พจ็๏ผไฝ็ป่ดจ๏ผ่ขไฝ็ธๅฝข๏ผๆๆ็ธๅฝข๏ผ็ป้ข่ฟ้ฅฑๅ๏ผ่กๅๆ๏ผไบบ่ธๆ ็ป่๏ผ่ฟๅบฆๅ ๆป๏ผ็ป้ขๅ ทๆAIๆใๆๅพๆททไนฑใๆๅญๆจก็ณ๏ผๆญๆฒใ" @dataclass diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index d169d7ddc..9dfc3462e 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -133,9 +133,7 @@ class SamplingParams: prompt: str | list[str] | None = field( default=None, metadata={"batch_sig_exclude": True} ) - negative_prompt: str = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" - ) + negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" prompt_path: str | None = field(default=None, metadata={"batch_sig_exclude": True}) output_path: str | None = field(default=None, metadata={"batch_sig_exclude": True}) output_file_name: str | None = field( @@ -263,12 +261,8 @@ class SamplingParams: ) return_trajectory_latents: bool = False # returns all latents for each timestep return_trajectory_decoded: bool = False # returns decoded latents for each timestep - rollout_return_denoising_env: bool = ( - False # populate ``denoising_env`` (image/pos/neg kwargs, guidance) for RL replay - ) - rollout_return_dit_trajectory: bool = ( - False # per-step noisy latents + final latent + timesteps (RolloutDitTrajectory) - ) + rollout_return_denoising_env: bool = False # populate ``denoising_env`` (image/pos/neg kwargs, guidance) for RL replay + rollout_return_dit_trajectory: bool = False # per-step noisy latents + final latent + timesteps (RolloutDitTrajectory) # 0-indexed denoising-loop step filters; None = all steps. rollout_sde_step_indices: list[int] | None = None rollout_return_step_indices: list[int] | None = None @@ -488,8 +482,7 @@ class SamplingParams: if self.quality not in QUALITY_LEVELS: raise ValueError( - f"quality must be one of {list(QUALITY_LEVELS)}, " - f"got {self.quality!r}" + f"quality must be one of {list(QUALITY_LEVELS)}, got {self.quality!r}" ) # These are always required to be sane regardless of pipeline. diff --git a/python/sglang/multimodal_gen/configs/sample/wan.py b/python/sglang/multimodal_gen/configs/sample/wan.py index 0464ed60e..dfea73974 100644 --- a/python/sglang/multimodal_gen/configs/sample/wan.py +++ b/python/sglang/multimodal_gen/configs/sample/wan.py @@ -51,9 +51,7 @@ class WanT2V_1_3B_SamplingParams(SamplingParams): # Denoising stage guidance_scale: float = 3.0 - negative_prompt: str = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" - ) + negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" num_inference_steps: int = 50 # Wan T2V 1.3B supported resolutions @@ -85,9 +83,7 @@ class WanT2V_14B_SamplingParams(SamplingParams): # Denoising stage guidance_scale: float = 5.0 - negative_prompt: str = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" - ) + negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" num_inference_steps: int = 50 # Wan T2V 14B supported resolutions diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index cf76d3422..a31842f16 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -784,9 +784,11 @@ def _register_configs(): hf_model_paths=["Lightricks/LTX-2"], model_detectors=[ lambda path: "ltx" in path.lower() and "video" in path.lower(), - lambda path: "ltx-2" in path.lower() - and "ltx-2.3" not in path.lower() - and "ltx-2.5" not in path.lower(), + lambda path: ( + "ltx-2" in path.lower() + and "ltx-2.3" not in path.lower() + and "ltx-2.5" not in path.lower() + ), ], ) register_configs( @@ -973,8 +975,9 @@ def _register_configs(): "MiniMax/MiniMax-H3", ], model_detectors=[ - lambda model_id: "minimaxh3" - in model_id.lower().replace("-", "").replace("_", "") + lambda model_id: ( + "minimaxh3" in model_id.lower().replace("-", "").replace("_", "") + ) ], ) register_configs( @@ -984,8 +987,9 @@ def _register_configs(): "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree", ], model_detectors=[ - lambda model_id: "fasth3" - in model_id.lower().replace("-", "").replace("_", "") + lambda model_id: ( + "fasth3" in model_id.lower().replace("-", "").replace("_", "") + ) ], ) # FLUX @@ -1191,7 +1195,7 @@ def _register_configs(): ], model_detectors=[ # Match "sana-wm" or "sana_wm" but NOT plain T2I "sana" checkpoints. - lambda hf_id: ("sana-wm" in hf_id.lower() or "sana_wm" in hf_id.lower()), + lambda hf_id: "sana-wm" in hf_id.lower() or "sana_wm" in hf_id.lower(), ], ) @@ -1203,9 +1207,7 @@ def _register_configs(): "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", ], model_detectors=[ - lambda hf_id: ( - "sana-video" in hf_id.lower() or "sana_video" in hf_id.lower() - ) + lambda hf_id: "sana-video" in hf_id.lower() or "sana_video" in hf_id.lower() ], ) @@ -1293,8 +1295,10 @@ def _register_configs(): "jdopensource/JoyAI-Echo", ], model_detectors=[ - lambda hf_id: ("joy-echo" in hf_id.lower() or "joyai-echo" in hf_id.lower()) - and "image-edit" not in hf_id.lower(), + lambda hf_id: ( + ("joy-echo" in hf_id.lower() or "joyai-echo" in hf_id.lower()) + and "image-edit" not in hf_id.lower() + ), ], ) @@ -1354,9 +1358,11 @@ def _register_configs(): "meituan-longcat/LongCat-Image-Edit-Turbo", ], model_detectors=[ - lambda hf_id: "longcat" in hf_id.lower() - and "edit" in hf_id.lower() - and "turbo" in hf_id.lower(), + lambda hf_id: ( + "longcat" in hf_id.lower() + and "edit" in hf_id.lower() + and "turbo" in hf_id.lower() + ), ], ) @@ -1368,9 +1374,11 @@ def _register_configs(): "meituan-longcat/LongCat-Image-Edit", ], model_detectors=[ - lambda hf_id: "longcat" in hf_id.lower() - and "edit" in hf_id.lower() - and "turbo" not in hf_id.lower(), + lambda hf_id: ( + "longcat" in hf_id.lower() + and "edit" in hf_id.lower() + and "turbo" not in hf_id.lower() + ), ], ) diff --git a/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py b/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py index 3f868280f..998241307 100644 --- a/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py +++ b/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py @@ -1083,7 +1083,7 @@ class SchedulerDisaggMixin: ) use_prefetch = self._compute_ready_queue is not None logger.info( - "Pool mode %s rank %d event loop started " "(multi_rank=%s, prefetch=%s)", + "Pool mode %s rank %d event loop started (multi_rank=%s, prefetch=%s)", role_name, self.gpu_id, is_multi_rank, diff --git a/python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py b/python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py index 5afba92e2..cd44a0a49 100644 --- a/python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py +++ b/python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py @@ -54,8 +54,7 @@ class TransferTensorBuffer: pool_location = "pinned CPU" if device == "cpu" else f"GPU ({device})" logger.info( - "TransferTensorBuffer[%s]: allocated %d MiB %s memory " - "(min_block=%d KiB)", + "TransferTensorBuffer[%s]: allocated %d MiB %s memory (min_block=%d KiB)", role_name, actual_size >> 20, pool_location, diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py index f76d53937..be3dd2552 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py @@ -165,9 +165,9 @@ class DistributedAutograd: if world_size == 1: return input_ - assert ( - input_.dim() == 4 - ), f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" + assert input_.dim() == 4, ( + f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" + ) if world_size == 2 and scatter_dim in (1, 2): fast = _ipc_all_to_all_4d(group, input_, scatter_dim) @@ -303,9 +303,9 @@ class DeviceCommunicatorBase: NOTE: `dst` is the local rank of the destination rank. """ world_size = self.world_size - assert ( - -input_.dim() <= dim < input_.dim() - ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + 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() diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py index 3345cc379..fd9b13070 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py @@ -15,7 +15,6 @@ from .base_device_communicator import DeviceCommunicatorBase class CpuCommunicator(DeviceCommunicatorBase): - def __init__( self, cpu_group: ProcessGroup, @@ -55,9 +54,9 @@ class CpuCommunicator(DeviceCommunicatorBase): NOTE: `dst` is the local rank of the destination rank. """ world_size = self.world_size - assert ( - -input_.dim() <= dim < input_.dim() - ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + 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() @@ -107,7 +106,6 @@ class CpuCommunicator(DeviceCommunicatorBase): class _CPUSHMDistributed: - def __init__(self, communicator: CpuCommunicator): instance_identifier = os.environ["VLLM_DIST_IDENT"] unique_name = communicator.unique_name diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py index a9a0625fb..307f381ef 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py @@ -13,7 +13,6 @@ from sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_ class CudaCommunicator(DeviceCommunicatorBase): - def __init__( self, cpu_group: ProcessGroup, diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py index 45a9ff40f..64d551b5c 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py @@ -28,7 +28,6 @@ logger = init_logger(__name__) class PyNcclCommunicator: - def __init__( self, group: ProcessGroup | StatelessProcessGroup, @@ -48,9 +47,9 @@ class PyNcclCommunicator: """ if not isinstance(group, StatelessProcessGroup): assert dist.is_initialized() - assert ( - dist.get_backend(group) != dist.Backend.NCCL - ), "PyNcclCommunicator should be attached to a non-NCCL group." + assert dist.get_backend(group) != dist.Backend.NCCL, ( + "PyNcclCommunicator should be attached to a non-NCCL group." + ) # note: this rank is the rank in the group self.rank = dist.get_rank(group) self.world_size = dist.get_world_size(group) diff --git a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py index 9e06b97fb..d4d502ac4 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py @@ -87,8 +87,7 @@ def _split_tensor_dict( tensor_list = [] for key, value in tensor_dict.items(): assert "%" not in key, ( - "Avoid having '%' in key " - "as it is used as a separator for nested entries." + "Avoid having '%' in key as it is used as a separator for nested entries." ) if isinstance(value, torch.Tensor): # Note: we cannot use `value.device` here, @@ -425,9 +424,9 @@ class GroupCoordinator: # 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()}" + 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() @@ -481,9 +480,9 @@ class GroupCoordinator: # 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()}" + 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() @@ -599,9 +598,9 @@ class GroupCoordinator: assert src < self.world_size, f"Invalid src rank ({src})" - assert ( - src != self.rank - ), "Invalid source rank. Source rank is the same as the current rank." + assert src != self.rank, ( + "Invalid source rank. Source rank is the same as the current rank." + ) size_tensor = torch.empty(1, dtype=torch.long, device="cpu") @@ -621,9 +620,9 @@ class GroupCoordinator: object_tensor, src=self.ranks[src], group=self.cpu_group ) - assert ( - rank_object == rank_size - ), "Received object sender rank does not match the size sender rank." + assert rank_object == rank_size, ( + "Received object sender rank does not match the size sender rank." + ) obj = pickle.loads(object_tensor.numpy().tobytes()) @@ -652,9 +651,9 @@ class GroupCoordinator: rank = self.rank if rank == src_global_rank: metadata_list: List[Tuple[Any, Any]] = [] - assert isinstance( - tensor_dict, dict - ), f"Expecting a dictionary, got {type(tensor_dict)}" + 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, @@ -736,9 +735,9 @@ class GroupCoordinator: assert dst < self.world_size, f"Invalid dst rank ({dst})" metadata_list: List[Tuple[Any, Any]] = [] - assert isinstance( - tensor_dict, dict - ), f"Expecting a dictionary, got {type(tensor_dict)}" + assert isinstance(tensor_dict, dict), ( + f"Expecting a dictionary, got {type(tensor_dict)}" + ) metadata_list, tensor_list = _split_tensor_dict(tensor_dict) # `metadata_list` lives in CPU memory. # `send_object_list` has serialization & deserialization, @@ -1193,14 +1192,14 @@ class PipelineGroupCoordinator(GroupCoordinator): def get_pipeline_recv_data( self, idx: int = -1, name: str = "latent" ) -> torch.Tensor: - assert ( - len(self.receiving_tasks) > 0 - ), "No tasks to receive, call add_pipeline_recv_task first" + assert len(self.receiving_tasks) > 0, ( + "No tasks to receive, call add_pipeline_recv_task first" + ) receiving_task = self.receiving_tasks.pop(0) receiving_task[0].wait() - assert ( - receiving_task[1] == name and receiving_task[2] == idx - ), "Received tensor does not match the requested" + assert receiving_task[1] == name and receiving_task[2] == idx, ( + "Received tensor does not match the requested" + ) return self.recv_buffer[name][idx] def _pipeline_irecv(self, tensor: torch.tensor): @@ -1255,14 +1254,14 @@ class PipelineGroupCoordinator(GroupCoordinator): self.recv_skip_tasks_queue.append(idx) def get_pipeline_recv_skip_data(self, idx: int = -1) -> torch.Tensor: - assert ( - len(self.receiving_skip_tasks) > 0 - ), "No tasks to receive, call add_pipeline_recv_skip_task first" + assert len(self.receiving_skip_tasks) > 0, ( + "No tasks to receive, call add_pipeline_recv_skip_task first" + ) receiving_skip_task = self.receiving_skip_tasks.pop(0) receiving_skip_task[0].wait() - assert ( - receiving_skip_task[2] == idx - ), "Received tensor does not match the requested" + assert receiving_skip_task[2] == idx, ( + "Received tensor does not match the requested" + ) return self.skip_tensor_recv_buffer[idx] def recv_skip_next(self): diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py index d1d3bc2f6..0c5476bb8 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py @@ -43,9 +43,9 @@ def set_seq_parallel_pg_by_sp_groups( """ sp_degree = sp_ring_degree * sp_ulysses_degree assert sp_degree > 0 - assert all( - len(g) == sp_degree for g in sp_groups - ), f"Each SP group must have size {sp_degree}, got sizes {[len(g) for g in sp_groups]}" + assert all(len(g) == sp_degree for g in sp_groups), ( + f"Each SP group must have size {sp_degree}, got sizes {[len(g) for g in sp_groups]}" + ) ulyssess_pg = None ring_pg = None diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index eb642a83b..a22b54dfc 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -298,7 +298,6 @@ def init_distributed_environment( ) if timeout is not None: - extra_args["timeout"] = datetime.timedelta(seconds=timeout) logger.info(f"Setting distributed timeout to {timeout} seconds") @@ -325,9 +324,9 @@ def init_distributed_environment( 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" + assert _WORLD.world_size == torch.distributed.get_world_size(), ( + "world group already initialized with a different world size" + ) _sync_srt_world_group() @@ -648,12 +647,12 @@ def maybe_init_distributed_environment_and_model_parallel( if _WORLD is not None and model_parallel_is_initialized(): # make sure the tp and sp sizes are correct - assert ( - get_tp_world_size() == tp_size - ), f"You are trying to initialize model parallel groups with size {tp_size}, but they are already initialized with size {get_tp_world_size()}" - assert ( - get_sp_world_size() == sp_size - ), f"You are trying to initialize model parallel groups with size {sp_size}, but they are already initialized with size {get_sp_world_size()}" + assert get_tp_world_size() == tp_size, ( + f"You are trying to initialize model parallel groups with size {tp_size}, but they are already initialized with size {get_tp_world_size()}" + ) + assert get_sp_world_size() == sp_size, ( + f"You are trying to initialize model parallel groups with size {sp_size}, but they are already initialized with size {get_sp_world_size()}" + ) return local_rank = int(os.environ.get("LOCAL_RANK", 0)) world_size = int(os.environ.get("WORLD_SIZE", 1)) @@ -773,9 +772,9 @@ def is_the_same_node_as( memory system (shared access to shared memory). """ if isinstance(pg, ProcessGroup): - assert ( - torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL - ), "in_the_same_node_as should be tested with a non-NCCL group." + 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) @@ -933,9 +932,9 @@ def is_pipeline_last_stage() -> bool: # CFG def get_cfg_group() -> GroupCoordinator: - assert ( - _CFG is not None - ), "classifier_free_guidance parallel group is not initialized" + assert _CFG is not None, ( + "classifier_free_guidance parallel group is not initialized" + ) return _CFG diff --git a/python/sglang/multimodal_gen/runtime/distributed/utils.py b/python/sglang/multimodal_gen/runtime/distributed/utils.py index ef481c9e3..a3577610d 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/utils.py +++ b/python/sglang/multimodal_gen/runtime/distributed/utils.py @@ -134,13 +134,13 @@ class StatelessProcessGroup: """ if self.rank == src: self.expire_data() - key = f"broadcast_from/{src}/" f"{self.broadcast_send_counter}" + key = f"broadcast_from/{src}/{self.broadcast_send_counter}" self.store.set(key, pickle.dumps(obj)) self.broadcast_send_counter += 1 self.entries.append((key, time.perf_counter())) return obj else: - key = f"broadcast_from/{src}/" f"{self.broadcast_recv_src_counter[src]}" + key = f"broadcast_from/{src}/{self.broadcast_recv_src_counter[src]}" recv_obj = pickle.loads(self.store.get(key)) self.broadcast_recv_src_counter[src] += 1 return recv_obj diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py index 1699fca94..955df8f3a 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py @@ -14,7 +14,6 @@ logger = init_logger(__name__) class RaiseNotImplementedAction(argparse.Action): - def __call__(self, parser, namespace, values, option_string=None): raise NotImplementedError(f"The {option_string} option is not yet implemented") diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py index 3b600bbba..71492523d 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py @@ -209,9 +209,9 @@ def _build_response( peak_memory_mb = result.peak_memory_mb if result.peak_memory_mb > 0 else None rollout_trajectory_data = result.rollout_trajectory_data if rollout: - assert ( - rollout_trajectory_data is not None - ), "rollout_trajectory_data must be present when rollout=True" + assert rollout_trajectory_data is not None, ( + "rollout_trajectory_data must be present when rollout=True" + ) serialized_dit_timesteps = None serialized_dit_sigmas = None diff --git a/python/sglang/multimodal_gen/runtime/launch_server.py b/python/sglang/multimodal_gen/runtime/launch_server.py index df6fed265..f4b655ae5 100644 --- a/python/sglang/multimodal_gen/runtime/launch_server.py +++ b/python/sglang/multimodal_gen/runtime/launch_server.py @@ -662,7 +662,7 @@ def launch_disagg_role(server_args: ServerArgs): role_type = server_args.disagg_role if server_args.disagg_server_addr is None: raise ValueError( - "--disagg-server-addr is required for --disagg-role " f"{role_type.value}" + f"--disagg-server-addr is required for --disagg-role {role_type.value}" ) # Derive endpoints diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py b/python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py index 9635a6740..5c17d0405 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py @@ -138,7 +138,9 @@ def configure_sta( print("\nStrategy usage counts:") total_heads = time_step_num * layer_num * head_num # Fixed dimensions for strategy, count in strategy_counts.items(): - print(f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)") + print( + f"Strategy {strategy}: {count} heads ({count / total_heads * 100:.2f}%)" + ) # Convert dictionary to 3D list with fixed dimensions mask_strategy_3d = dict_to_3d_list( @@ -221,7 +223,9 @@ def configure_sta( print("\nStrategy usage counts:") total_heads = time_step_num * layer_num * head_num # Fixed dimensions for strategy, count in strategy_counts.items(): - print(f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)") + print( + f"Strategy {strategy}: {count} heads ({count / total_heads * 100:.2f}%)" + ) # Convert dictionary to 3D list with fixed dimensions mask_strategy_3d = dict_to_3d_list( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py index 56c274c42..17ad922a0 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py @@ -13,7 +13,6 @@ from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum class AITERSageBackend(AttentionBackend): - @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.AITER_SAGE @@ -35,7 +34,6 @@ class AITERSageBackend(AttentionBackend): class AITERSageImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py index b25388ae4..98e74fada 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py @@ -173,7 +173,6 @@ class AscendFAMetadataBuilder(AttentionMetadataBuilder): class AscendFABackend(AttentionBackend): - @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.FA @@ -198,7 +197,6 @@ class AscendFABackend(AttentionBackend): class AscendFAImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py index 733b3bbe4..7f50900f3 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py @@ -125,7 +125,6 @@ class AttentionMetadataBuilder(ABC, Generic[T]): class AttentionLayer(Protocol): - _k_scale: torch.Tensor _v_scale: torch.Tensor _k_scale_float: float @@ -142,7 +141,6 @@ class AttentionLayer(Protocol): class AttentionImpl(ABC, Generic[T]): - @abstractmethod def __init__( self, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py index fa6df1ef2..f51ffa9ac 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py @@ -21,7 +21,6 @@ BSA_BLOCK_SIZE = 128 class BlockSparseAttentionBackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -126,7 +125,6 @@ class BlockSparseAttentionMetadataBuilder(AttentionMetadataBuilder): class BlockSparseAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/cube_sparse_attn/backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/cube_sparse_attn/backend.py index 26f2be21d..50c4d94ac 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/cube_sparse_attn/backend.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/cube_sparse_attn/backend.py @@ -36,7 +36,6 @@ from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum class CubeSparseAttentionBackend(AttentionBackend): - @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.CUBE_SPARSE_ATTN @@ -63,7 +62,6 @@ class CubeSparseAttentionMetadata(AttentionMetadata): class CubeSparseAttentionMetadataBuilder(AttentionMetadataBuilder): - def __init__(self): pass @@ -279,7 +277,6 @@ def cube_sparse_attention( class CubeSparseAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index 5710e321d..b0bb625c9 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py @@ -67,9 +67,9 @@ def flash_attn_varlen_func_fake_out( head_dim_v = v.shape[-1] if cu_seqlens_q is not None: - assert cu_seqlens_q.shape == ( - batch_size + 1, - ), "cu_seqlens_q must have shape (batch_size + 1,)" + assert cu_seqlens_q.shape == (batch_size + 1,), ( + "cu_seqlens_q must have shape (batch_size + 1,)" + ) assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32" assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous" @@ -129,9 +129,9 @@ def flash_attn_varlen_func_fake_out_lse( head_dim_v = v.shape[-1] if cu_seqlens_q is not None: - assert cu_seqlens_q.shape == ( - batch_size + 1, - ), "cu_seqlens_q must have shape (batch_size + 1,)" + assert cu_seqlens_q.shape == (batch_size + 1,), ( + "cu_seqlens_q must have shape (batch_size + 1,)" + ) assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32" assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous" @@ -329,7 +329,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder): class FlashAttentionBackend(AttentionBackend): - @classmethod def supports_ring_rotation(cls) -> bool: return True diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py index 62a1974ad..7b6b8df43 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py @@ -44,7 +44,6 @@ class FlashAttention2Backend(AttentionBackend): class FlashAttention2Impl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/laser_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/laser_attn.py index 487a4f936..5afb82b95 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/laser_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/laser_attn.py @@ -32,7 +32,6 @@ _BF16_LASER_SCALE = 256.0 class LaserAttentionBackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -49,7 +48,6 @@ class LaserAttentionBackend(AttentionBackend): class LaserAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/rain_fusion_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/rain_fusion_attn.py index 95ea1374b..48245874e 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/rain_fusion_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/rain_fusion_attn.py @@ -22,7 +22,6 @@ logger = init_logger(__name__) class RainFusionAttentionBackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -99,7 +98,6 @@ class RainFusionAttentionMetadataBuilder(AttentionMetadataBuilder): class RainFusionAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py index c4ffbaea5..a137d8017 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py @@ -33,7 +33,6 @@ def _trailing_padding_used_len( class SageAttentionBackend(AttentionBackend): - @classmethod def supports_ring_rotation(cls) -> bool: return True @@ -54,7 +53,6 @@ class SageAttentionBackend(AttentionBackend): class SageAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py index 02bbd5e6b..73f55158c 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py @@ -28,7 +28,6 @@ _MPS_VARLEN_QUERY_CHUNK_SIZE = 128 class SDPABackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -49,7 +48,6 @@ class SDPABackend(AttentionBackend): class SDPAImpl(AttentionImpl): - def __init__( self, num_heads: int, @@ -268,7 +266,7 @@ class DynamicCudnnSDPAImpl(SDPAImpl): # cuDNN raises "No available kernel" for some shapes; pin the # FA fail-safe path for this layer and keep going. logger.warning( - "cuDNN SDPA failed (%s); falling back to FlashAttention " "for %s.", + "cuDNN SDPA failed (%s); falling back to FlashAttention for %s.", e, type(self).__name__, ) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py index 37a1acf30..0db59e292 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py @@ -35,7 +35,6 @@ logger = init_logger(__name__) class RangeDict(dict): - def __getitem__(self, item: int) -> str: for key in self.keys(): if isinstance(key, tuple): @@ -81,7 +80,6 @@ class SlidingTileAttentionMetadata(AttentionMetadata): class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder): - def __init__(self): pass @@ -105,7 +103,6 @@ class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder): class SlidingTileAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py index b7be98848..5f984631a 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py @@ -99,7 +99,6 @@ class SolAttnBackend(AttentionBackend): class SolAttnImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py index 3acae28e6..b3dcd6dee 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py @@ -356,9 +356,9 @@ class SageSparseLinearAttentionImpl(AttentionImpl, nn.Module): ) -> None: nn.Module.__init__(self) - assert ( - SAGESLA_ENABLED - ), "Install spas_sage_attn(pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation) first to enable SageSLA." + assert SAGESLA_ENABLED, ( + "Install spas_sage_attn(pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation) first to enable SageSLA." + ) self.num_heads = num_heads self.head_size = head_size @@ -454,7 +454,9 @@ class SageSparseLinearAttentionImpl(AttentionImpl, nn.Module): assert headdim in [ 64, 128, - ], "headdim should be in [64, 128]. For other headdim, you can use padding and specify the softmax scale." + ], ( + "headdim should be in [64, 128]. For other headdim, you can use padding and specify the softmax scale." + ) # Quantize Q, K to INT8 q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant(q, k, km, BLKQ, BLKK) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_video_gen_2_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_video_gen_2_attn.py index 0d07259c0..f5d03072c 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_video_gen_2_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_video_gen_2_attn.py @@ -42,7 +42,6 @@ logger = init_logger(__name__) class SparseVideoGen2AttentionBackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -116,7 +115,6 @@ def _require_kwarg(kwargs: dict[str, Any], name: str) -> Any: class SparseVideoGen2AttentionMetadataBuilder(AttentionMetadataBuilder): - def __init__(self) -> None: pass @@ -180,7 +178,6 @@ class SparseVideoGen2AttentionMetadataBuilder(AttentionMetadataBuilder): class SparseVideoGen2AttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, @@ -461,9 +458,9 @@ class SparseVideoGen2AttentionImpl(AttentionImpl): if prompt_length is None: prompt_length = context_length - assert ( - seq_len == context_length + num_frame * frame_size - ), f"Query Shape: {seq_len} is not equivalent to {context_length} + {num_frame} * {frame_size}" + assert seq_len == context_length + num_frame * frame_size, ( + f"Query Shape: {seq_len} is not equivalent to {context_length} + {num_frame} * {frame_size}" + ) # Determine if we use Full Attention to calculate full_attention_flag = False diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py index 4f72dcc1f..6607321a8 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py @@ -98,9 +98,7 @@ def construct_variable_block_sizes( t_sizes[:, None, None] # [n_t, 1, 1] * h_sizes[None, :, None] # [1, n_h, 1] * w_sizes[None, None, :] # [1, 1, n_w] - ).reshape( - -1 - ) # [n_t * n_h * n_w] + ).reshape(-1) # [n_t * n_h * n_w] return block_sizes diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn_h3.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn_h3.py index 8367dce0f..2a0812284 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn_h3.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn_h3.py @@ -393,7 +393,7 @@ class VideoSparseAttentionH3Impl(AttentionImpl): attn_metadata: AttentionMetadata, ) -> torch.Tensor: raise NotImplementedError( - "VSA-H3 serves MiniMax-H3's packed varlen attention; use " "forward_varlen." + "VSA-H3 serves MiniMax-H3's packed varlen attention; use forward_varlen." ) def forward_varlen( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py index e07c74336..614c38307 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py @@ -26,7 +26,6 @@ logger = init_logger(__name__) class VMOBAAttentionBackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -89,7 +88,6 @@ def pad_input(hidden_states, indices, batch, seqlen): class VideoMobaAttentionMetadataBuilder(AttentionMetadataBuilder): - def __init__(self): pass @@ -124,7 +122,9 @@ class VideoMobaAttentionMetadataBuilder(AttentionMetadataBuilder): raw_latent_shape[0] % patch_size[0] == 0 and raw_latent_shape[1] % patch_size[1] == 0 and raw_latent_shape[2] % patch_size[2] == 0 - ), f"spatial patch_resolution {raw_latent_shape} should be divisible by patch_size {patch_size}" + ), ( + f"spatial patch_resolution {raw_latent_shape} should be divisible by patch_size {patch_size}" + ) patch_resolution = [ t // pt for t, pt in zip(raw_latent_shape, patch_size, strict=False) ] @@ -150,7 +150,6 @@ class VideoMobaAttentionMetadataBuilder(AttentionMetadataBuilder): class VMOBAAttentionImpl(AttentionImpl): - def __init__( self, num_heads, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py index 613f2df66..783010317 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py @@ -62,7 +62,6 @@ def _get_cu_seqlens(device_index: int, bsz: int, seqlen: int) -> torch.Tensor: class XPUAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index f0776fa5d..800d04400 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -562,9 +562,9 @@ class UlyssesAttention_VSA(UlyssesAttention): "K/V-gather SP does not support video sparse attention." ) # Check text tokens are not supported for VSA now - assert ( - replicated_q is None and replicated_k is None and replicated_v is None - ), "Replicated QKV is not supported for VSA now" + assert replicated_q is None and replicated_k is None and replicated_v is None, ( + "Replicated QKV is not supported for VSA now" + ) # Check input shapes assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors" @@ -1042,9 +1042,9 @@ class USPAttention(nn.Module): inv_indices = attn_mask_meta["inv_indices"] # Guard against a caller passing meta from a different # mask shape (silent corruption otherwise). - assert ( - inv_indices.shape[0] == bs * seq - ), "attn_mask_meta shape does not match attn_mask" + assert inv_indices.shape[0] == bs * seq, ( + "attn_mask_meta shape does not match attn_mask" + ) # All-False mask: FA varlen rejects zero-length input. # Fall through to SDPA which handles it via broadcast. # (Joint attention with an image side is always non-empty @@ -1169,9 +1169,9 @@ class USPAttention(nn.Module): # Zero-copy tail path: run varlen FA straight over the # padded layout, each row split into [valid | pad] segments # (contiguous reshapes only, no repacking). - assert ( - cu_tail.numel() == 2 * bs + 1 - ), "cu_seqlens_tail does not match the batch size" + assert cu_tail.numel() == 2 * bs + 1, ( + "cu_seqlens_tail does not match the batch size" + ) out = flash_attn_varlen_func( q=q.reshape(bs * seq, *q.shape[2:]), k=k.reshape(bs * seq, *k.shape[2:]), @@ -1259,9 +1259,9 @@ class USPAttention(nn.Module): gathered_mask_meta = build_varlen_mask_meta(gathered_mask) indices = gathered_mask_meta["indices"] inv_indices = gathered_mask_meta["inv_indices"] - assert ( - inv_indices.shape[0] == bs * seq - ), "gathered attn_mask shape does not match q/k/v" + assert inv_indices.shape[0] == bs * seq, ( + "gathered attn_mask shape does not match q/k/v" + ) if indices.shape[0] > 0: q_unpad, k_unpad, v_unpad = fused_pack_qkv(q, k, v, indices) out_unpad = flash_attn_varlen_func( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py index 8b78133fc..f86cb9b94 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py @@ -173,8 +173,7 @@ def _log_component_attn_backend_summary( backend_parts.append(backend_name) logger.info_once( - f"Attention backends for {context.component_name}: " - f"{', '.join(backend_parts)}" + f"Attention backends for {context.component_name}: {', '.join(backend_parts)}" ) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py index c40948d68..31e2c8103 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py @@ -49,9 +49,9 @@ def single_all_to_all(input, local_seq_2_local_head, group, async_op=False): # b, s, n, h if local_seq_2_local_head: bs, local_seq_len, num_total_head, head_dim = input.shape - assert ( - num_total_head % seq_world_size == 0 - ), f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!" + assert num_total_head % seq_world_size == 0, ( + f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!" + ) input_t = rearrange( input, "bs seq_len (w h) d -> w bs seq_len h d", diff --git a/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py b/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py index c8722b31c..e42284adc 100644 --- a/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py +++ b/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py @@ -221,8 +221,7 @@ class CausalSelfAttentionKVCache: self.k[:, sink_tokens : sink_tokens + num_rolled_tokens] = ( self.k[ :, - sink_tokens - + num_evicted_tokens : sink_tokens + sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, ].clone() @@ -230,8 +229,7 @@ class CausalSelfAttentionKVCache: self.v[:, sink_tokens : sink_tokens + num_rolled_tokens] = ( self.v[ :, - sink_tokens - + num_evicted_tokens : sink_tokens + sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, ].clone() @@ -244,8 +242,7 @@ class CausalSelfAttentionKVCache: :, ] = self.k[ :, - sink_tokens - + num_evicted_tokens : sink_tokens + sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, cache_head_slice, @@ -258,8 +255,7 @@ class CausalSelfAttentionKVCache: :, ] = self.v[ :, - sink_tokens - + num_evicted_tokens : sink_tokens + sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, cache_head_slice, diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 356d25ffa..c22de734a 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -446,7 +446,6 @@ class LayerNorm(CustomOp): # FSDP's MixedPrecisionPolicy @CustomOp.register("fp32_layer_norm") class FP32LayerNorm(CustomOp, nn.LayerNorm): - def __init__( self, normalized_shape, diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py index 7dd00c526..ef6f0621d 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py @@ -80,7 +80,7 @@ def register_quantization_config(quantization: str): ) if not issubclass(quant_config_cls, QuantizationConfig): raise ValueError( - "The quantization config must be a subclass of " "`QuantizationConfig`." + "The quantization config must be a subclass of `QuantizationConfig`." ) _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls QUANTIZATION_METHODS.append(quantization) diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py index cde021f64..357a35e4a 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py @@ -64,8 +64,7 @@ class KitchenW4A4Config(QuantizationConfig): continue if marker_format != "convrot_w4a4": raise ValueError( - f"Unsupported Comfy W4A4 format for {prefix!r}: " - f"{marker_format!r}" + f"Unsupported Comfy W4A4 format for {prefix!r}: {marker_format!r}" ) self._parse_marker(prefix, marker) diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a8_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a8_config.py index 38a12f667..f80d538ba 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a8_config.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a8_config.py @@ -55,8 +55,7 @@ class KitchenW4A8Config(QuantizationConfig): continue if marker_format != "asym_w4a8_int8": raise ValueError( - f"Unsupported Comfy W4A8 format for {prefix!r}: " - f"{marker_format!r}" + f"Unsupported Comfy W4A8 format for {prefix!r}: {marker_format!r}" ) if marker.get("convrot") is not True: raise ValueError( diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py index ca39b2f0e..78e37c64a 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py @@ -254,9 +254,9 @@ class Fp8LinearMethod(LinearMethodBase): ) layer.input_scale = None elif _is_cpu: - assert ( - _is_cpu_amx_available - ), "Fp8LinearMethod on CPU requires that CPU has AMX support" + assert _is_cpu_amx_available, ( + "Fp8LinearMethod on CPU requires that CPU has AMX support" + ) _amx_process_weight_after_loading(layer, ["weight"]) layer.weight_scale_inv = torch.nn.Parameter( layer.weight_scale_inv.data, requires_grad=False diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index eebcc1c7a..c1b36a427 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py @@ -122,7 +122,7 @@ def _prepare_nvfp4_swiglu_fusion_weights( ) if weight.shape[0] % 128 != 0: raise ValueError( - "Fused NVFP4 SwiGLU requires FC1 N % 128 == 0, " f"got N={weight.shape[0]}." + f"Fused NVFP4 SwiGLU requires FC1 N % 128 == 0, got N={weight.shape[0]}." ) # FLUX.2 stores [gate; up]. The kernel consumes 64-row groups in diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py index 62b84ee17..cd11a5539 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py @@ -201,7 +201,6 @@ class ModelSlimConfig(QuantizationConfig): class ModelSlimLinearMethod(LinearMethodBase): - def __init__(self, quantization_config: ModelSlimConfig): self.quantization_config = quantization_config diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py index 00bb925b9..e2a20a773 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py @@ -35,7 +35,6 @@ MXFP4_DUAL_LEVEL_RATIO = 16 class ModelSlimMXFP4Scheme(ModelSlimLinearScheme): - def create_weights( self, layer: torch.nn.Module, diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py index c4c5dfb7b..44be13f47 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py @@ -25,7 +25,6 @@ MXFP8_BLOCK_SIZE = 32 class ModelSlimMXFP8Scheme(ModelSlimLinearScheme): - def create_weights( self, layer: torch.nn.Module, diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py index e2ddc79b7..629a2f9a1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py @@ -185,8 +185,7 @@ class Mxfp4LinearMethod(LinearMethodBase): if any(fn is None for fn in (dynamic_mxfp4_quant, shuffle_weight, gemm_a4w4)): raise RuntimeError( - "aiter MXFP4 kernels not available. " - "Install aiter with MXFP4 support." + "aiter MXFP4 kernels not available. Install aiter with MXFP4 support." ) weight_data = layer.weight.data diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/factory.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/factory.py index 807660ea0..f03978264 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/factory.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/factory.py @@ -118,9 +118,9 @@ def get_rotary_pos_embed( if rope_dim_list is None: rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] - assert ( - sum(rope_dim_list) == head_dim - ), "sum(rope_dim_list) should equal to head_dim of attention layer" + assert sum(rope_dim_list) == head_dim, ( + "sum(rope_dim_list) should equal to head_dim of attention layer" + ) # Get SP info - now handled within NDRotaryEmbedding # sp_group = get_sp_group() diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py index bfe77777f..e284afab8 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py @@ -298,9 +298,9 @@ class NDRotaryEmbedding(torch.nn.Module): self.theta_rescale_factor = [theta_rescale_factor[0]] * self.ndim else: self.theta_rescale_factor = theta_rescale_factor - assert ( - len(self.theta_rescale_factor) == self.ndim - ), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + assert len(self.theta_rescale_factor) == self.ndim, ( + "len(theta_rescale_factor) should equal to len(rope_dim_list)" + ) if isinstance(interpolation_factor, (int, float)): self.interpolation_factor = [interpolation_factor] * self.ndim @@ -308,9 +308,9 @@ class NDRotaryEmbedding(torch.nn.Module): self.interpolation_factor = [interpolation_factor[0]] * self.ndim else: self.interpolation_factor = interpolation_factor - assert ( - len(self.interpolation_factor) == self.ndim - ), "len(interpolation_factor) should equal to len(rope_dim_list)" + assert len(self.interpolation_factor) == self.ndim, ( + "len(interpolation_factor) should equal to len(rope_dim_list)" + ) self.rope_generators: list[OneDRotaryEmbedding] = torch.nn.ModuleList() _config_to_gen_idx: dict[tuple, int] = {} diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py index 459a38620..57dc71211 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py @@ -195,7 +195,7 @@ def apply_flashinfer_rope_qk_inplace( raise ValueError("positions must be a 1D Tensor") if positions.numel() != bsz * seqlen: raise ValueError( - f"positions length must be bsz*seqlen={bsz*seqlen}, got {positions.numel()}" + f"positions length must be bsz*seqlen={bsz * seqlen}, got {positions.numel()}" ) positions = positions.to(device=q.device, dtype=torch.long) diff --git a/python/sglang/multimodal_gen/runtime/layers/usp.py b/python/sglang/multimodal_gen/runtime/layers/usp.py index 1ed06e785..8b2d66db1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/usp.py +++ b/python/sglang/multimodal_gen/runtime/layers/usp.py @@ -312,9 +312,9 @@ def _usp_input_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor: # Shape transition: [b, s_local, h_global, d] -> [h_global, b, s_local, d] permute_order = (2, 0, 1, 3) - assert ( - h_global % world_size == 0 - ), f"h_global ({h_global}) must be divisible by world_size ({world_size})" + assert h_global % world_size == 0, ( + f"h_global ({h_global}) must be divisible by world_size ({world_size})" + ) h_local, s_global = h_global // world_size, s_local * world_size @@ -488,9 +488,9 @@ def _usp_input_all_to_all_varlen( assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}" assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}" - assert ( - len(seq_lens) == world_size - ), f"seq_lens must have length {world_size}, got {len(seq_lens)}" + assert len(seq_lens) == world_size, ( + f"seq_lens must have length {world_size}, got {len(seq_lens)}" + ) rank = get_ulysses_parallel_rank() @@ -504,12 +504,12 @@ def _usp_input_all_to_all_varlen( # Shape transition: [b, s_local, h_global, d] -> [h_global, b, s_local, d] permute_order = (2, 0, 1, 3) - assert ( - s_local == seq_lens[rank] - ), f"s_local ({s_local}) must equal seq_lens[{rank}] ({seq_lens[rank]})" - assert ( - h_global % world_size == 0 - ), f"h_global ({h_global}) must be divisible by world_size ({world_size})" + assert s_local == seq_lens[rank], ( + f"s_local ({s_local}) must equal seq_lens[{rank}] ({seq_lens[rank]})" + ) + assert h_global % world_size == 0, ( + f"h_global ({h_global}) must be divisible by world_size ({world_size})" + ) h_local = h_global // world_size @@ -578,9 +578,9 @@ def _usp_output_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor: # Shape transition: [b, s_global, h_local, d] -> [s_global, b, h_local, d] permute_order = (1, 0, 2, 3) - assert ( - s_global % world_size == 0 - ), f"s_global ({s_global}) must be divisible by world_size ({world_size})" + assert s_global % world_size == 0, ( + f"s_global ({s_global}) must be divisible by world_size ({world_size})" + ) s_local, h_global = s_global // world_size, h_local * world_size @@ -632,9 +632,9 @@ def _usp_output_all_to_all_varlen( assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}" assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}" - assert ( - len(seq_lens) == world_size - ), f"seq_lens must have length {world_size}, got {len(seq_lens)}" + assert len(seq_lens) == world_size, ( + f"seq_lens must have length {world_size}, got {len(seq_lens)}" + ) rank = get_ulysses_parallel_rank() @@ -648,9 +648,9 @@ def _usp_output_all_to_all_varlen( # Shape transition: [b, s_global, h_local, d] -> [h_local, b, s_global, d] permute_order = (2, 0, 1, 3) - assert s_global == sum( - seq_lens - ), f"s_global ({s_global}) must equal sum(seq_lens) ({sum(seq_lens)})" + assert s_global == sum(seq_lens), ( + f"s_global ({s_global}) must equal sum(seq_lens) ({sum(seq_lens)})" + ) s_local = seq_lens[rank] diff --git a/python/sglang/multimodal_gen/runtime/layers/utils.py b/python/sglang/multimodal_gen/runtime/layers/utils.py index 2454eae78..b4e115073 100644 --- a/python/sglang/multimodal_gen/runtime/layers/utils.py +++ b/python/sglang/multimodal_gen/runtime/layers/utils.py @@ -147,7 +147,6 @@ class CustomOpWrapper: def real_impl(self) -> Callable: if self._impl is None: if not hasattr(torch.ops.sglang, self.op_name): - # NOTE(dark): if torch compile fail here, mark the decorator as eager # lazy registration does not work with torch compile direct_register_custom_op( @@ -240,15 +239,15 @@ def register_custom_op( """ extra_kwarg_keys = set(extra_kwargs.keys()) expected_kwarg_keys = set({"out_shape", "fake_impl"}) - assert ( - expected_kwarg_keys >= extra_kwarg_keys - ), f"Unexpected extra kwargs: {extra_kwarg_keys - expected_kwarg_keys}" + assert expected_kwarg_keys >= extra_kwarg_keys, ( + f"Unexpected extra kwargs: {extra_kwarg_keys - expected_kwarg_keys}" + ) has_out_shape = "out_shape" in extra_kwargs has_fake_impl = "fake_impl" in extra_kwargs - assert not ( - has_out_shape and has_fake_impl - ), "Only one of `out_shape` or `fake_impl` should be provided." + assert not (has_out_shape and has_fake_impl), ( + "Only one of `out_shape` or `fake_impl` should be provided." + ) # Assume inplace if neither out_shape nor fake_impl is provided if not (has_out_shape or has_fake_impl): extra_kwargs["out_shape"] = None diff --git a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py index 64efd15cb..2f3efe73e 100644 --- a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py @@ -258,9 +258,9 @@ class TimestepEmbedder(nn.Module): t, self.frequency_embedding_size, self.max_period, dtype=self.freq_dtype ).to(self.mlp.fc_in.weight.dtype) if timestep_seq_len is not None: - assert ( - t_freq.shape[0] % timestep_seq_len == 0 - ), "timestep length is not divisible by timestep_seq_len" + assert t_freq.shape[0] % timestep_seq_len == 0, ( + "timestep length is not divisible by timestep_seq_len" + ) batch_size = t_freq.shape[0] // timestep_seq_len t_freq = t_freq.unflatten(0, (batch_size, timestep_seq_len)) # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py index 9890b8bae..e1116d4d5 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py @@ -655,9 +655,9 @@ class ComponentLoader(ABC): ] expected_library = loader_cls.expected_library # Assert that the library matches what's expected for this component type - assert ( - transformers_or_diffusers == expected_library - ), f"{loader_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" + assert transformers_or_diffusers == expected_library, ( + f"{loader_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" + ) loader = loader_cls() loader.component_type = structural_component_name loader.component_architecture = component_architecture diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py index c51114fb9..756ec3236 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py @@ -92,7 +92,6 @@ class PEModelWrapper(nn.Module, LayerwiseOffloadableModuleMixin): class SGLangPEModelWrapper: - def __init__(self, model_url): self.model_url = model_url.rstrip("/") # Tokenizer is initialized separately during pipeline setup diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py index 44b5b4139..1bc0ae537 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py @@ -50,9 +50,9 @@ class SchedulerLoader(ComponentLoader): getattr(server_args.pipeline_config, "scheduler_class_override", None) or checkpoint_class_name ) - assert ( - class_name is not None - ), "Model config does not contain a _class_name attribute. Only diffusers format is supported." + assert class_name is not None, ( + "Model config does not contain a _class_name attribute. Only diffusers format is supported." + ) if checkpoint_class_name is not None and class_name != checkpoint_class_name: logger.info( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py index 94b5d237f..511066434 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py @@ -31,9 +31,9 @@ class SoundTokenizerLoader(PlainStateDictComponentLoader): component_model_path, server_args, component_name ) class_name = config.pop("_class_name", None) or self.component_architecture - assert ( - class_name is not None - ), "Sound tokenizer class name must be available from component config." + assert class_name is not None, ( + "Sound tokenizer class name must be available from component config." + ) server_args.model_paths[component_name] = component_model_path diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 590c2c1f0..c0c44e5ae 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -882,8 +882,9 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): model_device = local_torch_device encoder_tp_group = get_folding_tp_group(model_config) - with use_tensor_parallel_group(encoder_tp_group), set_default_torch_dtype( - PRECISION_TO_TYPE[dtype] + with ( + use_tensor_parallel_group(encoder_tp_group), + set_default_torch_dtype(PRECISION_TO_TYPE[dtype]), ): with model_device, skip_init_modules(): architectures = getattr(model_config, "architectures", []) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index a07a33e8a..5171b5402 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -467,9 +467,9 @@ class VAELoader(WeightOverrideComponentLoader): ) class_name = config.pop("_class_name", None) - assert ( - class_name is not None - ), "Model config does not contain a _class_name attribute. Only diffusers format is supported." + assert class_name is not None, ( + "Model config does not contain a _class_name attribute. Only diffusers format is supported." + ) component_type = self.structural_component_type(component_name) if component_type in ("vae", "video_vae"): @@ -577,9 +577,9 @@ class VAELoader(WeightOverrideComponentLoader): vae_precision, ) - assert ( - len(safetensors_list) >= 1 - ), f"Found no safetensors files in {component_weights_path}" + assert len(safetensors_list) >= 1, ( + f"Found no safetensors files in {component_weights_path}" + ) if direct_gpu_weight_loading: _assign_direct_gpu_vae_state( vae, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vl_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vl_encoder_loader.py index c684e0e2b..3b4159794 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vl_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vl_encoder_loader.py @@ -25,7 +25,6 @@ class VisionLanguageEncoderLoader(ComponentLoader): component_name: str = "vision_language_encoder", ) -> Any: if self.structural_component_type(component_name) == "vision_language_encoder": - if server_args.srt_encoder_url is not None: health_url = server_args.srt_encoder_url.rstrip("/") + "/health" try: diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py index b21ccaaae..e616ea4f8 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py @@ -29,9 +29,9 @@ class VocoderLoader(PlainStateDictComponentLoader): component_model_path, server_args, component_name ) class_name = config.pop("_class_name", None) or self.component_architecture - assert ( - class_name is not None - ), "Vocoder class name must be available from component config or pipeline config." + assert class_name is not None, ( + "Vocoder class name must be available from component config or pipeline config." + ) server_args.model_paths[component_name] = component_model_path diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index e0c14e77e..5676dd198 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -770,8 +770,7 @@ def resolve_transformer_quant_load_spec( ) if server_args.nunchaku_config is not None: raise ValueError( - "Per-layer checkpoint quantization and Nunchaku are mutually " - "exclusive" + "Per-layer checkpoint quantization and Nunchaku are mutually exclusive" ) quant_config = checkpoint_quant_config elif getattr(model_cls, "handles_checkpoint_quantization", False): diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py index a409abb6f..44f6309a5 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py @@ -77,7 +77,6 @@ temp_dir = tempfile.gettempdir() class DisabledTqdm(tqdm): - def __init__(self, *args, **kwargs): kwargs["disable"] = True super().__init__(*args, **kwargs) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index dab6e2855..2139b876f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -434,8 +434,8 @@ class GPUWorker(GPUWorkerPostTrainingMixin): ), log_reqs=[req], return_req=False, - save_output_paths=lambda output_batch, req=req: self._save_output_paths( - req, output_batch + save_output_paths=lambda output_batch, req=req: ( + self._save_output_paths(req, output_batch) ), error_context=f"grouped request {req.request_id}", execution_start_time=group_start_time, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency.py index cc5ffafb6..0835d331a 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency.py @@ -82,8 +82,7 @@ def normalize_component_residency( for raw_selector, raw_mode in entries: if not isinstance(raw_selector, str) or not isinstance(raw_mode, str): raise ComponentResidencyError( - "Invalid component residency assignment: " - f"{raw_selector!r}={raw_mode!r}" + f"Invalid component residency assignment: {raw_selector!r}={raw_mode!r}" ) selector = raw_selector.strip().replace("-", "_").lower() mode = raw_mode.strip().replace("_", "-").lower() @@ -149,8 +148,7 @@ def resolve_diffusers_pipeline_offload( return None if LAYERWISE_OFFLOAD in assignments.values(): raise ComponentResidencyError( - "--component-residency layerwise-offload requires the native SGLang " - "backend" + "--component-residency layerwise-offload requires the native SGLang backend" ) pipeline_mode = assignments.get(LAYERWISE_OFFLOAD_ALL_COMPONENTS) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py index 48c6619cc..220ad3b96 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py @@ -729,9 +729,9 @@ class LayerwiseOffloadManager: # below swaps that same object's storage for a (1,) # placeholder, leaving the placeholder in the store. # Keep an independent tensor over the mapped storage. - self._mapped_cpu_weights[layer_idx][ - name - ] = local_weight.detach().view_as(local_weight) + self._mapped_cpu_weights[layer_idx][name] = ( + local_weight.detach().view_as(local_weight) + ) self._weight_metadata[layer_idx][name] = { "dtype": local_weight.dtype, "shape": tuple(local_weight.shape), diff --git a/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py b/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py index c94191b29..240f25a88 100644 --- a/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py +++ b/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py @@ -118,10 +118,12 @@ class PerFrameAttentionPooling(nn.Module): """ B, L, D = x.shape T, H, W = grid_size - assert ( - D == self.dim - ), f"Input dimension D={D} does not match module dim={self.dim}" - assert L == T * H * W, f"Flattened length L={L} does not match T*H*W={T*H*W}" + assert D == self.dim, ( + f"Input dimension D={D} does not match module dim={self.dim}" + ) + assert L == T * H * W, ( + f"Flattened length L={L} does not match T*H*W={T * H * W}" + ) S = H * W x_bt_s_d = x.view(B, T, S, D).contiguous().view(B * T, S, D) # [B*T, S, D] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py index 6f3fc4494..30710d6d8 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py @@ -413,9 +413,10 @@ class CausalWanTransformerBlock(nn.Module): norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 2. Cross-attention attn_output = self.attn2( @@ -427,9 +428,10 @@ class CausalWanTransformerBlock(nn.Module): norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py index 3b6ebd693..be94f0f7e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py @@ -228,9 +228,9 @@ class ErnieImageSelfAttention(nn.Module): tp_size = get_tp_world_size() self.num_local_heads = num_heads // tp_size - assert ( - num_heads % tp_size == 0 - ), f"num_heads ({num_heads}) must be divisible by tp_size ({tp_size})" + assert num_heads % tp_size == 0, ( + f"num_heads ({num_heads}) must be divisible by tp_size ({tp_size})" + ) self.to_q = ColumnParallelLinear( hidden_size, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index 97fd79a73..56dbaa506 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py @@ -531,9 +531,9 @@ class GlmImageAttention(torch.nn.Module): self.out_dim = out_dim if out_dim is not None else query_dim tp_size = get_tp_world_size() - assert ( - self.heads % tp_size == 0 - ), f"heads ({self.heads}) must be divisible by tp_size ({tp_size})" + assert self.heads % tp_size == 0, ( + f"heads ({self.heads}) must be divisible by tp_size ({tp_size})" + ) self.num_local_heads = self.heads // tp_size self.num_local_kv_heads = self.num_local_heads @@ -673,9 +673,9 @@ class GlmImageAttention(torch.nn.Module): # 4. Attention if attention_mask is not None: text_attn_mask = attention_mask - assert ( - text_attn_mask.dim() == 2 - ), "the shape of text_attn_mask should be (batch_size, text_seq_length)" + assert text_attn_mask.dim() == 2, ( + "the shape of text_attn_mask should be (batch_size, text_seq_length)" + ) hidden_states = self.attn( query, key, value, num_replicated_prefix=text_seq_length ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index acdb6c7f9..9e064ccbf 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py @@ -960,9 +960,9 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi teacache_params = forward_batch.teacache_params assert teacache_params is not None, "teacache_params is not initialized" - assert isinstance( - teacache_params, TeaCacheParams - ), "teacache_params is not a TeaCacheParams" + assert isinstance(teacache_params, TeaCacheParams), ( + "teacache_params is not a TeaCacheParams" + ) num_inference_steps = forward_batch.num_inference_steps teache_thresh = teacache_params.teacache_thresh @@ -1006,9 +1006,7 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi img_mod2_shift, img_mod2_scale, img_mod2_gate, - ) = ( - self.double_blocks[0].img_mod(vec_).chunk(6, dim=-1) - ) + ) = self.double_blocks[0].img_mod(vec_).chunk(6, dim=-1) normed_inp = self.double_blocks[0].img_attn_norm.norm(inp) modulated_inp = modulate(normed_inp, shift=img_mod1_shift, scale=img_mod1_scale) if self.cnt == 0 or self.cnt == num_inference_steps - 1: @@ -1023,9 +1021,9 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi 9.61237896e-02, ] rescale_func = np.poly1d(coefficients) - assert ( - self.previous_modulated_input is not None - ), "previous_modulated_input is not initialized" + assert self.previous_modulated_input is not None, ( + "previous_modulated_input is not initialized" + ) self.accumulated_rel_l1_distance += rescale_func( ( (modulated_inp - self.previous_modulated_input).abs().mean() diff --git a/python/sglang/multimodal_gen/runtime/models/dits/krea2.py b/python/sglang/multimodal_gen/runtime/models/dits/krea2.py index eeff31c4e..2c0dc91a1 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/krea2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/krea2.py @@ -264,9 +264,9 @@ class Attention(nn.Module): # Parameter names match the released checkpoint (to_q/to_k/to_v/to_gate, # norm_q/norm_k, to_out.0) so the checkpoint loads with an identity mapping. tp = get_tp_world_size() - assert ( - self.heads % tp == 0 and self.kvheads % tp == 0 - ), f"heads={self.heads}, kvheads={self.kvheads} must be divisible by tp={tp}" + assert self.heads % tp == 0 and self.kvheads % tp == 0, ( + f"heads={self.heads}, kvheads={self.kvheads} must be divisible by tp={tp}" + ) self.local_heads = self.heads // tp self.local_kvheads = self.kvheads // tp diff --git a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py index b01a48ab7..012cd4781 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py @@ -427,9 +427,9 @@ class LingBotVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi hidden_size = config.hidden_size num_attention_heads = config.num_attention_heads head_dim = hidden_size // num_attention_heads - assert head_dim == sum( - config.axes_dims - ), f"head_dim {head_dim} != sum(axes_dims) {sum(config.axes_dims)}" + assert head_dim == sum(config.axes_dims), ( + f"head_dim {head_dim} != sum(axes_dims) {sum(config.axes_dims)}" + ) mlp_only_layers = tuple(config.mlp_only_layers) self.hidden_size = hidden_size diff --git a/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py b/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py index eb85876cd..986b793c7 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py @@ -268,9 +268,9 @@ class _LongCatJointAttention(nn.Module): super().__init__() tp_size = get_tp_world_size() self.num_local_heads = num_attention_heads // tp_size - assert ( - num_attention_heads % tp_size == 0 - ), f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})" + assert num_attention_heads % tp_size == 0, ( + f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})" + ) self.head_dim = attention_head_dim inner_dim = num_attention_heads * attention_head_dim @@ -429,9 +429,9 @@ class _LongCatSingleAttention(nn.Module): super().__init__() tp_size = get_tp_world_size() self.num_local_heads = num_attention_heads // tp_size - assert ( - num_attention_heads % tp_size == 0 - ), f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})" + assert num_attention_heads % tp_size == 0, ( + f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})" + ) self.head_dim = attention_head_dim inner_dim = num_attention_heads * attention_head_dim diff --git a/python/sglang/multimodal_gen/runtime/models/dits/longlive2.py b/python/sglang/multimodal_gen/runtime/models/dits/longlive2.py index 19b56c138..c47acabcd 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/longlive2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/longlive2.py @@ -128,9 +128,10 @@ class LongLive2CausalWanTransformerBlock(CausalWanTransformerBlock): norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) attn_output = self._cross_attn_with_cache( norm_hidden_states, @@ -140,9 +141,10 @@ class LongLive2CausalWanTransformerBlock(CausalWanTransformerBlock): norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 50f572b3d..a72bcd81e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -129,9 +129,12 @@ def _diffusers_h3_checkpoint( continue merge_dim = 1 if target_name.endswith((".qweight", ".qzeros", ".scales")) else 0 - yield target_name, torch.cat( - [pending[target_name][index] for index in range(merge_count)], - dim=merge_dim, + yield ( + target_name, + torch.cat( + [pending[target_name][index] for index in range(merge_count)], + dim=merge_dim, + ), ) del pending[target_name] @@ -1945,8 +1948,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): ): if value % tp_size: raise ValueError( - f"MiniMax H3 {name}={value} must be divisible by " - f"TP size {tp_size}." + f"MiniMax H3 {name}={value} must be divisible by TP size {tp_size}." ) @staticmethod @@ -2005,8 +2007,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): adaln_cache_path is not None or adaln_weight_files is not None ): raise ValueError( - "MiniMax H3 pruned curve checkpoints cannot use a separate " - "AdaLN cache" + "MiniMax H3 pruned curve checkpoints cannot use a separate AdaLN cache" ) self._adaln_precomputed = ( adaln_cache_path is not None or adaln_weight_files is not None diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 944b67161..c4167020f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -762,9 +762,9 @@ class QwenImageCrossAttention(nn.Module): self.inner_kv_dim = self.inner_dim tp_size = get_tp_world_size() - assert ( - self.num_heads % tp_size == 0 - ), f"num_heads ({self.num_heads}) must be divisible by tp_size ({tp_size})" + assert self.num_heads % tp_size == 0, ( + f"num_heads ({self.num_heads}) must be divisible by tp_size ({tp_size})" + ) self.local_num_heads = self.num_heads // tp_size self._unquantized_added_qkv_is_packed = False diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana.py b/python/sglang/multimodal_gen/runtime/models/dits/sana.py index 2af9923b2..f1e21cff8 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana.py @@ -494,7 +494,6 @@ class SanaTransformerBlock(nn.Module): class SanaTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): - _fsdp_shard_conditions = [ lambda n, m: isinstance(m, SanaTransformerBlock), ] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py index 1d9b1cc30..cad7a50c4 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py @@ -1871,9 +1871,9 @@ class BidirectionalGDNUCPESinglePathLiteLA(nn.Module): ) -> None: super().__init__() out_dim = heads * head_dim - assert ( - out_dim == in_dim - ), f"in_dim ({in_dim}) must equal heads*head_dim ({out_dim})" + assert out_dim == in_dim, ( + f"in_dim ({in_dim}) must equal heads*head_dim ({out_dim})" + ) self.in_dim = in_dim self.out_dim = out_dim self.heads = heads @@ -2219,9 +2219,7 @@ class BidirectionalGDNUCPESinglePathLiteLA(nn.Module): if beta.ndim == 3 and beta.shape != (B, heads, T): return f"requires beta shape {(B, heads, T)}, got {tuple(beta.shape)}" if beta.ndim == 4 and beta.shape != (B, heads, T, S): - return ( - f"requires beta shape {(B, heads, T, S)}, " f"got {tuple(beta.shape)}" - ) + return f"requires beta shape {(B, heads, T, S)}, got {tuple(beta.shape)}" if decay.shape != (B, heads, T): return f"requires decay shape {(B, heads, T)}, got {tuple(decay.shape)}" if head_dim > 128: @@ -2252,8 +2250,7 @@ class BidirectionalGDNUCPESinglePathLiteLA(nn.Module): if precheck_reason is not None: if self.gdn_backend == "triton": raise RuntimeError( - "SANA-WM Triton camera GDN backend unavailable: " - f"{precheck_reason}" + f"SANA-WM Triton camera GDN backend unavailable: {precheck_reason}" ) return None diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index b80b22702..898add578 100755 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -679,9 +679,10 @@ class WanTransformerBlock(nn.Module): query = q_sbhd.view(query_shape) key = k_sbhd.view(key_shape) else: - query, key = _apply_rotary_emb( - query, cos, sin, is_neox_style=False - ), _apply_rotary_emb(key, cos, sin, is_neox_style=False) + query, key = ( + _apply_rotary_emb(query, cos, sin, is_neox_style=False), + _apply_rotary_emb(key, cos, sin, is_neox_style=False), + ) attn_output = self.attn1(query, key, value) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) @@ -693,9 +694,10 @@ class WanTransformerBlock(nn.Module): norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 2. Cross-attention attn_output = self.attn2( @@ -704,9 +706,10 @@ class WanTransformerBlock(nn.Module): norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) @@ -930,9 +933,10 @@ class WanTransformerBlock_VSA(nn.Module): query = q_sbhd.view(query_shape) key = k_sbhd.view(key_shape) else: - query, key = _apply_rotary_emb( - query, cos, sin, is_neox_style=False - ), _apply_rotary_emb(key, cos, sin, is_neox_style=False) + query, key = ( + _apply_rotary_emb(query, cos, sin, is_neox_style=False), + _apply_rotary_emb(key, cos, sin, is_neox_style=False), + ) attn_output = self.attn1(query, key, value, gate_compress=gate_compress) attn_output = attn_output.flatten(2) @@ -943,9 +947,10 @@ class WanTransformerBlock_VSA(nn.Module): norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 2. Cross-attention attn_output = self.attn2( @@ -954,9 +959,10 @@ class WanTransformerBlock_VSA(nn.Module): norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index a3fd410d9..3519b2583 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -255,12 +255,12 @@ class ZImageAttention(nn.Module): self.enable_zimage_qk_fusion = quant_config is None tp_size = get_tp_world_size() - assert ( - num_heads % tp_size == 0 - ), f"num_heads {num_heads} must be divisible by tp world size {tp_size}" - assert ( - num_kv_heads % tp_size == 0 - ), f"num_kv_heads {num_kv_heads} must be divisible by tp world size {tp_size}" + assert num_heads % tp_size == 0, ( + f"num_heads {num_heads} must be divisible by tp world size {tp_size}" + ) + assert num_kv_heads % tp_size == 0, ( + f"num_kv_heads {num_kv_heads} must be divisible by tp world size {tp_size}" + ) self.local_num_heads = num_heads // tp_size self.local_num_kv_heads = num_kv_heads // tp_size @@ -704,9 +704,9 @@ class RopeEmbedder: self.theta = theta self.axes_dims = axes_dims self.axes_lens = axes_lens - assert len(axes_dims) == len( - axes_lens - ), "axes_dims and axes_lens must have the same length" + assert len(axes_dims) == len(axes_lens), ( + "axes_dims and axes_lens must have the same length" + ) self.cos_cached = None self.sin_cached = None diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py index 75df48a7f..ce42af1d8 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py @@ -36,7 +36,6 @@ def _srt_clip_param_name(name: str) -> str: class CLIPTextTransformer(nn.Module): - def __init__( self, config: CLIPTextConfig, @@ -145,7 +144,6 @@ class CLIPTextTransformer(nn.Module): class CLIPTextModel(TextEncoder): - def __init__( self, config: CLIPTextConfig, @@ -254,7 +252,6 @@ class CLIPTextModelWithProjection(CLIPTextModel): class CLIPVisionTransformer(nn.Module): - def __init__( self, config: CLIPVisionConfig, diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py b/python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py index dc355f46e..f5548aa54 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py @@ -209,7 +209,7 @@ def build_image_encoder(config): elif config["type"] == "DinoImageEncoderMV": return DinoImageEncoderMV(**config["kwargs"]) else: - raise ValueError(f'Unknown image encoder type: {config["type"]}') + raise ValueError(f"Unknown image encoder type: {config['type']}") class DualImageEncoder(nn.Module, LayerwiseOffloadableModuleMixin): diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/llama.py b/python/sglang/multimodal_gen/runtime/models/encoders/llama.py index 27ad1a518..5debd9735 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/llama.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/llama.py @@ -58,7 +58,6 @@ from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder class LlamaMLP(nn.Module): - def __init__( self, hidden_size: int, @@ -86,8 +85,7 @@ class LlamaMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -99,7 +97,6 @@ class LlamaMLP(nn.Module): class LlamaAttention(nn.Module): - def __init__( self, config: LlamaConfig, @@ -218,7 +215,6 @@ class LlamaAttention(nn.Module): class LlamaDecoderLayer(nn.Module): - def __init__( self, config: LlamaConfig, @@ -296,7 +292,6 @@ class LlamaDecoderLayer(nn.Module): class LlamaModel(TextEncoder): - def __init__( self, config: LlamaConfig, diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py index ab8b74f0e..c995291ca 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py @@ -1329,7 +1329,9 @@ class Qwen2_5_VLForConditionalGeneration(TextEncoder): eos_token_ids = ( [] if eos_token_id is None - else [eos_token_id] if isinstance(eos_token_id, int) else list(eos_token_id) + else [eos_token_id] + if isinstance(eos_token_id, int) + else list(eos_token_id) ) if pad_token_id is None: raise ValueError("pad_token_id must be set for generation") diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen_vl_rope.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen_vl_rope.py index 85077a925..4e96370fd 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen_vl_rope.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/qwen_vl_rope.py @@ -38,8 +38,7 @@ def apply_qwen_vl_text_rope( """Apply three-axis MRoPE to batched attention tensors.""" if query.ndim != 4 or key.ndim != 4: raise ValueError( - "Qwen-VL query and key must have shape " - "[batch, heads, sequence, head_dim]" + "Qwen-VL query and key must have shape [batch, heads, sequence, head_dim]" ) if position_ids.ndim != 3 or position_ids.shape[0] != 3: raise ValueError( diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/t5.py b/python/sglang/multimodal_gen/runtime/models/encoders/t5.py index 3834fb69e..7eb86621e 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/t5.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/t5.py @@ -72,7 +72,6 @@ class AttentionMetadata: class T5DenseActDense(nn.Module): - def __init__( self, config: T5Config, quant_config: QuantizationConfig | None = None ): @@ -98,7 +97,6 @@ class T5DenseActDense(nn.Module): class T5DenseGatedActDense(nn.Module): - def __init__( self, config: T5Config, quant_config: QuantizationConfig | None = None ): @@ -138,7 +136,6 @@ class T5DenseGatedActDense(nn.Module): class T5LayerFF(nn.Module): - def __init__( self, config: T5Config, quant_config: QuantizationConfig | None = None ): @@ -161,7 +158,6 @@ class T5LayerFF(nn.Module): # T5 has attn_bias and does not use softmax scaling class T5MultiHeadAttention(nn.Module): - def __init__(self) -> None: super().__init__() @@ -178,7 +174,6 @@ class T5MultiHeadAttention(nn.Module): class T5Attention(nn.Module): - def __init__( self, config: T5Config, @@ -378,7 +373,6 @@ class T5Attention(nn.Module): class T5LayerSelfAttention(nn.Module): - def __init__( self, config, @@ -416,7 +410,6 @@ class T5LayerSelfAttention(nn.Module): class T5LayerCrossAttention(nn.Module): - def __init__( self, config, quant_config: QuantizationConfig | None = None, prefix: str = "" ): @@ -445,7 +438,6 @@ class T5LayerCrossAttention(nn.Module): class T5Block(nn.Module): - def __init__( self, config: T5Config, @@ -505,7 +497,6 @@ class T5Block(nn.Module): class T5Stack(nn.Module): - def __init__( self, config: T5Config, diff --git a/python/sglang/multimodal_gen/runtime/models/registry.py b/python/sglang/multimodal_gen/runtime/models/registry.py index cb44d9fb6..0330acc75 100644 --- a/python/sglang/multimodal_gen/runtime/models/registry.py +++ b/python/sglang/multimodal_gen/runtime/models/registry.py @@ -177,7 +177,6 @@ class _ModelInfo: class _BaseRegisteredModel(ABC): - @abstractmethod def inspect_model_cls(self) -> _ModelInfo: raise NotImplementedError @@ -231,7 +230,7 @@ def _run_in_subprocess(fn: Callable[[], _T]) -> _T: except Exception as e: # wrap raised exception to provide more information raise RuntimeError( - f"Error raised in subprocess:\n" f"{returned.stderr.decode()}" + f"Error raised in subprocess:\n{returned.stderr.decode()}" ) from e with open(output_filepath, "rb") as f: diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py index 5d9ea035d..c78d3ddf7 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py @@ -194,9 +194,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): assert num_inference_steps is not None sigmas = np.linspace( self.sigma_max, self.sigma_min, num_inference_steps + 1 - ).copy()[ - :-1 - ] # pyright: ignore + ).copy()[:-1] # pyright: ignore if self.config.use_dynamic_shifting: assert mu is not None @@ -217,9 +215,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): ) timesteps = sigmas * self.config.num_train_timesteps - sigmas = np.concatenate([sigmas, [sigma_last]]).astype( - np.float32 - ) # pyright: ignore + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) # pyright: ignore self.sigmas = torch.from_numpy(sigmas).to(device=device) self.timesteps = torch.from_numpy(timesteps).to( @@ -476,18 +472,14 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): if self.predict_x0: x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 if D1s is not None: - pred_res = torch.einsum( - "k,bkc...->bc...", rhos_p, D1s - ) # pyright: ignore + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) # pyright: ignore else: pred_res = 0 x_t = x_t_ - alpha_t * B_h * pred_res else: x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 if D1s is not None: - pred_res = torch.einsum( - "k,bkc...->bc...", rhos_p, D1s - ) # pyright: ignore + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) # pyright: ignore else: pred_res = 0 x_t = x_t_ - sigma_t * B_h * pred_res diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_minimax_h3_euler_ancestral.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_minimax_h3_euler_ancestral.py index 6c067b0a7..81629fe14 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_minimax_h3_euler_ancestral.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_minimax_h3_euler_ancestral.py @@ -138,8 +138,7 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter: def __init__(self, **config: Any) -> None: if config: raise ValueError( - f"{type(self).__name__} does not accept config fields: " - f"{sorted(config)}" + f"{type(self).__name__} does not accept config fields: {sorted(config)}" ) def set_shift(self, _flow_shift: float) -> None: diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/common.py b/python/sglang/multimodal_gen/runtime/models/vaes/common.py index cbe9b6af1..28b032153 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/common.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/common.py @@ -749,7 +749,6 @@ class ParallelTiledVAE(ABC, nn.Module, LayerwiseOffloadableModuleMixin): # adapted from https://github.com/huggingface/diffusers/blob/e7ffeae0a191f710881d1fbde00cd6ff025e81f2/src/diffusers/models/autoencoders/vae.py#L691 class DiagonalGaussianDistribution: - def __init__(self, parameters: torch.Tensor, deterministic: bool = False): self.parameters = parameters self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py index 1de3203c4..8ad0a992a 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py @@ -475,7 +475,6 @@ class Transformer(nn.Module): class CrossAttentionDecoder(nn.Module): - def __init__( self, *, diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py index 5067a890e..93d424f71 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py @@ -100,7 +100,6 @@ def _apply_group_norm_silu( class HunyuanVAEAttention(nn.Module): - def __init__( self, in_channels, heads, dim_head, eps, norm_num_groups, bias ) -> None: @@ -167,7 +166,6 @@ class HunyuanVAEAttention(nn.Module): class HunyuanVideoCausalConv3d(nn.Module): - def __init__( self, in_channels: int, @@ -237,7 +235,6 @@ class HunyuanVideoCausalConv3d(nn.Module): class HunyuanVideoUpsampleCausal3D(nn.Module): - def __init__( self, in_channels: int, @@ -286,7 +283,6 @@ class HunyuanVideoUpsampleCausal3D(nn.Module): class HunyuanVideoDownsampleCausal3D(nn.Module): - def __init__( self, channels: int, @@ -309,7 +305,6 @@ class HunyuanVideoDownsampleCausal3D(nn.Module): class HunyuanVideoResnetBlockCausal3D(nn.Module): - def __init__( self, in_channels: int, @@ -361,7 +356,6 @@ class HunyuanVideoResnetBlockCausal3D(nn.Module): class HunyuanVideoMidBlock3D(nn.Module): - def __init__( self, in_channels: int, @@ -473,7 +467,6 @@ class HunyuanVideoMidBlock3D(nn.Module): class HunyuanVideoDownBlock3D(nn.Module): - def __init__( self, in_channels: int, @@ -537,7 +530,6 @@ class HunyuanVideoDownBlock3D(nn.Module): class HunyuanVideoUpBlock3D(nn.Module): - def __init__( self, in_channels: int, diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/processor.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/processor.py index 3fa00d9a7..9de0a673a 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/processor.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/processor.py @@ -54,7 +54,6 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name class VAEProcessor: - def __init__( self, *, diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py index 39d37376e..9bcbade80 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py @@ -241,8 +241,7 @@ class ViT3DDecoder(ViTBase): if dtype not in (torch.float16, torch.bfloat16): raise ValueError( - "MiniMax H3 decoder autocast weights require fp16 or bf16, " - f"got {dtype}" + f"MiniMax H3 decoder autocast weights require fp16 or bf16, got {dtype}" ) if self._autocast_linear_dtype == dtype: return 0 diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py index f631bdccd..c18bdf181 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py @@ -777,7 +777,6 @@ class WanMidBlock(nn.Module): class WanResidualDownBlock(nn.Module): - def __init__( self, in_dim, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py index cd3c9f39f..db7305d6a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py @@ -365,10 +365,22 @@ class ComfyUIFluxPipeline(LoRAPipeline, ComposedPipelineBase): k_tensor = qkv_tensor[hidden_size : 2 * hidden_size, :] v_tensor = qkv_tensor[2 * hidden_size : 3 * hidden_size, :] - yield f"single_transformer_blocks.{block_idx}.attn.to_q.{param_type}", q_tensor - yield f"single_transformer_blocks.{block_idx}.attn.to_k.{param_type}", k_tensor - yield f"single_transformer_blocks.{block_idx}.attn.to_v.{param_type}", v_tensor - yield f"single_transformer_blocks.{block_idx}.proj_mlp.{param_type}", mlp_tensor + yield ( + f"single_transformer_blocks.{block_idx}.attn.to_q.{param_type}", + q_tensor, + ) + yield ( + f"single_transformer_blocks.{block_idx}.attn.to_k.{param_type}", + k_tensor, + ) + yield ( + f"single_transformer_blocks.{block_idx}.attn.to_v.{param_type}", + v_tensor, + ) + yield ( + f"single_transformer_blocks.{block_idx}.proj_mlp.{param_type}", + mlp_tensor, + ) elif name == "final_layer.adaLN_modulation.1.weight": # ComfyUI: output order is [shift, scale] # AdaLayerNormContinuous: expects [scale, shift] diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ernie_image.py b/python/sglang/multimodal_gen/runtime/pipelines/ernie_image.py index 207653035..63f9bd533 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ernie_image.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ernie_image.py @@ -27,7 +27,6 @@ logger = init_logger(__name__) class ErnieImagePipeline(LoRAPipeline, ComposedPipelineBase): - pipeline_name = "ErnieImagePipeline" _required_config_modules = [ @@ -153,9 +152,9 @@ class ErnieImagePipeline(LoRAPipeline, ComposedPipelineBase): hasattr(pipeline_config, "text_encoder_extra_args") and pipeline_config.text_encoder_extra_args ): - pipeline_config.text_encoder_extra_args[0][ - "max_length" - ] = text_model_max_length + pipeline_config.text_encoder_extra_args[0]["max_length"] = ( + text_model_max_length + ) logger.info( "Set text encoder model_max_length=%d from tokenizer/tokenizer_config.json", text_model_max_length, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py index ef7fd15ce..75ea74255 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py @@ -19,7 +19,6 @@ from sglang.multimodal_gen.runtime.server_args import ServerArgs class HunyuanVideoPipeline(ComposedPipelineBase): - pipeline_name = "HunyuanVideoPipeline" _required_config_modules = [ diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 72266cfff..8d601c781 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -480,9 +480,9 @@ class ComposedPipelineBase(ABC): self._validate_direct_gpu_component_selection(model_index, server_args) # some sanity checks - assert ( - len(model_index) > 1 - ), "model_index.json must contain at least one pipeline module" + assert len(model_index) > 1, ( + "model_index.json must contain at least one pipeline module" + ) # In disagg mode, read HF config for skipped components (e.g., VAE) # so that update_model_arch + post_init can derive pipeline_config. diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/format_adapter.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/format_adapter.py index bc610cce1..557756d05 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/format_adapter.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/format_adapter.py @@ -337,8 +337,7 @@ def _convert_kohya_flux_via_diffusers( no_converter_warning="[LoRAFormatAdapter] No Kohya FLUX converter found.", success_info="[LoRAFormatAdapter] Converted Kohya FLUX LoRA using {name}", all_failed_warning=( - "[LoRAFormatAdapter] Kohya FLUX conversion failed; " - "last error: {last_err}" + "[LoRAFormatAdapter] Kohya FLUX conversion failed; last error: {last_err}" ), ) @@ -485,8 +484,7 @@ def _convert_ai_toolkit_flux_lora( sample = _sample_keys(final_out.keys(), 20) log.info( - "[LoRAFormatAdapter] after AI_TOOLKIT_FLUX conversion, " - "sample keys (<=20): %s", + "[LoRAFormatAdapter] after AI_TOOLKIT_FLUX conversion, sample keys (<=20): %s", ", ".join(sample), ) return final_out diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py index a6aed9cfa..68ab86c5d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py @@ -264,7 +264,9 @@ def scale_fused_sections( alpha = float( alpha_parts[index].item() if index in alpha_parts - else default_alpha if default_alpha is not None else rank + else default_alpha + if default_alpha is not None + else rank ) scale = alpha / rank weight = b_parts[index] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index 107c28e43..70d43d582 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py @@ -215,8 +215,7 @@ class PipelineStage(StageDedupMixin, ABC): return replace(use, target_dtype=target_dtype) return use raise ValueError( - f"{self.__class__.__name__} did not declare component use: " - f"{component_name}" + f"{self.__class__.__name__} did not declare component use: {component_name}" ) @contextmanager diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index ca65b357d..ef91dfbe9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -344,9 +344,9 @@ class DecodingStage(PipelineStage): # decode trajectory latents if needed if batch.return_trajectory_decoded: - assert ( - batch.trajectory_latents is not None - ), "batch should have trajectory latents" + assert batch.trajectory_latents is not None, ( + "batch should have trajectory latents" + ) # 1. Batch trajectory decoding to improve GPU utilization # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 03cc84fa0..69e61bc75 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -1521,9 +1521,9 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): # 1. Prepare latent inputs in the model's compute dtype. latent_model_input = ctx.latents.to(ctx.target_dtype) if batch.image_latent is not None: - assert ( - not server_args.pipeline_config.task_type == ModelTaskType.TI2V - ), "image latents should not be provided for TI2V task" + assert not server_args.pipeline_config.task_type == ModelTaskType.TI2V, ( + "image latents should not be provided for TI2V task" + ) latent_model_input = torch.cat( [latent_model_input, batch.image_latent], dim=1 ).to(ctx.target_dtype) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py index 808c55559..5cefe0a7f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py @@ -134,9 +134,9 @@ class DmdDenoisingStage(DenoisingStage): ], dim=2, ).to(target_dtype) - assert not torch.isnan( - latent_model_input - ).any(), "latent_model_input contains nan" + assert not torch.isnan(latent_model_input).any(), ( + "latent_model_input contains nan" + ) # Prepare inputs for transformer t_expand = t.repeat(latent_model_input.shape[0]) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index 56dd05431..4b7e29e68 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py @@ -359,7 +359,9 @@ class InputValidationStage(PipelineStage): neg_prompt_state = ( "not set" if batch.negative_prompt is None - else "empty" if batch.negative_prompt == "" else "set" + else "empty" + if batch.negative_prompt == "" + else "set" ) raise ValueError( f"Server was launched with --enable-cfg-parallel but this " @@ -446,8 +448,10 @@ class InputValidationStage(PipelineStage): result.add_check( "prompt_or_embeds", None, - lambda _: V.string_or_list_strings(batch.prompt) - or V.list_not_empty(batch.prompt_embeds), + lambda _: ( + V.string_or_list_strings(batch.prompt) + or V.list_not_empty(batch.prompt_embeds) + ), ) if server_args.pipeline_config.task_type != ModelTaskType.I2M: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py index 6b7f35cb6..ce07e10f8 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py @@ -342,9 +342,7 @@ class LatentPreparationStage(PipelineStage): server_args.pipeline_config.vae_config.use_temporal_scaling_frames ) if use_temporal_scaling_frames: - temporal_scale_factor = ( - server_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio - ) + temporal_scale_factor = server_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio latent_num_frames = (video_length - 1) // temporal_scale_factor + 1 return int(latent_num_frames) @@ -354,8 +352,10 @@ class LatentPreparationStage(PipelineStage): result.add_check( "prompt_or_embeds", None, - lambda _: V.string_or_list_strings(batch.prompt) - or V.list_not_empty(batch.prompt_embeds), + lambda _: ( + V.string_or_list_strings(batch.prompt) + or V.list_not_empty(batch.prompt_embeds) + ), ) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_of_tensors) result.add_check( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 220916090..02e382703 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -152,8 +152,7 @@ def _resize_center_crop_uint8_cthw( """Resize and center-crop ``uint8 [3, T, H, W]`` transfer frames.""" if frames.ndim != 4 or frames.shape[0] != 3: raise ValueError( - "Transfer frames must have shape [3, T, H, W], got " - f"{tuple(frames.shape)}" + f"Transfer frames must have shape [3, T, H, W], got {tuple(frames.shape)}" ) orig_h, orig_w = int(frames.shape[2]), int(frames.shape[3]) scale = max(width / orig_w, height / orig_h) @@ -178,8 +177,7 @@ def _pad_transfer_frames(video: torch.Tensor, target_frames: int) -> torch.Tenso """Pad ``[1, 3, T, H, W]`` with reflected temporal content.""" if video.ndim != 5 or video.shape[0] != 1 or video.shape[1] != 3: raise ValueError( - "Transfer video must have shape [1, 3, T, H, W], got " - f"{tuple(video.shape)}" + f"Transfer video must have shape [1, 3, T, H, W], got {tuple(video.shape)}" ) if target_frames <= 0: raise ValueError("Transfer target frame count must be positive") @@ -242,8 +240,7 @@ class Cosmos3ImagePreprocessStage(PipelineStage): stride = frames_per_chunk - conditional_frames if stride <= 0: raise ValueError( - "num_conditional_frames must be smaller than " - "num_video_frames_per_chunk" + "num_conditional_frames must be smaller than num_video_frames_per_chunk" ) remaining = total_frames - frames_per_chunk return 1 + math.ceil(remaining / stride), stride diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ernie_image_pe.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ernie_image_pe.py index 7eab9e803..a713d569d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ernie_image_pe.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ernie_image_pe.py @@ -19,7 +19,6 @@ logger = init_logger(__name__) class PromptEnhancementStage(PipelineStage): - def __init__(self, pe_model, pe_tokenizer): super().__init__() self.pe_model = pe_model diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py index e6c60bd96..162767f5d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py @@ -2719,8 +2719,8 @@ class LTX2DenoisingStage(DenoisingStage): def _get_negative_prompt_embeds_validator(self, batch: Req): """Allow either tensor or list negative prompt embeddings for LTX-2 CFG.""" - return ( - lambda x: (not batch.do_classifier_free_guidance) + return lambda x: ( + (not batch.do_classifier_free_guidance) or V.is_tensor(x) or V.list_not_empty(x) ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/latent_preparation_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/latent_preparation_av.py index c2066c078..1fafae6d9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/latent_preparation_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/latent_preparation_av.py @@ -40,9 +40,11 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage): result.add_check( "prompt_or_embeds", None, - lambda _: V.string_or_list_strings(batch.prompt) - or V.list_not_empty(batch.prompt_embeds) - or V.is_tensor(batch.prompt_embeds), + lambda _: ( + V.string_or_list_strings(batch.prompt) + or V.list_not_empty(batch.prompt_embeds) + or V.is_tensor(batch.prompt_embeds) + ), ) if isinstance(batch.prompt_embeds, list): diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py index e4d3d4e5d..8babe1856 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py @@ -55,8 +55,7 @@ class LTX2TextConnectorStage(PipelineStage): if prompt_embeds is None or prompt_attention_mask is None: raise ValueError( - "LTX2TextConnectorStage requires prompt embeddings and " - "attention mask." + "LTX2TextConnectorStage requires prompt embeddings and attention mask." ) if batch.do_classifier_free_guidance: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/condition_noise.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/condition_noise.py index 3a5e038a1..cf7f5aad7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/condition_noise.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/condition_noise.py @@ -95,8 +95,7 @@ def minimax_h3_imgvid_cond_noise_aug_rows( full_t = target_latent_t + imgvid_cond_num_frames if full_t < latent_t: raise ValueError( - f"condition latent_t {latent_t} exceeds the noise draw " - f"length {full_t}" + f"condition latent_t {latent_t} exceeds the noise draw length {full_t}" ) generator = torch.Generator(device="cpu").manual_seed(int(seed)) noise = torch.randn( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py index bd8d0e7d0..8df5e5a72 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py @@ -63,7 +63,7 @@ class MiniMaxH3ReleaseMetadata: partition = raw.get("partition") if partition not in {"fl2va", "ref2va"}: raise ValueError( - "model_index.json._minimax_h3.partition must be one of " "fl2va, ref2va" + "model_index.json._minimax_h3.partition must be one of fl2va, ref2va" ) tasks = _string_list(raw.get("tasks"), "model_index.json._minimax_h3.tasks") aliases = raw.get("task_aliases", {}) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/request_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/request_validation.py index d18f0d7f2..707be26ff 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/request_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/request_validation.py @@ -196,7 +196,7 @@ def _validate_conditions( MINIMAX_H3_CONDITION_ROLE_REFERENCE, ): raise ValueError( - f"{cpath}.role must be keyframe or reference, " f"got {role!r}" + f"{cpath}.role must be keyframe or reference, got {role!r}" ) cond_type = _require_str(cond.get("type"), f"{cpath}.type") try: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/latent_preparation.py index 86adbe91b..a3e9f67f1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/latent_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/latent_preparation.py @@ -145,8 +145,10 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): result.add_check( "prompt_or_embeds", None, - lambda _: V.string_or_list_strings(batch.prompt) - or V.list_not_empty(batch.prompt_embeds), + lambda _: ( + V.string_or_list_strings(batch.prompt) + or V.list_not_empty(batch.prompt_embeds) + ), ) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_of_tensors) result.add_check( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/timestep_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/timestep_preparation.py index 512fb5aa9..bc4712420 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/timestep_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/timestep_preparation.py @@ -120,8 +120,7 @@ class MiniMaxH3TimestepPreparationStage(PipelineStage): ) -> float: value = request_value source = ( - "request " - f"{'flow_shift' if modality == 'video' else 'audio_flow_shift'}" + f"request {'flow_shift' if modality == 'video' else 'audio_flow_shift'}" ) if value is None and model_scales is not None: value = model_scales.get(modality) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py index 32dfd26d1..dc4da9c46 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py @@ -70,8 +70,7 @@ class MiniMaxH3VideoModelAdapter: def validate_task_gate(self, task: Any, *, provided: bool) -> None: if not provided or task is None: raise ValueError( - "task is required for MiniMax H3; supported tasks: " - "fl2va, ref2va, t2va" + "task is required for MiniMax H3; supported tasks: fl2va, ref2va, t2va" ) if not isinstance(task, str): raise ValueError("task must be a non-empty string for MiniMax H3") diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py index 32c3285e0..44aee43d7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py @@ -88,7 +88,7 @@ def _activation_time(P: float, delta: float) -> float: denom = P * (1.0 + P - delta) if denom <= 0 or delta >= 1.0 + P: raise ValueError( - f"delta={delta} >= 1+P={1+P:.4f}; criterion trivially satisfied." + f"delta={delta} >= 1+P={1 + P:.4f}; criterion trivially satisfied." ) return 1.0 / (1.0 + math.sqrt(delta / denom)) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py index 8662f2cfb..733968065 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py @@ -480,9 +480,11 @@ class TextEncodingStage(ConditionEncodingStage): result.add_check( "negative_prompt", batch.negative_prompt, - lambda x: not batch.do_classifier_free_guidance - or V.string_not_none(x) - or isinstance(x, str), + lambda x: ( + not batch.do_classifier_free_guidance + or V.string_not_none(x) + or isinstance(x, str) + ), ) result.add_check( "do_classifier_free_guidance", @@ -868,8 +870,10 @@ class TextEncodingStage(ConditionEncodingStage): result.add_check( "negative_prompt_embeds", batch.negative_prompt_embeds, - lambda x: not batch.do_classifier_free_guidance - or V.list_of_tensors_with_min_dims(x, 2), + lambda x: ( + not batch.do_classifier_free_guidance + or V.list_of_tensors_with_min_dims(x, 2) + ), ) if batch.debug: logger.debug(f"{batch.prompt_embeds=}") diff --git a/python/sglang/multimodal_gen/runtime/platforms/__init__.py b/python/sglang/multimodal_gen/runtime/platforms/__init__.py index 91d7c6a9c..e47eeac80 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/__init__.py +++ b/python/sglang/multimodal_gen/runtime/platforms/__init__.py @@ -247,7 +247,7 @@ def resolve_current_platform_cls_qualname() -> str: if platform_cls_qualname is not None: return platform_cls_qualname - raise RuntimeError("No platform plugin found. Please check your " "installation.") + raise RuntimeError("No platform plugin found. Please check your installation.") _current_platform: Platform | None = None diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index c6efec9ca..04eb9baf2 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/python/sglang/multimodal_gen/runtime/platforms/interface.py @@ -392,8 +392,7 @@ class Platform: """ if cls.supported_quantization and quant not in cls.supported_quantization: raise ValueError( - f"{quant} quantization is currently not supported in " - f"{cls.device_name}." + f"{quant} quantization is currently not supported in {cls.device_name}." ) @classmethod diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py index 8b7545198..674fb1a87 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py @@ -37,7 +37,6 @@ def _kwargs_to_cpu(d: Any) -> Any: class RolloutDenoisingMixin: - def _maybe_prepare_rollout(self, batch: Req): """Prepare denoising loop for rollout.""" scheduler = batch.scheduler diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py index 3777ad4c4..aa52a25e6 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py @@ -88,9 +88,9 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): assert B == 1, "Generator must be a list if batch size is not 1" generator = [generator] else: - assert ( - len(generator) == B - ), "Generator list must have the same length as batch size" + assert len(generator) == B, ( + "Generator list must have the same length as batch size" + ) buffer = self._get_or_create_rollout_noise_buffer( rollout_session_data, rollout_session_data.latents_shape, device, dtype @@ -136,9 +136,9 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): debug_mode = bool(getattr(batch, "rollout_debug_mode", False)) if not log_prob_no_const and sde_type != "ode": - assert ( - noise_level > 0 - ), "True log-probability computation requires a non-zero noise level." + assert noise_level > 0, ( + "True log-probability computation requires a non-zero noise level." + ) dt = next_sigma - current_sigma @@ -230,9 +230,9 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): # Only enforce the "no full log-prob with ODE" constraint when the # user explicitly chose ODE globally. if sde_type == "ode": - assert ( - log_prob_no_const - ), "p_ode is always 0, true log_prob is meaningless, set rollout_log_prob_no_const to True to enable log_prob computation" + assert log_prob_no_const, ( + "p_ode is always 0, true log_prob is meaningless, set rollout_log_prob_no_const to True to enable log_prob computation" + ) else: raise ValueError(f"Unsupported sde_type: {sde_type}") diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index bd0bc7f73..41eef744e 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -2545,7 +2545,7 @@ class ServerArgs(DisaggServerArgsMixin): type=int, default=None, choices=[0, 1], - help="Quantize the attention sink too (1, default) " "or keep it bf16 (0).", + help="Quantize the attention sink too (1, default) or keep it bf16 (0).", ) parser.add_argument( "--kv-cache-quant-sink-keep", diff --git a/python/sglang/multimodal_gen/runtime/utils/distributed.py b/python/sglang/multimodal_gen/runtime/utils/distributed.py index 80afa0721..70f9a4e47 100644 --- a/python/sglang/multimodal_gen/runtime/utils/distributed.py +++ b/python/sglang/multimodal_gen/runtime/utils/distributed.py @@ -131,9 +131,9 @@ def generate_masked_orthogonal_rank_groups( idx = [(index // d) % s for s, d in zip(shape, stride)] # stride is a prefix_product result. And the value of stride[-1] # is not used. - assert ( - sum([x * y for x, y in zip(idx, stride[:-1])]) == index - ), "idx {} with shape {} mismatch the return idx {}".format(index, shape, idx) + assert sum([x * y for x, y in zip(idx, stride[:-1])]) == index, ( + "idx {} with shape {} mismatch the return idx {}".format(index, shape, idx) + ) return idx masked_shape = [s for s, m in zip(parallel_size, mask) if m] diff --git a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py index 10ccc11cc..18b80cbd8 100644 --- a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -454,13 +454,13 @@ def enable_trace_function_call(log_file_path: str, root_dir: str | None = None): def set_uvicorn_logging_configs(server_args=None): from uvicorn.config import LOGGING_CONFIG - LOGGING_CONFIG["formatters"]["default"][ - "fmt" - ] = "[%(asctime)s] %(levelprefix)s %(message)s" + LOGGING_CONFIG["formatters"]["default"]["fmt"] = ( + "[%(asctime)s] %(levelprefix)s %(message)s" + ) LOGGING_CONFIG["formatters"]["default"]["datefmt"] = "%Y-%m-%d %H:%M:%S" - LOGGING_CONFIG["formatters"]["access"][ - "fmt" - ] = '[%(asctime)s] %(levelprefix)s %(client_addr)s - "%(request_line)s" %(status_code)s' + LOGGING_CONFIG["formatters"]["access"]["fmt"] = ( + '[%(asctime)s] %(levelprefix)s %(client_addr)s - "%(request_line)s" %(status_code)s' + ) LOGGING_CONFIG["formatters"]["access"]["datefmt"] = "%Y-%m-%d %H:%M:%S" # Install access log path filter into LOGGING_CONFIG so it survives diff --git a/python/sglang/multimodal_gen/runtime/weights/source.py b/python/sglang/multimodal_gen/runtime/weights/source.py index d5412158c..3b246c615 100644 --- a/python/sglang/multimodal_gen/runtime/weights/source.py +++ b/python/sglang/multimodal_gen/runtime/weights/source.py @@ -263,8 +263,7 @@ def _select_named_file(candidates: tuple[str, ...], weight_name: str) -> str: if not basename_matches: raise FileNotFoundError(f"Requested weight {weight_name!r} was not found") raise ValueError( - f"Weight name {weight_name!r} matches multiple files: " - f"{list(basename_matches)}" + f"Weight name {weight_name!r} matches multiple files: {list(basename_matches)}" ) diff --git a/python/sglang/multimodal_gen/test/run_suite.py b/python/sglang/multimodal_gen/test/run_suite.py index aa8f33a02..c5aec5f82 100644 --- a/python/sglang/multimodal_gen/test/run_suite.py +++ b/python/sglang/multimodal_gen/test/run_suite.py @@ -213,9 +213,7 @@ def _format_standalone_estimate_snippet( suite: str, standalone_file: str, measured_full_test_time_s: float ) -> str: return ( - f'"{suite}": {{\n' - f' "{standalone_file}": {measured_full_test_time_s:.1f},\n' - f"}}" + f'"{suite}": {{\n "{standalone_file}": {measured_full_test_time_s:.1f},\n}}' ) @@ -228,13 +226,13 @@ def _print_missing_standalone_estimate_message( suite, standalone_file, measured_full_test_time_s ) logger.error( - f'\n{"=" * 60}\n' + f"\n{'=' * 60}\n" f'Add standalone estimate for suite "{suite}" and file "{standalone_file}":\n\n' f"File: python/sglang/multimodal_gen/test/run_suite.py\n\n" f"Current partition used fallback estimate: " f"{DEFAULT_STANDALONE_EST_TIME_SECONDS:.1f}s\n\n" f"{snippet}\n" - f'{"=" * 60}\n' + f"{'=' * 60}\n" ) @@ -579,8 +577,7 @@ def _run_partition_assignment( else "" ) print( - f" - standalone: {standalone_file} " - f"({est_time:.1f}s{fallback_suffix})" + f" - standalone: {standalone_file} ({est_time:.1f}s{fallback_suffix})" ) for standalone_file in assignment.standalone_files: diff --git a/python/sglang/multimodal_gen/test/server/common/slack.py b/python/sglang/multimodal_gen/test/server/common/slack.py index 040df1ec9..4b59eed92 100644 --- a/python/sglang/multimodal_gen/test/server/common/slack.py +++ b/python/sglang/multimodal_gen/test/server/common/slack.py @@ -145,15 +145,13 @@ def upload_file_to_slack( title = ( "Original Image" if len(final_origin_paths) == 1 - else f"Original Image {i+1}" + else f"Original Image {i + 1}" ) uploads.append({"file": path, "title": title}) uploads.append({"file": file_path, "title": "Generated Image"}) - message = ( - f"*Case ID:* `{case_id}`\n" f"*Model:* `{model}`\n" f"*Prompt:* {prompt}" - ) + message = f"*Case ID:* `{case_id}`\n*Model:* `{model}`\n*Prompt:* {prompt}" client = WebClient(token=token, timeout=60) channel_id = "C0A02NDF7UY" diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index d1f524daf..c409ac675 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -244,14 +244,14 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: if needs_estimated_time and not is_baseline_generation_mode: _MISSING_ESTIMATED_TIME_CASES.add(case.id) logger.error( - f'\n{"=" * 60}\n' + f"\n{'=' * 60}\n" f'Add "estimated_full_test_time_s" to scenario "{case.id}":\n\n' f"File: {get_perf_baseline_update_path()}\n\n" f' "{case.id}": {{\n' f" ...\n" f' "estimated_full_test_time_s": {_measured_full_time:.1f}\n' f" }}\n" - f'{"=" * 60}\n' + f"{'=' * 60}\n" ) _print_case_log_separator(case.id, "END diffusion testcase") @@ -319,8 +319,7 @@ class DiffusionServerBase: tail = ctx.log_tail() message = ( - f"{case_id}: server process exited during generation " - f"(code {returncode})." + f"{case_id}: server process exited during generation (code {returncode})." ) if tail: message += f"\n\nServer log tail:\n{tail}" @@ -1184,7 +1183,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} ), "loaded_adapters should be a non-empty list" assert any( a.get("nickname") == "default" for a in lora_info["loaded_adapters"] - ), f"nickname 'default' not found in loaded_adapters: {lora_info['loaded_adapters']}" + ), ( + f"nickname 'default' not found in loaded_adapters: {lora_info['loaded_adapters']}" + ) logger.info("[LoRA E2E] list_loras returned expected LoRA adapters") logger.info("[LoRA E2E] All LoRA API E2E tests passed for %s", case.id) @@ -1226,9 +1227,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} json={"lora_nickname": "lora2", "lora_path": second_lora_path}, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 200 - ), f"set_lora to second adapter failed: {resp.text}" + assert resp.status_code == 200, ( + f"set_lora to second adapter failed: {resp.text}" + ) logger.info( "[LoRA Switch E2E] Verifying generation with second LoRA for %s", case.id @@ -1320,9 +1321,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} }, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 200 - ), f"set_lora with multiple adapters failed: {resp.text}" + assert resp.status_code == 200, ( + f"set_lora with multiple adapters failed: {resp.text}" + ) rid, _ = self._run_generation_with_server_watchdog( ctx, case.id, generate_fn, client ) @@ -1339,9 +1340,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} }, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 200 - ), f"set_lora with different strengths failed: {resp.text}" + assert resp.status_code == 200, ( + f"set_lora with different strengths failed: {resp.text}" + ) rid, _ = self._run_generation_with_server_watchdog( ctx, case.id, generate_fn, client ) @@ -1363,9 +1364,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} }, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 200 - ), f"set_lora with cached adapters failed: {resp.text}" + assert resp.status_code == 200, ( + f"set_lora with cached adapters failed: {resp.text}" + ) rid, _ = self._run_generation_with_server_watchdog( ctx, case.id, generate_fn, client ) @@ -1377,9 +1378,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} json={"lora_nickname": "default"}, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 200 - ), f"set_lora back to single adapter failed: {resp.text}" + assert resp.status_code == 200, ( + f"set_lora back to single adapter failed: {resp.text}" + ) rid, content = self._run_generation_with_server_watchdog( ctx, case.id, generate_fn, client ) @@ -1403,28 +1404,28 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} assert resp.status_code == 200, f"/v1/models failed: {resp.text}" data = resp.json() - assert ( - data["object"] == "list" - ), f"Expected object='list', got {data.get('object')}" + assert data["object"] == "list", ( + f"Expected object='list', got {data.get('object')}" + ) assert len(data["data"]) >= 1, "Expected at least one model in response" model = data["data"][0] assert "id" in model, "Model missing 'id' field" - assert ( - model["object"] == "model" - ), f"Expected object='model', got {model.get('object')}" - assert ( - model["id"] == case.server_args.model_path - ), f"Model ID mismatch: expected {case.server_args.model_path}, got {model['id']}" + assert model["object"] == "model", ( + f"Expected object='model', got {model.get('object')}" + ) + assert model["id"] == case.server_args.model_path, ( + f"Model ID mismatch: expected {case.server_args.model_path}, got {model['id']}" + ) # Verify extended diffusion-specific fields assert "num_gpus" in model, "Model missing 'num_gpus' field" assert "task_type" in model, "Model missing 'task_type' field" assert "dit_precision" in model, "Model missing 'dit_precision' field" assert "vae_precision" in model, "Model missing 'vae_precision' field" - assert ( - model["num_gpus"] == case.server_args.num_gpus - ), f"num_gpus mismatch: expected {case.server_args.num_gpus}, got {model['num_gpus']}" + assert model["num_gpus"] == case.server_args.num_gpus, ( + f"num_gpus mismatch: expected {case.server_args.num_gpus}, got {model['num_gpus']}" + ) expected_task_type = get_model_task_type_for_server_args(case.server_args).name assert model["task_type"] == expected_task_type, ( f"task_type mismatch: expected {expected_task_type}, " @@ -1466,9 +1467,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} assert resp.status_code == 404, f"Expected 404, got {resp.status_code}" error_data = resp.json() assert "error" in error_data, "404 response missing 'error' field" - assert ( - error_data["error"]["code"] == "model_not_found" - ), f"Incorrect error code: {error_data['error'].get('code')}" + assert error_data["error"]["code"] == "model_not_found", ( + f"Incorrect error code: {error_data['error'].get('code')}" + ) logger.info("[Models API] GET /v1/models/non_existent returns 404 as expected") logger.info("[Models API] All /v1/models tests passed for %s", case.id) @@ -1500,13 +1501,13 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} json=payload, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 400 - ), f"Expected 400 for T2V input_reference, got {resp.status_code}: {resp.text}" + assert resp.status_code == 400, ( + f"Expected 400 for T2V input_reference, got {resp.status_code}: {resp.text}" + ) detail = resp.json().get("detail", "") - assert ( - "input_reference is not supported" in detail - ), f"Unexpected error detail for T2V input_reference: {detail}" + assert "input_reference is not supported" in detail, ( + f"Unexpected error detail for T2V input_reference: {detail}" + ) def test_diffusion_generation( self, diff --git a/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/engine.py b/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/engine.py index 4dfb4ade1..cf2a79055 100644 --- a/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/engine.py +++ b/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/engine.py @@ -431,9 +431,9 @@ class AccuracyEngine: ).item() rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 logger.info("[%s] Rank %s CosSim=%.6f", name, rank, cos_sim) - assert ( - cos_sim > threshold - ), f"Accuracy failure in {name}: CosSim {cos_sim:.4f} < {threshold}" + assert cos_sim > threshold, ( + f"Accuracy failure in {name}: CosSim {cos_sim:.4f} < {threshold}" + ) @staticmethod def transfer_weights( diff --git a/python/sglang/multimodal_gen/test/single_test_file/test_ar_models.py b/python/sglang/multimodal_gen/test/single_test_file/test_ar_models.py index 4a4b0f27d..768aafc40 100644 --- a/python/sglang/multimodal_gen/test/single_test_file/test_ar_models.py +++ b/python/sglang/multimodal_gen/test/single_test_file/test_ar_models.py @@ -94,8 +94,7 @@ class ARCluster(DisaggCluster): ) except Exception as e: raise RuntimeError( - f"AR model failed to start for {self.name}. Log tail:\n" - f"{_tail_log(log)}" + f"AR model failed to start for {self.name}. Log tail:\n{_tail_log(log)}" ) from e def _launch_server_head(self) -> None: @@ -144,7 +143,6 @@ class ARCluster(DisaggCluster): class _ARTestBase(_DisaggTestBase): - @classmethod def setUpClass(cls) -> None: super(CustomTestCase, cls).setUpClass() diff --git a/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_disk.py b/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_disk.py index 1f60b0772..11d30af1d 100644 --- a/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_disk.py +++ b/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_disk.py @@ -217,9 +217,9 @@ def _compute_checksum_from_disk(model_path: str, module_name: str) -> str: """ local_path = maybe_download_model(model_path) weights_dir = os.path.join(local_path, module_name) - assert os.path.exists( - weights_dir - ), f"No weights dir for {module_name} in {local_path}" + assert os.path.exists(weights_dir), ( + f"No weights dir for {module_name} in {local_path}" + ) safetensors_files = _list_safetensors_files(weights_dir) assert safetensors_files, f"No safetensors files in {weights_dir}" @@ -323,9 +323,9 @@ class _UpdateWeightsApiMixin: json=payload, timeout=timeout, ) - assert ( - response.status_code == 200 - ), f"get_weights_checksum failed: {response.status_code} {response.text}" + assert response.status_code == 200, ( + f"get_weights_checksum failed: {response.status_code} {response.text}" + ) return response.json() def _assert_server_matches_model( @@ -542,9 +542,9 @@ class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin): and perturbed_checksums.get(name) != "not_found" and base_checksums.get(name) != "not_found" ) - assert ( - text_encoder_modules - ), "Expected at least one text encoder module checksum" + assert text_encoder_modules, ( + "Expected at least one text encoder module checksum" + ) # perturbed โ corrupted (should fail and rollback) rollback_targets = [_TRANSFORMER_MODULE, _VAE_MODULE] @@ -553,18 +553,18 @@ class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin): corrupted_vae_model_dir, target_modules=rollback_targets, ) - assert ( - status_code == 400 - ), f"Expected 400 on corrupted weights, got {status_code}" + assert status_code == 400, ( + f"Expected 400 on corrupted weights, got {status_code}" + ) assert not result.get("success", True) message = result.get("message", "") assert "rolled back" in message.lower() # The updater reports the first failing module in the error message. # With ordered target_modules=[transformer, vae], this makes the # failure point explicit: transformer is processed first, then vae fails. - assert ( - "Failed to update module 'vae'" in message - ), f"Expected vae to be the explicit failure point, got: {message}" + assert "Failed to update module 'vae'" in message, ( + f"Expected vae to be the explicit failure point, got: {message}" + ) rolled_back_checksums = self._get_weights_checksum(base_url) # 1) transformer: server == perturbed != base @@ -583,12 +583,12 @@ class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin): # 3) text encoder(s): server == base == perturbed for name in text_encoder_modules: - assert rolled_back_checksums.get(name) == perturbed_checksums.get( - name - ), f"Text encoder module '{name}' should stay equal to perturbed" - assert rolled_back_checksums.get(name) == base_checksums.get( - name - ), f"Text encoder module '{name}' should stay equal to base" + assert rolled_back_checksums.get(name) == perturbed_checksums.get(name), ( + f"Text encoder module '{name}' should stay equal to perturbed" + ) + assert rolled_back_checksums.get(name) == base_checksums.get(name), ( + f"Text encoder module '{name}' should stay equal to base" + ) class TestUpdateWeightsFromDiskWithOffload(_UpdateWeightsApiMixin): diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index b43a29fbb..727e51516 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -575,9 +575,9 @@ def validate_image_file( assert os.path.exists(file_path), f"Image file does not exist: {file_path}" # 2. Extension check - assert file_path.endswith( - f".{expected_ext}" - ), f"Expected .{expected_ext} extension, got: {file_path}" + assert file_path.endswith(f".{expected_ext}"), ( + f"Expected .{expected_ext} extension, got: {file_path}" + ) # 3. File size > 0 file_size = os.path.getsize(file_path) @@ -585,9 +585,9 @@ def validate_image_file( # 4. Filename validation actual_filename = os.path.basename(file_path) - assert ( - actual_filename == expected_filename - ), f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'" + assert actual_filename == expected_filename, ( + f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'" + ) # 5. Image format validation (magic bytes check based on expected format) with open(file_path, "rb") as f: @@ -603,12 +603,12 @@ def validate_image_file( if expected_width is not None and expected_height is not None: with Image.open(file_path) as img: width, height = img.size - assert ( - width == expected_width - ), f"Width mismatch: expected {expected_width}, got {width}" - assert ( - height == expected_height - ), f"Height mismatch: expected {expected_height}, got {height}" + assert width == expected_width, ( + f"Width mismatch: expected {expected_width}, got {width}" + ) + assert height == expected_height, ( + f"Height mismatch: expected {expected_height}, got {height}" + ) def _get_video_dimensions_from_metadata( @@ -707,9 +707,9 @@ def validate_video_file( # 4. Filename validation actual_filename = os.path.basename(file_path) - assert ( - actual_filename == expected_filename - ), f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'" + assert actual_filename == expected_filename, ( + f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'" + ) # 5. Video format validation (reuse is_mp4) with open(file_path, "rb") as f: @@ -719,12 +719,12 @@ def validate_video_file( # 6. Video dimension validation (using OpenCV) if expected_width is not None and expected_height is not None: actual_width, actual_height = get_video_dimensions(file_path) - assert ( - actual_width == expected_width - ), f"Video width mismatch: expected {expected_width}, got {actual_width}" - assert ( - actual_height == expected_height - ), f"Video height mismatch: expected {expected_height}, got {actual_height}" + assert actual_width == expected_width, ( + f"Video width mismatch: expected {expected_width}, got {actual_width}" + ) + assert actual_height == expected_height, ( + f"Video height mismatch: expected {expected_height}, got {actual_height}" + ) @dataclass(frozen=True) @@ -776,9 +776,9 @@ def probe_audio_stream(file_path: str) -> AudioStreamInfo: ) assert sample_rate > 0, f"Audio stream has invalid sample rate: {sample_rate}" assert channels > 0, f"Audio stream has invalid channel count: {channels}" - assert ( - math.isfinite(duration) and duration > 0 - ), f"Audio stream has invalid duration: {duration}" + assert math.isfinite(duration) and duration > 0, ( + f"Audio stream has invalid duration: {duration}" + ) return AudioStreamInfo(sample_rate, channels, duration) diff --git a/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py b/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py index 83daa8316..c44ad2d41 100644 --- a/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py +++ b/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py @@ -26,7 +26,7 @@ def _fake_backend_cls(enum, *, ring_capable=True): return SimpleNamespace( get_enum=lambda: enum, supports_ring_rotation=lambda: ring_capable, - get_impl_cls=lambda: (lambda **kwargs: f"{enum.name.lower()}_impl"), + get_impl_cls=lambda: lambda **kwargs: f"{enum.name.lower()}_impl", ) diff --git a/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py b/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py index 9f1077188..df61480f6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py @@ -101,8 +101,9 @@ class TestComponentQuantizationAdmission(unittest.TestCase): component_precisions={}, component_quantizations={}, component_weights_paths={}, - should_direct_gpu_weight_load_component=lambda component: component - == "vocoder", + should_direct_gpu_weight_load_component=lambda component: ( + component == "vocoder" + ), ) with self.assertRaisesRegex( @@ -117,8 +118,9 @@ class TestComponentQuantizationAdmission(unittest.TestCase): component_precisions={}, component_quantizations={}, component_weights_paths={}, - should_direct_gpu_weight_load_component=lambda component: component - == "audio_vae", + should_direct_gpu_weight_load_component=lambda component: ( + component == "audio_vae" + ), ) with self.assertRaisesRegex( diff --git a/python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py b/python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py index e00572b24..1383b38db 100644 --- a/python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py @@ -263,8 +263,10 @@ def test_remote_platform_video_gt_prefers_platform_sglang_before_default_officia monkeypatch.setattr( test_utils, "_remote_file_exists", - lambda url: url.startswith(sglang_platform_prefix) - or url.startswith(official_default_prefix), + lambda url: ( + url.startswith(sglang_platform_prefix) + or url.startswith(official_default_prefix) + ), ) files = test_utils._find_remote_consistency_gt_files( diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py index 70ca35733..ffd1788b7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py @@ -48,7 +48,6 @@ def _prepare(serving, batch, explicit_shift): class TestPrepareRolloutRequestScheduler(unittest.TestCase): - def test_inherits_serving_grid_without_explicit_shift(self): serving = _serving_scheduler() batch = _rollout_batch() @@ -126,7 +125,6 @@ class _FusedParamModule(torch.nn.Module): class TestWeightsUpdaterFusedParams(unittest.TestCase): - def test_merge_index_reaches_weight_loader_as_shard_id(self): module = _FusedParamModule() calls = [] diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py index 1ce940dea..ebc591f02 100644 --- a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py +++ b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py @@ -58,15 +58,19 @@ class TestBCGTPGraphCapture(CustomTestCase): custom_allreduce_cls = MagicMock(return_value=expected) group = SimpleNamespace(cpu_group=object(), device=torch.device("cuda:0")) - with patch( - "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda", - return_value=True, - ), patch( - "sglang.srt.distributed.device_communicators.custom_all_reduce.dispatch_custom_allreduce", - return_value=custom_allreduce_cls, - ) as dispatch, patch( - "sglang.srt.distributed.device_communicators.custom_all_reduce_v2.CustomAllReduceV2", - custom_allreduce_cls, + with ( + patch( + "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda", + return_value=True, + ), + patch( + "sglang.srt.distributed.device_communicators.custom_all_reduce.dispatch_custom_allreduce", + return_value=custom_allreduce_cls, + ) as dispatch, + patch( + "sglang.srt.distributed.device_communicators.custom_all_reduce_v2.CustomAllReduceV2", + custom_allreduce_cls, + ), ): GroupCoordinator._init_srt_custom_allreduce(group) @@ -83,12 +87,15 @@ class TestBCGTPGraphCapture(CustomTestCase): custom_allreduce_cls = MagicMock(return_value=expected) group = SimpleNamespace(cpu_group=object(), device=object()) - with patch( - "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda", - return_value=False, - ), patch( - "sglang.srt.distributed.device_communicators.custom_all_reduce.CustomAllreduce", - custom_allreduce_cls, + with ( + patch( + "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda", + return_value=False, + ), + patch( + "sglang.srt.distributed.device_communicators.custom_all_reduce.CustomAllreduce", + custom_allreduce_cls, + ), ): GroupCoordinator._init_srt_custom_allreduce(group) @@ -120,10 +127,14 @@ class TestBCGTPGraphCapture(CustomTestCase): """Drive the CUDA branch of graph_capture() with a fake custom AR.""" group = SimpleNamespace(srt_custom_allreduce=custom_ar) ctx = GraphCaptureContext(MagicMock()) - with patch( - "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda_alike", - return_value=True, - ), patch("torch.cuda.stream"), patch("torch.cuda.current_stream"): + with ( + patch( + "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda_alike", + return_value=True, + ), + patch("torch.cuda.stream"), + patch("torch.cuda.current_stream"), + ): with GroupCoordinator.graph_capture(group, ctx) as yielded: events.append("body") return yielded, ctx @@ -157,12 +168,15 @@ class TestBCGTPGraphCapture(CustomTestCase): tp_group.graph_capture = _recording_context(events, "tp") runner = SimpleNamespace(_capture_stream=capture_stream) - with patch( - "sglang.multimodal_gen.runtime.distributed.parallel_state.model_parallel_is_initialized", - return_value=initialized, - ), patch( - "sglang.multimodal_gen.runtime.distributed.parallel_state.get_tp_group", - return_value=tp_group, + with ( + patch( + "sglang.multimodal_gen.runtime.distributed.parallel_state.model_parallel_is_initialized", + return_value=initialized, + ), + patch( + "sglang.multimodal_gen.runtime.distributed.parallel_state.get_tp_group", + return_value=tp_group, + ), ): with BaseBreakableCudaGraphRunner._tp_graph_capture(runner): events.append("body") @@ -190,12 +204,15 @@ class TestBCGTPGraphCapture(CustomTestCase): tp_group.graph_capture = _graph_capture runner = SimpleNamespace(_capture_stream=capture_stream) - with patch( - "sglang.multimodal_gen.runtime.distributed.parallel_state.model_parallel_is_initialized", - return_value=True, - ), patch( - "sglang.multimodal_gen.runtime.distributed.parallel_state.get_tp_group", - return_value=tp_group, + with ( + patch( + "sglang.multimodal_gen.runtime.distributed.parallel_state.model_parallel_is_initialized", + return_value=True, + ), + patch( + "sglang.multimodal_gen.runtime.distributed.parallel_state.get_tp_group", + return_value=tp_group, + ), ): with BaseBreakableCudaGraphRunner._tp_graph_capture(runner): pass @@ -232,20 +249,23 @@ class TestBCGTPGraphCapture(CustomTestCase): graph._segments = [] kwargs = {"hidden_states": torch.zeros(1)} - with patch.object( - BaseBreakableCudaGraphRunner, - "_tp_graph_capture", - _recording_context(events, "tp"), - ), patch.object( - runner_mod, "BreakableCUDAGraph", return_value=graph - ), patch.object( - runner_mod, - "enable_breakable_cuda_graph", - _recording_context(events, "bcg_enable"), - ), patch.object( - runner_mod, - "BreakableCUDAGraphCapture", - _recording_context(events, "bcg_capture"), + with ( + patch.object( + BaseBreakableCudaGraphRunner, + "_tp_graph_capture", + _recording_context(events, "tp"), + ), + patch.object(runner_mod, "BreakableCUDAGraph", return_value=graph), + patch.object( + runner_mod, + "enable_breakable_cuda_graph", + _recording_context(events, "bcg_enable"), + ), + patch.object( + runner_mod, + "BreakableCUDAGraphCapture", + _recording_context(events, "bcg_capture"), + ), ): runner._capture(kwargs, key=runner_mod._signature_kwargs(kwargs)) diff --git a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py b/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py index 5caeb91e3..4888bbf44 100644 --- a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py @@ -296,6 +296,6 @@ def test_the_forced_host_size_behaves_like_a_machine_of_that_size(monkeypatch): # and a larger pretend machine reports more room, same process monkeypatch.setenv("SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB", "64") larger = host_memory_budget.host_memory_available_bytes() - assert ( - abs((larger - available) - 32 * 1024**3) < 512 * 1024**2 - ), "the same process on a machine twice the size has one machine more room" + assert abs((larger - available) - 32 * 1024**3) < 512 * 1024**2, ( + "the same process on a machine twice the size has one machine more room" + ) diff --git a/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py index e0c4ef961..4f0861a81 100644 --- a/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py @@ -110,8 +110,9 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase): def test_native_only_quantized_architecture_does_not_fall_back(self): self.server_args.pipeline_config.native_only_components = ("image_encoder",) config = self._component_config("UnknownVisionModel", quantized=True) - with self._config_patch(config), self.assertRaises( - NativeComponentLoaderRequired + with ( + self._config_patch(config), + self.assertRaises(NativeComponentLoaderRequired), ): self._load() self.load_native.assert_not_called() @@ -121,8 +122,9 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase): "architectures": ["UnknownVisionModel"], "quantization_config": {"quant_method": "not-a-format"}, } - with self._config_patch(config), self.assertRaises( - ComponentCheckpointUnsupportedError + with ( + self._config_patch(config), + self.assertRaises(ComponentCheckpointUnsupportedError), ): self._load() self.load_native.assert_not_called() @@ -159,18 +161,22 @@ class TestImageEncoderNativeLoading(unittest.TestCase): ) loader = ImageEncoderLoader() - with mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.get_hf_config", - return_value=component_config, - ), mock.patch.object( - loader, - "resolve_native_transformers_model_class", - return_value=model_class, - ), mock.patch.object( - loader, - "target_device", - return_value=torch.device("cuda:0"), + with ( + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_hf_config", + return_value=component_config, + ), + mock.patch.object( + loader, + "resolve_native_transformers_model_class", + return_value=model_class, + ), + mock.patch.object( + loader, + "target_device", + return_value=torch.device("cuda:0"), + ), ): component = loader.load_native( "/model/image_encoder", @@ -211,16 +217,20 @@ class TestImageEncoderNativeLoading(unittest.TestCase): ) loader = ImageEncoderLoader() - with mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.get_hf_config", - return_value=component_config, - ), mock.patch.object( - loader, - "resolve_native_transformers_model_class", - return_value=model_class, - ), self.assertRaisesRegex( - ComponentCheckpointUnsupportedError, "requires resident placement" + with ( + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_hf_config", + return_value=component_config, + ), + mock.patch.object( + loader, + "resolve_native_transformers_model_class", + return_value=model_class, + ), + self.assertRaisesRegex( + ComponentCheckpointUnsupportedError, "requires resident placement" + ), ): loader.load_native( "/model/image_encoder", @@ -266,34 +276,42 @@ class TestImageEncoderNativeLoading(unittest.TestCase): ) loader = ImageEncoderLoader() - with mock.patch.object( - loader, - "load_customized", - side_effect=NativeComponentLoaderRequired("use Transformers"), - ), mock.patch.object( - loader, - "resolve_native_transformers_model_class", - return_value=model_class, - ), mock.patch.object( - loader, - "target_device", - return_value=torch.device("cuda:0"), - ), mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.get_hf_config", - return_value=component_config, - ), mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.current_platform.get_available_gpu_memory", - return_value=10.0, - ), mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.get_memory_usage_of_component", - return_value=0.0, - ), mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.format_component_residency", - return_value="resident", + with ( + mock.patch.object( + loader, + "load_customized", + side_effect=NativeComponentLoaderRequired("use Transformers"), + ), + mock.patch.object( + loader, + "resolve_native_transformers_model_class", + return_value=model_class, + ), + mock.patch.object( + loader, + "target_device", + return_value=torch.device("cuda:0"), + ), + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_hf_config", + return_value=component_config, + ), + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.current_platform.get_available_gpu_memory", + return_value=10.0, + ), + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_memory_usage_of_component", + return_value=0.0, + ), + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.format_component_residency", + return_value="resident", + ), ): component, _ = loader.load( "/model/image_encoder", diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index 83fce212b..b8320f96b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -1375,15 +1375,15 @@ def test_mapped_layers_ship_through_the_courier(tmp_path, monkeypatch): manager.prefetch_layer(0, non_blocking=True) assert 0 in manager._courier_inflight, "an async prefetch hands the layer over" - assert ( - 0 not in manager._gpu_layers - ), "the layer is not ready until its tensors are bound on this thread" + assert 0 not in manager._gpu_layers, ( + "the layer is not ready until its tensors are bound on this thread" + ) manager.prefetch_layer(0, non_blocking=False) assert 0 in manager._gpu_layers and not manager._courier_inflight - assert torch.equal( - model.blocks[0].weight.detach().cpu(), expected - ), "the bytes that went through the courier's slot must be the checkpoint's" + assert torch.equal(model.blocks[0].weight.detach().cpu(), expected), ( + "the bytes that went through the courier's slot must be the checkpoint's" + ) def test_the_courier_kill_switch_forces_the_synchronous_path(tmp_path, monkeypatch): @@ -1397,9 +1397,9 @@ def test_the_courier_kill_switch_forces_the_synchronous_path(tmp_path, monkeypat manager.release_all() manager.prefetch_layer(0, non_blocking=True) assert not manager._courier_inflight and manager._mapped_courier is None - assert ( - 0 in manager._gpu_layers - ), "with the courier disabled the direct synchronous path serves the layer" + assert 0 in manager._gpu_layers, ( + "with the courier disabled the direct synchronous path serves the layer" + ) def test_release_all_drains_the_courier(tmp_path, monkeypatch): diff --git a/python/sglang/multimodal_gen/test/unit/test_lora_format_adapter.py b/python/sglang/multimodal_gen/test/unit/test_lora_format_adapter.py index 7dec429dd..57a2c87fe 100644 --- a/python/sglang/multimodal_gen/test/unit/test_lora_format_adapter.py +++ b/python/sglang/multimodal_gen/test/unit/test_lora_format_adapter.py @@ -339,9 +339,9 @@ def main() -> None: class TestLoRAFormatAdapter: def test_lora_format_adapter_all_formats(self): results = _run_all_tests() - assert all( - r["pass"] for r in results - ), "At least one LoRA format adapter case failed" + assert all(r["pass"] for r in results), ( + "At least one LoRA format adapter case failed" + ) if __name__ == "__main__": diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py index 3f54206de..473bf04df 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py @@ -62,9 +62,7 @@ def test_unvalidated_decode_modes_are_rejected(mode): def test_vit_attention_uses_local_usp_backend_dispatch(): - module = ( - "sglang.multimodal_gen.runtime.models.vaes." "minimax_h3_video_vae.attention" - ) + module = "sglang.multimodal_gen.runtime.models.vaes.minimax_h3_video_vae.attention" with ( mock.patch(f"{module}.current_platform.is_cuda", return_value=True), mock.patch(f"{module}.USPAttention", autospec=True) as usp_attention, @@ -98,9 +96,7 @@ def test_audio_vae_attention_defaults_to_local_sdpa_and_allows_fa(): self.input_dtype = query.dtype return query - module = ( - "sglang.multimodal_gen.runtime.models.vaes." "minimax_h3_audio_vae.audio_vae" - ) + module = "sglang.multimodal_gen.runtime.models.vaes.minimax_h3_audio_vae.audio_vae" recording_fa = RecordingFA() with ( mock.patch(f"{module}.current_platform.is_cuda", return_value=True), diff --git a/python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py b/python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py index 9922743a8..cf94ed302 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py +++ b/python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py @@ -105,8 +105,9 @@ def test_native_ministral3_matches_hf_prefill_and_generation(): torch.testing.assert_close(native_output.logits, reference_output.logits) assert len(native_output.past_key_values.layers) == config.num_hidden_layers - with torch.no_grad(), set_forward_context( - current_timestep=0, attn_metadata=None + with ( + torch.no_grad(), + set_forward_context(current_timestep=0, attn_metadata=None), ): native_ids = native.generate(input_ids, max_new_tokens=2, do_sample=False) with torch.no_grad(): diff --git a/python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py index db3a9109c..cecdb0a18 100644 --- a/python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py +++ b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py @@ -12,8 +12,9 @@ ENV_PATH = ( def _backend(env_value, *, sm120): CudaPlatform.get_modelopt_flashinfer_fp4_backend.cache_clear() try: - with patch(ENV_PATH, env_value), patch.object( - CudaPlatform, "is_sm120", classmethod(lambda cls: sm120) + with ( + patch(ENV_PATH, env_value), + patch.object(CudaPlatform, "is_sm120", classmethod(lambda cls: sm120)), ): return CudaPlatform.get_modelopt_flashinfer_fp4_backend() finally: diff --git a/python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py b/python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py index fdd3fe1dd..864b62509 100644 --- a/python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py +++ b/python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py @@ -236,8 +236,8 @@ class TestComponentResidencyNvtxHooks(unittest.TestCase): manager.strategy_for = lambda _component_name, _module: ResidentStrategy() self.assertTrue(manager._should_keep_single_dit("transformer", module)) - manager.strategy_for = ( - lambda _component_name, _module: ComponentOffloadStrategy() + manager.strategy_for = lambda _component_name, _module: ( + ComponentOffloadStrategy() ) self.assertFalse(manager._should_keep_single_dit("transformer", module)) diff --git a/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py b/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py index d4ff2de69..733885e5e 100644 --- a/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py +++ b/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py @@ -688,10 +688,10 @@ def test_sample_actions_only_hoists_denoise_layout_for_eager(): model.denoise_step = lambda _ctx, x_t, _t, **_kwargs: torch.zeros_like(x_t) layout_calls = [] model.core_model = SimpleNamespace( - prepare_denoise_layout=lambda *args, **kwargs: layout_calls.append( - (args, kwargs) + prepare_denoise_layout=lambda *args, **kwargs: ( + layout_calls.append((args, kwargs)) + or (None, torch.zeros(1, 2, dtype=torch.long)) ) - or (None, torch.zeros(1, 2, dtype=torch.long)) ) observation = SimpleNamespace(batch_size=1) prefix_context = _prefix_context(1.0, "prompt") diff --git a/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py b/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py index 878b74b8f..80eeff9de 100644 --- a/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py +++ b/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py @@ -28,8 +28,9 @@ class _CompilableModule(nn.Module): class _RegionalModel(_CompilableModule): _compile_conditions = [ - lambda name, _module: name.startswith("transformer_blocks.") - and name.count(".") == 1 + lambda name, _module: ( + name.startswith("transformer_blocks.") and name.count(".") == 1 + ) ] def __init__(self): diff --git a/python/sglang/multimodal_gen/test/unit/test_rollout_api.py b/python/sglang/multimodal_gen/test/unit/test_rollout_api.py index 8f31e6531..64b0298c0 100644 --- a/python/sglang/multimodal_gen/test/unit/test_rollout_api.py +++ b/python/sglang/multimodal_gen/test/unit/test_rollout_api.py @@ -21,7 +21,6 @@ from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( class TestTensorToBytesRoundtrip(unittest.TestCase): - def _roundtrip(self, t: torch.Tensor): encoded = tensor_to_bytes(t) self.assertIsInstance(encoded, bytes) diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index e7f8f0b30..21cdb2136 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -427,9 +427,9 @@ class TestSchedulerFlowGRPOStepAlignmentUnit(unittest.TestCase): def _mock_variance_noise(_batch, *_args, **_kwargs): variance_noise_call_count["n"] += 1 - scheduler._get_rollout_session_data(_batch).noise_buffer = ( - variance_noise_ref - ) + scheduler._get_rollout_session_data( + _batch + ).noise_buffer = variance_noise_ref return variance_noise_ref scheduler._rollout_variance_noise = ( # type: ignore[method-assign] diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index ac41dceed..470516511 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -870,7 +870,6 @@ class TestDiffusionModelDetection(unittest.TestCase): class TestMiniMaxH3Routing(unittest.TestCase): - def test_semantic_variants_map_to_checkpoint_partitions(self): self.assertEqual( MiniMaxH3Pipeline.model_subfolder_for_variant("fl2va"), "FL2VA" diff --git a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py index e349daaa7..c6fb023ca 100644 --- a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py @@ -166,18 +166,22 @@ class TestTextEncoderClassResolution(unittest.TestCase): } loader = TextEncoderLoader() - with mock.patch.object( - TextEncoderLoader, - "resolve_native_transformers_model_class", - return_value=transformers_model_class, - ), mock.patch.object( - loader, - "target_device", - return_value=torch.device("cuda:0"), - ), mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.get_hf_config", - return_value=component_config, + with ( + mock.patch.object( + TextEncoderLoader, + "resolve_native_transformers_model_class", + return_value=transformers_model_class, + ), + mock.patch.object( + loader, + "target_device", + return_value=torch.device("cuda:0"), + ), + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_hf_config", + return_value=component_config, + ), ): encoder = loader.load_native( "/model/text_encoder", @@ -719,9 +723,12 @@ class TestTextEncoderQuantization(unittest.TestCase): "CLIPTextModel", "ThirdPartyTextEncoder", ): - with self.subTest(architecture=architecture), self.assertRaisesRegex( - NativeComponentLoaderRequired, - "delegates serialized quant_method='bitsandbytes' checkpoint", + with ( + self.subTest(architecture=architecture), + self.assertRaisesRegex( + NativeComponentLoaderRequired, + "delegates serialized quant_method='bitsandbytes' checkpoint", + ), ): _resolve_and_configure_encoder_quantization( SimpleNamespace(architectures=[architecture], quant_config=None), diff --git a/python/sglang/multimodal_gen/test/unit/test_usp_ipc_a2a_guard.py b/python/sglang/multimodal_gen/test/unit/test_usp_ipc_a2a_guard.py index e2553c650..840b9818b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_usp_ipc_a2a_guard.py +++ b/python/sglang/multimodal_gen/test/unit/test_usp_ipc_a2a_guard.py @@ -37,9 +37,10 @@ class TestIpcInputA2AQkvGuard(unittest.TestCase): # group lookup to prove the guard is not what rejected it. q = torch.zeros(1, 128, 8, 64) group = mock.MagicMock(return_value=None) - with mock.patch.object( - usp, "get_ulysses_parallel_world_size", lambda: 2 - ), mock.patch.object(usp, "_ipc_ready_group", group): + with ( + mock.patch.object(usp, "get_ulysses_parallel_world_size", lambda: 2), + mock.patch.object(usp, "_ipc_ready_group", group), + ): self.assertIsNone(usp._ipc_input_a2a_qkv(q, q.clone(), q.clone())) # Reached the group lookup, so the shape guard did not reject it. self.assertEqual(group.call_count, 1) diff --git a/python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py b/python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py index d3a985cb1..2a557d35c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py +++ b/python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py @@ -77,8 +77,8 @@ class RingReplicatedBase(unittest.TestCase): ), patch( f"{_LAYER}.ring_attn", - side_effect=lambda q, k, v, impl, return_softmax_lse: _ring_pair_via_impl( - impl, q, k, v + side_effect=lambda q, k, v, impl, return_softmax_lse: ( + _ring_pair_via_impl(impl, q, k, v) ), ), ) @@ -142,14 +142,19 @@ class TestRingReplicatedSuffix(RingReplicatedBase): o.copy_(t) ps = self._patches(ring_ws=1) - with ps[0], ps[4], ps[5], ps[6], patch( - f"{_LAYER}.get_ring_parallel_world_size", return_value=1 - ), patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=2), patch( - f"{_LAYER}.get_ulysses_parallel_world_size", return_value=2 - ), patch( - f"{_LAYER}.get_sp_group", return_value=SimpleNamespace(ulysses_group=None) - ), patch( - "torch.distributed.all_gather", side_effect=_fake_gather + with ( + ps[0], + ps[4], + ps[5], + ps[6], + patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=1), + patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=2), + patch(f"{_LAYER}.get_ulysses_parallel_world_size", return_value=2), + patch( + f"{_LAYER}.get_sp_group", + return_value=SimpleNamespace(ulysses_group=None), + ), + patch("torch.distributed.all_gather", side_effect=_fake_gather), ): # Identity-mocked collectives don't reproduce head-shard shapes, # so the final concat may fail โ the kernel K order is recorded diff --git a/python/sglang/multimodal_gen/utils.py b/python/sglang/multimodal_gen/utils.py index 6729138ee..be12a1551 100644 --- a/python/sglang/multimodal_gen/utils.py +++ b/python/sglang/multimodal_gen/utils.py @@ -135,7 +135,6 @@ def current_stream() -> torch.cuda.Stream | None: class StoreBoolean(argparse.Action): - def __init__(self, option_strings, dest, default=False, required=False, help=None): super().__init__( option_strings=option_strings, @@ -157,7 +156,7 @@ class StoreBoolean(argparse.Action): setattr(namespace, self.dest, False) else: raise ValueError( - f"Invalid boolean value: {values}. " "Expected 'true' or 'false'." + f"Invalid boolean value: {values}. Expected 'true' or 'false'." ) else: setattr(namespace, self.dest, bool(values)) @@ -295,8 +294,7 @@ class FlexibleArgumentParser(argparse.ArgumentParser): if args[0] == "serve": if index == 1: raise ValueError( - "No model_tag specified! Please check your command-line" - " arguments." + "No model_tag specified! Please check your command-line arguments." ) command = args_before_config[0] model_tag = args_before_config[1] @@ -484,7 +482,7 @@ def update_environment_variables(envs: dict[str, str]): for k, v in envs.items(): if k in os.environ and os.environ[k] != v: logger.warning( - "Overwriting environment variable %s " "from '%s' to '%s'", + "Overwriting environment variable %s from '%s' to '%s'", k, os.environ[k], v, @@ -509,7 +507,7 @@ def run_method( func = getattr(obj, method) except AttributeError: raise NotImplementedError( - f"Method {method!r} is not" " implemented." + f"Method {method!r} is not implemented." ) from None else: func = partial(method, obj) # type: ignore @@ -549,7 +547,6 @@ def get_exception_traceback() -> str: class TypeBasedDispatcher: - def __init__(self, mapping: list[tuple[type, Callable]]): self._mapping = mapping @@ -626,9 +623,9 @@ def dict_to_3d_list( """ # Case 1: no data, but fixed shape requested if mask_strategy is None: - assert ( - t_max is not None and l_max is not None and h_max is not None - ), "If mask_strategy is None, you must provide t_max, l_max, and h_max" + assert t_max is not None and l_max is not None and h_max is not None, ( + "If mask_strategy is None, you must provide t_max, l_max, and h_max" + ) return [ [[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max) ] diff --git a/python/sglang/srt/arg_groups/attention_hook.py b/python/sglang/srt/arg_groups/attention_hook.py index 282563416..d9bfab8c2 100644 --- a/python/sglang/srt/arg_groups/attention_hook.py +++ b/python/sglang/srt/arg_groups/attention_hook.py @@ -91,9 +91,9 @@ def handle_attention_backend_compatibility(server_args: Any): cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED ), ) - assert ( - cfg.speculative_algorithm is None - ), "Speculative decoding is currently not supported with Flex Attention backend" + assert cfg.speculative_algorithm is None, ( + "Speculative decoding is currently not supported with Flex Attention backend" + ) # Whisper's encoder token padding conflicts with prefix caching. # Only disable for Whisper; other encoder-decoder models (e.g., mllama) use radix cache. diff --git a/python/sglang/srt/arg_groups/cuda_graph_hook.py b/python/sglang/srt/arg_groups/cuda_graph_hook.py index 788f08292..4b3a5dfba 100644 --- a/python/sglang/srt/arg_groups/cuda_graph_hook.py +++ b/python/sglang/srt/arg_groups/cuda_graph_hook.py @@ -148,7 +148,7 @@ def apply_cuda_graph_compatibility(server_args: Any): and attention_backends_of(resolved_view(server_args))[0] != "trtllm_mla" ): logger.info( - "Using tc_piecewise CUDA graph for validated multimodal " "decoder prefill." + "Using tc_piecewise CUDA graph for validated multimodal decoder prefill." ) declare_resolution( server_args, @@ -183,16 +183,20 @@ def disable_tc_piecewise_cudagraph_if_incompatible(server_args: Any): ("pipeline parallelism (pp_size > 1)", lambda: cfg.pp_size > 1), ( "non-CUDA hardware (HIP/NPU/CPU/MPS/XPU)", - lambda: get_platform().is_hip - or get_platform().is_npu - or is_cpu() - or is_mps() - or get_platform().is_xpu, + lambda: ( + get_platform().is_hip + or get_platform().is_npu + or is_cpu() + or is_mps() + or get_platform().is_xpu + ), ), ( "OOT platform without piecewise support", - lambda: current_platform.is_out_of_tree() - and not current_platform.support_piecewise_cuda_graph(), + lambda: ( + current_platform.is_out_of_tree() + and not current_platform.support_piecewise_cuda_graph() + ), ), ( "MoE A2A backend", @@ -203,16 +207,20 @@ def disable_tc_piecewise_cudagraph_if_incompatible(server_args: Any): ("LoRA", lambda: bool(cfg.lora_paths) or cfg.enable_lora), ( "multimodal model", - lambda: model_config_of(server_args).is_multimodal - and not model_config_of( - server_args - ).is_multimodal_piecewise_cuda_graph_supported, + lambda: ( + model_config_of(server_args).is_multimodal + and not model_config_of( + server_args + ).is_multimodal_piecewise_cuda_graph_supported + ), ), ( "GGUF quantization", - lambda: cfg.load_format == "gguf" - or resolved_view(server_args).quantization == "gguf" - or check_gguf_file(cfg.model_path), + lambda: ( + cfg.load_format == "gguf" + or resolved_view(server_args).quantization == "gguf" + or check_gguf_file(cfg.model_path) + ), ), ("DLLM (diffusion LLM)", lambda: cfg.dllm_algorithm is not None), ( @@ -227,8 +235,9 @@ def disable_tc_piecewise_cudagraph_if_incompatible(server_args: Any): ("symmetric memory", lambda: cfg.enable_symm_mem), ( "expert distribution recorder", - lambda: cfg.enable_eplb - or cfg.expert_distribution_recorder_mode is not None, + lambda: ( + cfg.enable_eplb or cfg.expert_distribution_recorder_mode is not None + ), ), ( "context parallel (attn_cp_size > 1)", @@ -279,8 +288,10 @@ def disable_breakable_cudagraph_if_incompatible(server_args: Any): # CP all_gather replay size mismatch under BCG. ( "context parallel (attn_cp_size > 1)", - lambda: resolved_view(server_args).attn_cp_size > 1 - and not supports_prefill_cp_bcg(server_args), + lambda: ( + resolved_view(server_args).attn_cp_size > 1 + and not supports_prefill_cp_bcg(server_args) + ), ), # Capture builds a dummy extend forward with attn_dcp_metadata=None. ( @@ -294,16 +305,20 @@ def disable_breakable_cudagraph_if_incompatible(server_args: Any): ), ( "unvalidated a2a backend", - lambda: resolved_view(server_args).moe_a2a_backend - not in ("none", "deepep", "megamoe", "flashinfer"), + lambda: ( + resolved_view(server_args).moe_a2a_backend + not in ("none", "deepep", "megamoe", "flashinfer") + ), ), # Multimodal prefill replay faults under BCG; allowlisted archs opt back in. ( "multimodal model", - lambda: model_config_of(server_args).is_multimodal - and not model_config_of( - server_args - ).is_multimodal_breakable_cuda_graph_supported, + lambda: ( + model_config_of(server_args).is_multimodal + and not model_config_of( + server_args + ).is_multimodal_breakable_cuda_graph_supported + ), ), ] for name, predicate in rules: diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 2c991845f..551625c6d 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py @@ -148,11 +148,13 @@ def apply_deepseek_v4_defaults(server_args: ServerArgs, model_arch: str) -> None assert cfg.speculative_algorithm in ( "EAGLE", "DSPARK", - ), f"Only EAGLE and DSPARK speculative algorithms are supported for {model_arch}" + ), ( + f"Only EAGLE and DSPARK speculative algorithms are supported for {model_arch}" + ) if cfg.speculative_algorithm == "EAGLE": - assert ( - cfg.speculative_eagle_topk == 1 - ), f"Only EAGLE speculative algorithm with topk == 1 is supported for {model_arch}" + assert cfg.speculative_eagle_topk == 1, ( + f"Only EAGLE speculative algorithm with topk == 1 is supported for {model_arch}" + ) def validate_deepseek_v4_cp(server_args: ServerArgs) -> None: @@ -163,7 +165,7 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None: if cfg.cp_strategy != "interleave": raise ValueError( - "DeepSeekV4 only supports interleave CP strategy, " f"got {cfg.cp_strategy}" + f"DeepSeekV4 only supports interleave CP strategy, got {cfg.cp_strategy}" ) declare_resolution( @@ -196,12 +198,12 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None: "validate_deepseek_v4_cp", attn_cp_size=cfg.tp_size // cfg.dp_size, ) - assert ( - cfg.dp_size == 1 - ), "For round-robin split mode, dp attention is not supported." - assert ( - cfg.tp_size <= 8 - ), "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." + assert cfg.dp_size == 1, ( + "For round-robin split mode, dp attention is not supported." + ) + assert cfg.tp_size <= 8, ( + "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." + ) supported_a2a_backends = ("none", "deepep", "megamoe", "mori") if cfg.moe_a2a_backend not in supported_a2a_backends: raise ValueError( diff --git a/python/sglang/srt/arg_groups/hisparse_hook.py b/python/sglang/srt/arg_groups/hisparse_hook.py index 3b38b1e40..cdb841dff 100644 --- a/python/sglang/srt/arg_groups/hisparse_hook.py +++ b/python/sglang/srt/arg_groups/hisparse_hook.py @@ -96,9 +96,9 @@ def validate_hisparse(server_args: ServerArgs) -> None: "models (e.g., DeepSeek V3.2, GLM-5) and DeepSeek V4 now. " ) - assert ( - cfg.disable_radix_cache - ), "Hierarchical sparse attention currently requires --disable-radix-cache." + assert cfg.disable_radix_cache, ( + "Hierarchical sparse attention currently requires --disable-radix-cache." + ) # DSv4 hisparse handles its own dtype/backend pairing elsewhere; the dtype- # aware checks below only apply to the DSA hisparse path. diff --git a/python/sglang/srt/arg_groups/lora_hook.py b/python/sglang/srt/arg_groups/lora_hook.py index c349db689..82df754e1 100644 --- a/python/sglang/srt/arg_groups/lora_hook.py +++ b/python/sglang/srt/arg_groups/lora_hook.py @@ -76,9 +76,9 @@ def check_lora_server_args(server_args: Any): pinned=False, ) elif isinstance(lora_path, dict): - assert ( - "lora_name" in lora_path and "lora_path" in lora_path - ), f"When providing LoRA paths as a list of dict, each dict should contain 'lora_name' and 'lora_path' keys. Got: {lora_path}" + assert "lora_name" in lora_path and "lora_path" in lora_path, ( + f"When providing LoRA paths as a list of dict, each dict should contain 'lora_name' and 'lora_path' keys. Got: {lora_path}" + ) lora_ref = LoRARef( lora_id=LoRARef.deterministic_id( lora_path["lora_name"], lora_path["lora_path"] @@ -129,14 +129,14 @@ def check_lora_server_args(server_args: Any): lora_target_modules=set(cfg.lora_target_modules), ) if "all" in cfg.lora_target_modules: - assert ( - len(cfg.lora_target_modules) == 1 - ), "If 'all' is specified in --lora-target-modules, it should be the only module specified." + assert len(cfg.lora_target_modules) == 1, ( + "If 'all' is specified in --lora-target-modules, it should be the only module specified." + ) # Ensure sufficient information is provided for LoRA initialization. - assert cfg.lora_paths or ( - cfg.max_lora_rank and cfg.lora_target_modules - ), "When no initial --lora-paths is provided, you need to specify both --max-lora-rank and --lora-target-modules for LoRA initialization." + assert cfg.lora_paths or (cfg.max_lora_rank and cfg.lora_target_modules), ( + "When no initial --lora-paths is provided, you need to specify both --max-lora-rank and --lora-target-modules for LoRA initialization." + ) # Validate max_loaded_loras if cfg.max_loaded_loras is not None: @@ -158,9 +158,9 @@ def check_lora_server_args(server_args: Any): if cfg.lora_use_virtual_experts: logger.info("Virtual expert computation enabled.") - assert ( - cfg.lora_drain_wait_threshold >= 0.0 - ), "--lora-drain-wait-threshold must be non-negative." + assert cfg.lora_drain_wait_threshold >= 0.0, ( + "--lora-drain-wait-threshold must be non-negative." + ) def check_lora_speculative_compatibility(server_args: Any): diff --git a/python/sglang/srt/arg_groups/mamba_hook.py b/python/sglang/srt/arg_groups/mamba_hook.py index 39e28cc9e..9107f432e 100644 --- a/python/sglang/srt/arg_groups/mamba_hook.py +++ b/python/sglang/srt/arg_groups/mamba_hook.py @@ -97,9 +97,9 @@ def handle_int8_mamba_checkpoint(server_args: Any): def validate_mamba_extra_buffer(view, model_arch: str, *, mamba_cache_chunk_size_of): - assert supports_mamba_cache_extra_buffer( - view, model_arch - ), f"extra_buffer is not supported for {model_arch}; use no_buffer." + assert supports_mamba_cache_extra_buffer(view, model_arch), ( + f"extra_buffer is not supported for {model_arch}; use no_buffer." + ) assert ( get_platform().is_cuda or get_platform().is_musa @@ -142,9 +142,9 @@ def validate_mamba_extra_buffer(view, model_arch: str, *, mamba_cache_chunk_size def validate_mamba_no_buffer(view, model_arch: str): assert view.page_size in (1, None), "no_buffer only supports page_size=1." - assert ( - view.disable_overlap_schedule - ), "no_buffer do not support overlap schedule. Try to set disable_overlap_schedule=True." - assert ( - view.attention_backend != "trtllm_mha" - ), "no_buffer do not support trtllm_mha attention backend." + assert view.disable_overlap_schedule, ( + "no_buffer do not support overlap schedule. Try to set disable_overlap_schedule=True." + ) + assert view.attention_backend != "trtllm_mha", ( + "no_buffer do not support trtllm_mha attention backend." + ) diff --git a/python/sglang/srt/arg_groups/memory_hook.py b/python/sglang/srt/arg_groups/memory_hook.py index 39c5b2ad8..0685fe805 100644 --- a/python/sglang/srt/arg_groups/memory_hook.py +++ b/python/sglang/srt/arg_groups/memory_hook.py @@ -177,9 +177,9 @@ def handle_gpu_memory_settings(server_args: Any, gpu_mem): ) decode_cuda_graph_config.bs = generate_cpu_graph_batch_sizes(server_args) - assert ( - cfg.torch_compile_max_bs > 0 - ), "cuda_graph_config[decode].bs should contain positive batch sizes" + assert cfg.torch_compile_max_bs > 0, ( + "cuda_graph_config[decode].bs should contain positive batch sizes" + ) decode_cuda_graph_config.max_bs = cfg.torch_compile_max_bs if prefill_cuda_graph_config.max_bs is None: diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index 56bfd3b87..a56fe3f31 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -218,9 +218,9 @@ def handle_model_specific_adjustments(server_args: Any): run_post_process_pass(server_args, _dsa_split_backend_resolution) if cfg.enable_prefill_cp: - assert ( - cfg.disaggregation_mode != "decode" - ), "CP is only supported for prefill when PD disaggregation, please remove --enable-prefill-cp." + assert cfg.disaggregation_mode != "decode", ( + "CP is only supported for prefill when PD disaggregation, please remove --enable-prefill-cp." + ) if ( cfg.enable_dsa_cache_layer_split and cfg.disaggregation_mode != "prefill" @@ -423,9 +423,9 @@ def handle_model_specific_adjustments(server_args: Any): # (arg_groups/overrides.py: _gpt_oss_overrides). if resolved_view(server_args).moe_runner_backend == "triton_kernel": - assert ( - resolved_view(server_args).ep_size == 1 - ), "Triton kernel MoE is only supported when ep_size == 1" + assert resolved_view(server_args).ep_size == 1, ( + "Triton kernel MoE is only supported when ep_size == 1" + ) elif model_arch in ("MiMoV2ForCausalLM", "MiMoV2FlashForCausalLM"): if model_arch == "MiMoV2ForCausalLM" and not cfg.encoder_only: @@ -481,7 +481,9 @@ def handle_model_specific_adjustments(server_args: Any): "ascend", "trtllm_mha", "intel_xpu", - }, f"fa3, aiter, triton, ascend, trtllm_mha or intel_xpu is required for Llama4 model but got {attention_backend}" + }, ( + f"fa3, aiter, triton, ascend, trtllm_mha or intel_xpu is required for Llama4 model but got {attention_backend}" + ) # The moe_runner_backend selection moved to the override registry # (arg_groups/overrides.py: _llama4_overrides). # Gemma2/Gemma3 (disable_hybrid_swa_memory) moved to the override registry @@ -523,9 +525,9 @@ def handle_model_specific_adjustments(server_args: Any): # https://docs.sglang.ai/advanced_features/attention_backend.html accepted_backends = ["fa3", "triton", "trtllm_mha"] attention_backend = resolved_view(server_args).attention_backend - assert ( - attention_backend in accepted_backends - ), f"One of the attention backends in {accepted_backends} is required for {model_arch}, but got {attention_backend}" + assert attention_backend in accepted_backends, ( + f"One of the attention backends in {accepted_backends} is required for {model_arch}, but got {attention_backend}" + ) elif model_arch in ["Olmo2ForCausalLM"]: # disable_hybrid_swa_memory + attention backend selection moved to # the override registry (arg_groups/overrides.py: _olmo2_overrides). @@ -534,9 +536,9 @@ def handle_model_specific_adjustments(server_args: Any): # is used for the Olmo2 architecture. Olmo2 does not use sliding window attention # but Olmo3 does. attention_backend = resolved_view(server_args).attention_backend - assert ( - attention_backend != "flashinfer" - ), "FlashInfer backend can significantly degrade the performance of Olmo3 models." + assert attention_backend != "flashinfer", ( + "FlashInfer backend can significantly degrade the performance of Olmo3 models." + ) logger.info(f"Using {attention_backend} as attention backend for {model_arch}.") elif model_arch in [ diff --git a/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py b/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py index e50d65be1..19d48b659 100644 --- a/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py +++ b/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py @@ -59,12 +59,12 @@ def _deepseek_family_overrides(server_args: Any, hf_config: Any) -> dict: "moe_a2a_backend=deepep, ep_size=tp_size, batch_size=1." ) else: - assert ( - cfg.dp_size == 1 - ), "interleave DSA CP does not support DP attention." - assert ( - cfg.tp_size <= 8 - ), "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." + assert cfg.dp_size == 1, ( + "interleave DSA CP does not support DP attention." + ) + assert cfg.tp_size <= 8, ( + "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." + ) # Note(kpham-sgl): Keep attn_tp_size == 1 under DSA CP. # DSACPLayerCommunicator does not all-reduce attention-TP # partial o_proj outputs before replicated dense FFNs. diff --git a/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py b/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py index 558260893..3aa0c4711 100644 --- a/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py +++ b/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py @@ -70,7 +70,5 @@ def _deepseek_v4_overrides(server_args: Any, hf_config: Any) -> dict: ) ): overrides["moe_runner_backend"] = "flashinfer_mxfp4" - logger.info( - "Use flashinfer_mxfp4 as MoE runner backend for " f"{model_arch}." - ) + logger.info(f"Use flashinfer_mxfp4 as MoE runner backend for {model_arch}.") return overrides diff --git a/python/sglang/srt/arg_groups/model_overrides/gpt_oss.py b/python/sglang/srt/arg_groups/model_overrides/gpt_oss.py index 447702749..0bee69abc 100644 --- a/python/sglang/srt/arg_groups/model_overrides/gpt_oss.py +++ b/python/sglang/srt/arg_groups/model_overrides/gpt_oss.py @@ -67,7 +67,6 @@ def _gpt_oss_overrides(server_args: Any, hf_config: Any) -> dict: # use bf16 for mxfp4 triton kernels overrides["dtype"] = "bfloat16" if cfg.moe_runner_backend == "auto": - if get_platform().is_sm100 and is_mxfp4_quant_format: overrides["moe_runner_backend"] = "flashinfer_mxfp4" logger.warning( diff --git a/python/sglang/srt/arg_groups/moe_hook.py b/python/sglang/srt/arg_groups/moe_hook.py index 7eeafb0c4..b1146dd32 100644 --- a/python/sglang/srt/arg_groups/moe_hook.py +++ b/python/sglang/srt/arg_groups/moe_hook.py @@ -47,7 +47,9 @@ def handle_moe_kernel_config(server_args: Any): "modelopt_fp8", "modelopt_mixed", None, - ], f"Invalid quantization '{view.quantization}'. \nFlashInfer Cutlass MOE supports only: 'modelopt_fp4', 'modelopt_fp8', 'modelopt_mixed', or bfloat16 (None)." + ], ( + f"Invalid quantization '{view.quantization}'. \nFlashInfer Cutlass MOE supports only: 'modelopt_fp4', 'modelopt_fp8', 'modelopt_mixed', or bfloat16 (None)." + ) assert view.ep_size in [ 1, cfg.tp_size, @@ -58,7 +60,9 @@ def handle_moe_kernel_config(server_args: Any): assert ( view.quantization in ["modelopt_fp4", "modelopt_mixed", "nvfp4_online"] or model_config_of(server_args).nvfp4_moe_meta is not None - ), f"Invalid quantization '{view.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4', 'modelopt_mixed' (with NVFP4 MoE layers), 'nvfp4_online', or hybrid NVFP4 models." + ), ( + f"Invalid quantization '{view.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4', 'modelopt_mixed' (with NVFP4 MoE layers), 'nvfp4_online', or hybrid NVFP4 models." + ) assert view.ep_size in [ 1, cfg.tp_size, @@ -90,7 +94,9 @@ def handle_moe_kernel_config(server_args: Any): "modelopt_mixed", "compressed-tensors", None, - ], f"Invalid quantization '{view.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'nvfp4_online', 'fp8', 'modelopt_fp8', 'modelopt_mixed', 'compressed-tensors', or bfloat16 (None)." + ], ( + f"Invalid quantization '{view.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'nvfp4_online', 'fp8', 'modelopt_fp8', 'modelopt_mixed', 'compressed-tensors', or bfloat16 (None)." + ) if view.moe_runner_backend == "flashinfer_trtllm_routed": assert view.quantization in [ @@ -100,7 +106,9 @@ def handle_moe_kernel_config(server_args: Any): "modelopt_mixed", "nvfp4_online", None, - ], f"Invalid quantization '{view.quantization}'. \nFlashInfer TRTLLM routed MOE supports only: 'fp8', 'mxfp8', 'modelopt_fp4', 'modelopt_mixed', 'nvfp4_online', or bfloat16 (None)." + ], ( + f"Invalid quantization '{view.quantization}'. \nFlashInfer TRTLLM routed MOE supports only: 'fp8', 'mxfp8', 'modelopt_fp4', 'modelopt_mixed', 'nvfp4_online', or bfloat16 (None)." + ) # The runner-driven shared-experts fusion disables moved to the # pipeline (arg_groups/overrides.py: _moe_runner_fusion_disable), @@ -113,9 +121,9 @@ def handle_moe_kernel_config(server_args: Any): "fp8", "mxfp8", ]: - assert ( - resolved_view(server_args).ep_size == 1 - ), "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1" + assert resolved_view(server_args).ep_size == 1, ( + "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1" + ) def handle_a2a_moe(server_args: Any): @@ -256,7 +264,9 @@ def handle_a2a_moe(server_args: Any): assert ( resolved_view(server_args).enable_dp_attention and cfg.dp_size == cfg.tp_size - ), "Flashinfer MoE A2A is only supported with dp_size == tp_size and --enable-dp-attention" + ), ( + "Flashinfer MoE A2A is only supported with dp_size == tp_size and --enable-dp-attention" + ) if cfg.deepep_mode != "auto": logger.warning("--deepep-mode is ignored for Flashinfer MoE A2A") use_cutedsl_w4a16 = ( diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 9debcb6d7..9ebae0e1c 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -449,8 +449,9 @@ import sglang.srt.arg_groups.model_overrides # noqa: F401 @register_model_override_predicate( - lambda arch: "Step3p5ForCausalLM" in arch - or "Step3p7ForConditionalGeneration" in arch + lambda arch: ( + "Step3p5ForCausalLM" in arch or "Step3p7ForConditionalGeneration" in arch + ) ) def _step3p_overrides(server_args: Any, hf_config: Any) -> dict: cfg = resolving_view(server_args) @@ -1261,9 +1262,9 @@ def _cutedsl_prefill_backend_fill(view: Any) -> dict: or view.prefill_attention_backend == "cutedsl_mla" ): return {} - assert ( - view.prefill_attention_backend != "cutedsl_mla" - ), "CuteDSL MLA only supports decoding for now" + assert view.prefill_attention_backend != "cutedsl_mla", ( + "CuteDSL MLA only supports decoding for now" + ) if not get_platform().is_sm100: raise ValueError( "CuteDSL MLA backend is only supported on Blackwell GPUs (SM100). Please use a different backend." @@ -1435,12 +1436,12 @@ def _dp_lm_head_validation(view: Any) -> dict: dp LM head and the TP LM-head all-to-all path. Reads the mid-resolution values through the view.""" if view.enable_dp_lm_head: - assert ( - view.enable_dp_attention - ), "Please enable dp attention when setting enable_dp_lm_head. " + assert view.enable_dp_attention, ( + "Please enable dp attention when setting enable_dp_lm_head. " + ) if view.enable_tp_lm_head_all_to_all: assert view.enable_dp_attention, ( - "Please enable dp attention when setting " "enable_tp_lm_head_all_to_all." + "Please enable dp attention when setting enable_tp_lm_head_all_to_all." ) assert not view.enable_dp_lm_head, ( "--enable-tp-lm-head-all-to-all uses a TP-sharded LM head and is " @@ -1500,7 +1501,7 @@ def _moe_runner_backend_quant_constraints(view: Any) -> dict: moe_runner_backend = mxfp8_default elif moe_runner_backend not in allowed: logger.warning( - "mxfp8 quantization supports only %s backends. " "Overriding %r.", + "mxfp8 quantization supports only %s backends. Overriding %r.", ", ".join(allowed), moe_runner_backend, ) @@ -1845,7 +1846,6 @@ def mamba_cache_chunk_size(server_args: Any) -> int: from sglang.srt.arg_groups.overrides import model_config_of if not hasattr(server_args, "_mamba_cache_chunk_size"): - try: from sglang.kernels.ops.attention.fla.chunk_delta_h import ( CHUNK_SIZE as FLA_CHUNK_SIZE, @@ -1857,9 +1857,9 @@ def mamba_cache_chunk_size(server_args: Any) -> int: hf_config = model_config_of(server_args).hf_config chunk_size = getattr(hf_config, "mamba_chunk_size", FLA_CHUNK_SIZE) page_size = resolved_view(server_args).page_size - assert ( - max(chunk_size, page_size) % min(chunk_size, page_size) == 0 - ), f"For SSM models, either chunk_size or page_size must be divisible by the other, got {chunk_size=}, {page_size=}" + assert max(chunk_size, page_size) % min(chunk_size, page_size) == 0, ( + f"For SSM models, either chunk_size or page_size must be divisible by the other, got {chunk_size=}, {page_size=}" + ) if not getattr(server_args, "_resolution_finished", False): return max(chunk_size, page_size) server_args._mamba_cache_chunk_size = max(chunk_size, page_size) diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index b066a56cb..7573a92ed 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -59,8 +59,7 @@ def handle_context_parallelism(server_args: Any): and not cfg.language_model_only ): raise ValueError( - "MiMo V2 CP-v2 only supports text inference; add " - "--language-only." + "MiMo V2 CP-v2 only supports text inference; add --language-only." ) if cfg.enable_prefill_cp and cfg.cp_strategy is None: @@ -81,40 +80,40 @@ def handle_context_parallelism(server_args: Any): view = resolved_view(server_args) if view.attn_cp_size > 1: # The tp_size is the world size, not the real tensor parallel size - assert ( - cfg.tp_size % view.attn_cp_size == 0 - ), "tp_size must be divisible by attn_cp_size" - assert ( - cfg.tp_size % (cfg.dp_size * view.attn_cp_size) == 0 - ), "tp_size must be divisible by dp_size * attn_cp_size" + assert cfg.tp_size % view.attn_cp_size == 0, ( + "tp_size must be divisible by attn_cp_size" + ) + assert cfg.tp_size % (cfg.dp_size * view.attn_cp_size) == 0, ( + "tp_size must be divisible by dp_size * attn_cp_size" + ) - assert ( - not cfg.enable_aiter_allreduce_fusion - ), "Aiter allreduce fusion is not supported with context parallelism" + assert not cfg.enable_aiter_allreduce_fusion, ( + "Aiter allreduce fusion is not supported with context parallelism" + ) if cfg.moe_dp_size > 1: # The tp_size is the world size, not the real tensor parallel size - assert ( - cfg.tp_size % cfg.moe_dp_size == 0 - ), "tp_size must be divisible by moe_dp_size" - assert ( - view.ep_size * cfg.moe_dp_size <= cfg.tp_size - ), "ep_size * moe_dp_size must be less than or equal to tp_size" + assert cfg.tp_size % cfg.moe_dp_size == 0, ( + "tp_size must be divisible by moe_dp_size" + ) + assert view.ep_size * cfg.moe_dp_size <= cfg.tp_size, ( + "ep_size * moe_dp_size must be less than or equal to tp_size" + ) assert cfg.pp_size == 1, "PP is not supported with context parallelism" if view.ep_size > 1: - assert ( - view.ep_size * cfg.moe_dp_size == cfg.tp_size - ), "ep_size * moe_dp_size must be equal to tp_size" + assert view.ep_size * cfg.moe_dp_size == cfg.tp_size, ( + "ep_size * moe_dp_size must be equal to tp_size" + ) - assert ( - not cfg.enable_aiter_allreduce_fusion - ), "Aiter allreduce fusion is not supported with context parallelism" + assert not cfg.enable_aiter_allreduce_fusion, ( + "Aiter allreduce fusion is not supported with context parallelism" + ) if view.attn_cp_size != cfg.moe_dp_size: - assert ( - cfg.moe_dp_size == 1 - ), "attn_cp_size != moe_dp_size is only supported when moe_dp_size == 1" + assert cfg.moe_dp_size == 1, ( + "attn_cp_size != moe_dp_size is only supported when moe_dp_size == 1" + ) from sglang.srt.layers.cp.base import init_cp_strategy @@ -244,26 +243,26 @@ def handle_dwdp(server_args: Any): if cfg.dwdp_size <= 1: return - assert ( - cfg.dwdp_size >= 2 - ), f"dwdp_size must be >= 2 when enabled, got {cfg.dwdp_size}" - assert ( - cfg.dwdp_size == cfg.tp_size - ), f"dwdp_size ({cfg.dwdp_size}) must equal tp_size ({cfg.tp_size})" + assert cfg.dwdp_size >= 2, ( + f"dwdp_size must be >= 2 when enabled, got {cfg.dwdp_size}" + ) + assert cfg.dwdp_size == cfg.tp_size, ( + f"dwdp_size ({cfg.dwdp_size}) must equal tp_size ({cfg.tp_size})" + ) assert cfg.disaggregation_mode in ( "null", "prefill", ), "DWDP requires --disaggregation-mode null or prefill" - assert ( - not cfg.enable_eplb - ), "EPLB dynamic migration conflicts with static DWDP partitioning" - assert ( - cfg.speculative_algorithm is None - ), "DWDP does not support speculative decoding (MTP/draft workers)" + assert not cfg.enable_eplb, ( + "EPLB dynamic migration conflicts with static DWDP partitioning" + ) + assert cfg.speculative_algorithm is None, ( + "DWDP does not support speculative decoding (MTP/draft workers)" + ) assert cfg.pp_size == 1, "DWDP requires pp_size == 1" - assert ( - not cfg.enable_two_batch_overlap - ), "DWDP's prefetch event protocol does not support two-batch overlap" + assert not cfg.enable_two_batch_overlap, ( + "DWDP's prefetch event protocol does not support two-batch overlap" + ) if cfg.disaggregation_mode == "null": logger.warning( @@ -359,7 +358,9 @@ def handle_elastic_ep(server_args: Any): assert cfg.eplb_algorithm in [ "elasticity_aware", "elasticity_aware_hierarchical", - ], "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'." + ], ( + "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'." + ) assert cfg.pp_size == 1, "PP size should be set to 1 under elastic EP" @@ -370,9 +371,9 @@ def handle_elastic_ep(server_args: Any): mooncake_ib_device=validate_ib_devices(cfg.mooncake_ib_device), ) if cfg.ep_join_mode is not None: - assert ( - cfg.elastic_ep_backend is not None - ), "--elastic-ep-join-mode requires --elastic-ep-backend to be set." + assert cfg.elastic_ep_backend is not None, ( + "--elastic-ep-join-mode requires --elastic-ep-backend to be set." + ) if cfg.ep_join_mode == "scale": assert cfg.node_rank == 1, ( "Elastic EP scale-up requires one joining TP group at " @@ -390,9 +391,9 @@ def handle_elastic_ep(server_args: Any): ) assert cfg.ep_join_rank_offset >= 0, "elastic EP join rank offset must be >= 0." if cfg.max_ep_size is not None: - assert ( - cfg.elastic_ep_backend is not None - ), "--max-ep-size requires --elastic-ep-backend to be set." + assert cfg.elastic_ep_backend is not None, ( + "--max-ep-size requires --elastic-ep-backend to be set." + ) assert cfg.max_ep_size > 0, "--max-ep-size must be a positive integer." scaling_active = ( @@ -407,16 +408,15 @@ def handle_elastic_ep(server_args: Any): ) if scaling_active: resolved = resolved_view(server_args) - assert ( - cfg.elastic_ep_scale_timeout > 0 - ), "--elastic-ep-scale-timeout must be greater than zero." + assert cfg.elastic_ep_scale_timeout > 0, ( + "--elastic-ep-scale-timeout must be greater than zero." + ) assert cfg.tokenizer_worker_num == 1, ( - "Elastic EP runtime scale-up currently requires " - "--tokenizer-worker-num 1." + "Elastic EP runtime scale-up currently requires --tokenizer-worker-num 1." + ) + assert not cfg.use_ray, ( + "Elastic EP runtime scale-up does not support --use-ray." ) - assert ( - not cfg.use_ray - ), "Elastic EP runtime scale-up does not support --use-ray." assert not cfg.enable_elastic_expert_backup, ( "Elastic EP runtime scale-up does not support " "--enable-elastic-expert-backup." diff --git a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py index 7e05b9176..a615e59ea 100644 --- a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py +++ b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py @@ -129,9 +129,9 @@ def handle_pd_disaggregation(server_args: ServerArgs) -> None: ) elif cfg.disaggregation_mode == "prefill": - assert ( - cfg.disaggregation_transfer_backend != "fake" - ), "Prefill server does not support 'fake' as the transfer backend" + assert cfg.disaggregation_transfer_backend != "fake", ( + "Prefill server does not support 'fake' as the transfer backend" + ) if envs.SGLANG_RUST_SERVER.get(): _alias_bootstrap_port_to_api_port(server_args) diff --git a/python/sglang/srt/arg_groups/serving_hook.py b/python/sglang/srt/arg_groups/serving_hook.py index 39c3f63ca..404509bb3 100644 --- a/python/sglang/srt/arg_groups/serving_hook.py +++ b/python/sglang/srt/arg_groups/serving_hook.py @@ -75,7 +75,7 @@ def handle_ssl_validation(server_args: Any): if cfg.enable_http2: if not 0 < cfg.http2_max_concurrent_streams < 2**32: raise ValueError( - "--http2-max-concurrent-streams must be between 1 and " "4294967295." + "--http2-max-concurrent-streams must be between 1 and 4294967295." ) if not 1024 <= cfg.http2_initial_connection_window_size < 2**31: raise ValueError( @@ -343,8 +343,7 @@ def handle_deprecated_args(server_args: Any): ) if cfg.grpc_worker_threads is not None and cfg.grpc_worker_threads < 1: raise ValueError( - "SGLANG_GRPC_WORKER_THREADS " - f"({cfg.grpc_worker_threads}) must be >= 1" + f"SGLANG_GRPC_WORKER_THREADS ({cfg.grpc_worker_threads}) must be >= 1" ) # Native gRPC is incompatible with launch paths it doesn't wire into. @@ -482,8 +481,7 @@ def handle_other_validations(server_args: Any): ) elif resolved_view(server_args).uses_mamba_radix_cache: logger.warning( - "Optimistic prefill does not support models that use " - "mamba radix cache." + "Optimistic prefill does not support models that use mamba radix cache." ) declare_resolution( server_args, @@ -851,8 +849,7 @@ def handle_multimodal_feature_transport(server_args: Any): raise ValueError("--mm-feature-transport=cuda_vmm requires NVIDIA CUDA.") if cfg.pp_size != 1: raise ValueError( - "--mm-feature-transport=cuda_vmm does not support pipeline " - "parallelism." + "--mm-feature-transport=cuda_vmm does not support pipeline parallelism." ) if envs.SGLANG_RUST_SERVER.get(): raise ValueError( diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index b7aefdc4a..3726bf44e 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -559,7 +559,6 @@ def _resolve_dflash_draft_attention_backend(server_args: ServerArgs) -> None: draft_backend = cfg.speculative_draft_attention_backend if draft_backend is None: - draft_backend, _ = attention_backends_of(resolved_view(server_args)) if draft_backend is None: draft_backend = fallback_backend diff --git a/python/sglang/srt/arg_groups/validation_hook.py b/python/sglang/srt/arg_groups/validation_hook.py index e02cccded..eef15b96d 100644 --- a/python/sglang/srt/arg_groups/validation_hook.py +++ b/python/sglang/srt/arg_groups/validation_hook.py @@ -31,9 +31,9 @@ def check_server_args(server_args: Any): # Check parallel size constraints if cfg.ep_join_mode != "scale": - assert ( - cfg.tp_size * cfg.pp_size - ) % cfg.nnodes == 0, "tp_size must be divisible by number of nodes" + assert (cfg.tp_size * cfg.pp_size) % cfg.nnodes == 0, ( + "tp_size must be divisible by number of nodes" + ) assert cfg.pp_max_micro_batch_size is None or cfg.pp_max_micro_batch_size >= 1, ( "pp_max_micro_batch_size must be a positive integer or None (for auto-compute). " @@ -49,18 +49,18 @@ def check_server_args(server_args: Any): ) if cfg.pp_size > 1: - assert ( - cfg.disable_overlap_schedule and cfg.speculative_algorithm is None - ), "Pipeline parallelism is not compatible with overlap schedule, speculative decoding" + assert cfg.disable_overlap_schedule and cfg.speculative_algorithm is None, ( + "Pipeline parallelism is not compatible with overlap schedule, speculative decoding" + ) assert cfg.min_free_slots_delay is None, ( "--min-free-slots-delay is not supported with pipeline " "parallelism: allocatable slots per microbatch are bounded by " "pp-max-micro-batch-size, so the threshold may never be reached" ) - assert not ( - cfg.dp_size > 1 and cfg.nnodes != 1 and not cfg.enable_dp_attention - ), "multi-node data parallel is not supported unless dp attention!" + assert not (cfg.dp_size > 1 and cfg.nnodes != 1 and not cfg.enable_dp_attention), ( + "multi-node data parallel is not supported unless dp attention!" + ) assert cfg.base_gpu_id >= 0, "base_gpu_id must be non-negative" assert cfg.gpu_id_step >= 1, "gpu_id_step must be positive" @@ -102,24 +102,24 @@ def check_server_args(server_args: Any): # Skip validation if chunked prefill is disabled (i.e., size <= 0). # Skip validation if disaggregation mode is decode. if cfg.chunked_prefill_size > 0 and cfg.disaggregation_mode != "decode": - assert ( - cfg.chunked_prefill_size % cfg.page_size == 0 - ), "chunked_prefill_size must be divisible by page_size" + assert cfg.chunked_prefill_size % cfg.page_size == 0, ( + "chunked_prefill_size must be divisible by page_size" + ) # Check pdmux if cfg.enable_pdmux: - assert ( - cfg.pp_size == 1 - ), "PD-Multiplexing is only supported with pipeline parallelism disabled (pp_size=1)." - assert ( - cfg.chunked_prefill_size == -1 - ), "PD-Multiplexing is not compatible with chunked prefill." - assert ( - cfg.disaggregation_mode == "null" - ), "PD-Multiplexing is not compatible with disaggregation mode." - assert ( - cfg.disable_overlap_schedule - ), "PD-Multiplexing is not compatible with overlap schedule." + assert cfg.pp_size == 1, ( + "PD-Multiplexing is only supported with pipeline parallelism disabled (pp_size=1)." + ) + assert cfg.chunked_prefill_size == -1, ( + "PD-Multiplexing is not compatible with chunked prefill." + ) + assert cfg.disaggregation_mode == "null", ( + "PD-Multiplexing is not compatible with disaggregation mode." + ) + assert cfg.disable_overlap_schedule, ( + "PD-Multiplexing is not compatible with overlap schedule." + ) # NOTE: CUDA Green Context may encounter potential issues with CudaGraph on torch 2.7.x โ 2.8.x, leading to performance degradation. import torch @@ -143,7 +143,9 @@ def check_server_args(server_args: Any): assert cfg.schedule_policy in [ "fcfs", "lof", - ], f"To use priority scheduling, schedule_policy must be 'fcfs' or 'lof'. '{cfg.schedule_policy}' is not supported." + ], ( + f"To use priority scheduling, schedule_policy must be 'fcfs' or 'lof'. '{cfg.schedule_policy}' is not supported." + ) if cfg.default_priority_value is None: logger.warning( "--default-priority-value is not set while --enable-priority-scheduling is enabled. " @@ -170,14 +172,14 @@ def check_server_args(server_args: Any): run_post_process_pass(server_args, _hisparse_validation) - assert ( - cfg.schedule_conservativeness >= 0 - ), "schedule_conservativeness must be non-negative" + assert cfg.schedule_conservativeness >= 0, ( + "schedule_conservativeness must be non-negative" + ) if cfg.model_impl == "mindspore": - assert ( - get_platform().is_npu - ), "MindSpore model impl is only supported on Ascend npu." + assert get_platform().is_npu, ( + "MindSpore model impl is only supported on Ascend npu." + ) # Check metrics labels if ( @@ -239,43 +241,45 @@ def validate_buckets_rule(arg_name: str, buckets_rule: List[str]): "tse", "default", "custom", - ], f"Unsupported {arg_name} rule type: '{rule}'. Must be one of: 'tse', 'default', 'custom'" + ], ( + f"Unsupported {arg_name} rule type: '{rule}'. Must be one of: 'tse', 'default', 'custom'" + ) if rule == "tse": - assert ( - len(buckets_rule) == 4 - ), f"{arg_name} TSE rule requires exactly 4 parameters: ['tse', middle, base, count], got {len(buckets_rule)}" + assert len(buckets_rule) == 4, ( + f"{arg_name} TSE rule requires exactly 4 parameters: ['tse', middle, base, count], got {len(buckets_rule)}" + ) try: middle = float(buckets_rule[1]) base = float(buckets_rule[2]) count = int(buckets_rule[3]) except (ValueError, IndexError): - assert ( - False - ), f"{arg_name} TSE rule parameters must be: ['tse',