minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[CI][RFC] Replace black-jupyter with ruff-format (#37210)

Browse Source 
Co-authored-by: Alison Shao <a.shao@wustl.edu>
This commit is contained in:
Alex Nails
2026-09-02 19:46:08 -07:00
committed by GitHub
co-authored by Alison Shao
parent 2641e427be
commit 28262c20df
1411 changed files with 7766 additions and 8176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py
+4 -6
View File
@@ -541,7 +541,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
"fused_flashmla_metadata", "fused_flashmla_metadata",
), ),
rationale_hint=( 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, min_share=0.02,
likely_share=0.2, likely_share=0.2,
@@ -787,7 +787,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
("softmax", "sampling"), ("softmax", "sampling"),
), ),
rationale_hint=( 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, min_share=0.05,
likely_share=0.5, likely_share=0.5,
@@ -1218,8 +1218,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
FusionPatternSpec( FusionPatternSpec(
pattern="vLLM fused residual add + RMSNorm", pattern="vLLM fused residual add + RMSNorm",
candidate_path=( candidate_path=(
"vllm/_custom_ops.py" "vllm/_custom_ops.py<br>vllm/compilation/passes/fusion/rms_quant_fusion.py"
"<br>vllm/compilation/passes/fusion/rms_quant_fusion.py"
), ),
active_keywords=( active_keywords=(
"fused_add_rms_norm", "fused_add_rms_norm",
@@ -1236,8 +1235,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
FusionPatternSpec( FusionPatternSpec(
pattern="vLLM fused activation-and-mul", pattern="vLLM fused activation-and-mul",
candidate_path=( candidate_path=(
"vllm/_custom_ops.py" "vllm/_custom_ops.py<br>vllm/compilation/passes/fusion/act_quant_fusion.py"
"<br>vllm/compilation/passes/fusion/act_quant_fusion.py"
), ),
active_keywords=( active_keywords=(
"silu_and_mul", "silu_and_mul",
.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_proof_generator.py
+3 -1
View File
@@ -256,7 +256,9 @@ def _module_assign_names(text: str) -> set:
targets = ( targets = (
node.targets node.targets
if isinstance(node, ast.Assign) 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)} names |= {t.id for t in targets if isinstance(t, ast.Name)}
return names return names
.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_reproduction_utils.py
+33 -29
View File
@@ -173,9 +173,9 @@ def _find_unique_def(
if isinstance(node, definition) and node.name == name if isinstance(node, definition) and node.name == name
] ]
assert matches, f"{name} not found in {where}" assert matches, f"{name} not found in {where}"
assert ( assert len(matches) == 1, (
len(matches) == 1 f"{len(matches)} defs named {name} in {where}; pass from_class to disambiguate"
), f"{len(matches)} defs named {name} in {where}; pass from_class to disambiguate" )
return matches[0] return matches[0]
@@ -287,9 +287,9 @@ def _lowered_call_text(text: str, node: ast.Call) -> str:
""" """
receiver = node.args[0] receiver = node.args[0]
receiver_src = _node_slice(text, receiver) receiver_src = _node_slice(text, receiver)
assert ( assert "\n" not in receiver_src and "#" not in receiver_src, (
"\n" not in receiver_src and "#" not in receiver_src f"receiver {receiver_src!r} must be single-line and comment-free"
), f"receiver {receiver_src!r} must be single-line and comment-free" )
opener = _slice_span( opener = _slice_span(
text, text,
node.func.end_lineno, node.func.end_lineno,
@@ -722,9 +722,9 @@ class Repro:
) )
existing = [alias_text(a.name, a.asname) for a in node.names] existing = [alias_text(a.name, a.asname) for a in node.names]
added = alias_text(name, asname) added = alias_text(name, asname)
assert ( assert added not in existing, (
added not in existing f"{name!r} already imported from {module!r} in {rel}"
), f"{name!r} already imported from {module!r} in {rel}" )
rebuilt = f"from {module} import " + ", ".join(existing + [added]) + nl rebuilt = f"from {module} import " + ", ".join(existing + [added]) + nl
lines[node.lineno - 1 : node.end_lineno] = [rebuilt] lines[node.lineno - 1 : node.end_lineno] = [rebuilt]
_write_source(path, "".join(lines)) _write_source(path, "".join(lines))
@@ -825,9 +825,9 @@ class Repro:
for node in tree.body for node in tree.body
if isinstance(node, (ast.Import, ast.ImportFrom)) if isinstance(node, (ast.Import, ast.ImportFrom))
] ]
assert ( assert imports, (
imports f"no imports to anchor a new `if TYPE_CHECKING:` block in {rel}"
), f"no imports to anchor a new `if TYPE_CHECKING:` block in {rel}" )
insert_at = imports[-1].end_lineno insert_at = imports[-1].end_lineno
lines[insert_at:insert_at] = [ lines[insert_at:insert_at] = [
nl, nl,
@@ -864,9 +864,9 @@ class Repro:
replaced = lines[node.lineno - 1].replace( replaced = lines[node.lineno - 1].replace(
f"from {spelled} import", f"from {new_module} import", 1 f"from {spelled} import", f"from {new_module} import", 1
) )
assert ( assert replaced != lines[node.lineno - 1], (
replaced != lines[node.lineno - 1] f"import spelling {spelled!r} not found on its line in {rel}"
), f"import spelling {spelled!r} not found on its line in {rel}" )
lines[node.lineno - 1] = replaced lines[node.lineno - 1] = replaced
changed = True changed = True
assert changed, f"nested import of {name} from {old_module} not in {rel}" 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 ``self: Target`` annotation is dropped (redundant inside the class). The body is moved
verbatim; the formatter normalises the surrounding blank lines. verbatim; the formatter normalises the surrounding blank lines.
""" """
assert ( assert before is None or after is None, (
before is None or after is None "move_symbol: before and after are mutually exclusive"
), "move_symbol: before and after are mutually exclusive" )
def op(root: Path) -> None: def op(root: Path) -> None:
src_path = root / src src_path = root / src
@@ -1255,15 +1255,17 @@ class Repro:
targets = ( targets = (
node.targets node.targets
if isinstance(node, ast.Assign) 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)} names = {t.id for t in targets if isinstance(t, ast.Name)}
hit = names & dropped hit = names & dropped
if not hit: if not hit:
continue continue
assert len(names) == len( assert len(names) == len(targets), (
targets f"drop_assigns {sorted(hit)}: non-name targets in {src}"
), f"drop_assigns {sorted(hit)}: non-name targets in {src}" )
value_src = ast.unparse(node.value) if node.value is not None else None value_src = ast.unparse(node.value) if node.value is not None else None
for dropped_name in hit: for dropped_name in hit:
removed_assigns[dropped_name] = value_src removed_assigns[dropped_name] = value_src
@@ -1289,15 +1291,17 @@ class Repro:
else: else:
assign_spans.append((node.lineno, node.end_lineno)) assign_spans.append((node.lineno, node.end_lineno))
found_assigns |= hit found_assigns |= hit
assert ( assert found_assigns == dropped, (
found_assigns == dropped f"{dropped - found_assigns} not assigned in {src}"
), f"{dropped - found_assigns} not assigned in {src}" )
rederivable: dict[str, str | None] = {} rederivable: dict[str, str | None] = {}
for node in tree.body: for node in tree.body:
targets = ( targets = (
node.targets node.targets
if isinstance(node, ast.Assign) 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)] names = [t.id for t in targets if isinstance(t, ast.Name)]
if not names or set(names) & dropped: if not names or set(names) & dropped:
@@ -1383,9 +1387,9 @@ class Repro:
src_text = _read_source(src_path) src_text = _read_source(src_path)
assert src_text.count(body) == 1, f"block not found uniquely in {src}" assert src_text.count(body) == 1, f"block not found uniquely in {src}"
at = src_text.find(body) at = src_text.find(body)
assert ( assert at == 0 or src_text[at - 1] == "\n", (
at == 0 or src_text[at - 1] == "\n" f"block matches mid-line in {src}; it must start at a line boundary"
), f"block matches mid-line in {src}; it must start at a line boundary" )
_write_source(src_path, src_text.replace(body, call, 1)) _write_source(src_path, src_text.replace(body, call, 1))
dst_path = root / dst dst_path = root / dst
.claude/skills/mechanical-refactor-verify/scripts/tests/proof_generator/test_infer_moves.py
+3 -12
View File
@@ -370,12 +370,7 @@ def test_infer_recipe_module_level_def_shadowed_by_method_name(repo: Path) -> No
" return foo(x=self.x)\n" " return foo(x=self.x)\n"
), ),
"util.py": ( "util.py": (
"def keep():\n" "def keep():\n return 1\n\n\ndef foo(*, x):\n return x + 1\n"
" return 1\n"
"\n"
"\n"
"def foo(*, x):\n"
" return x + 1\n"
), ),
}, },
) )
@@ -437,9 +432,7 @@ def test_infer_recipe_move_leaving_a_forwarding_delegate(repo: Path) -> None:
_write( _write(
repo, repo,
**{ **{
"model.py": ( "model.py": ("class M:\n def work(self, x):\n return x + 1\n"),
"class M:\n" " def work(self, x):\n" " return x + 1\n"
),
"comp.py": "class C:\n def keep(self):\n return 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, repo,
**{ **{
"model.py": ( "model.py": (
"class M:\n" "class M:\n def work(self, x):\n return self.comp.work(x)\n"
" def work(self, x):\n"
" return self.comp.work(x)\n"
), ),
"comp.py": ( "comp.py": (
"class C:\n" "class C:\n"
.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_cli/cli_testlib.py
+2 -6
View File
@@ -2,13 +2,9 @@ import subprocess
from pathlib import Path from pathlib import Path
_PASSING_PROOF = ( _PASSING_PROOF = (
"import sys\n" 'import sys\nprint("PASS: reproduces the commit byte-for-byte.")\nsys.exit(0)\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"
) )
_FAILING_PROOF = 'import sys\nprint("RESIDUAL (2 lines):\\n+x\\n-y")\nsys.exit(1)\n'
def _git(repo: Path, *args: str) -> str: def _git(repo: Path, *args: str) -> str:
.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_add_imports.py
+2 -13
View File
@@ -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: 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.""" """With no TYPE_CHECKING block, one is created after the trailing module import."""
(tmp_path / "m.py").write_text( (tmp_path / "m.py").write_text(
"from typing import TYPE_CHECKING\n" "from typing import TYPE_CHECKING\n\nfrom a import X\n\n\ndef f():\n pass\n"
"\n"
"from a import X\n"
"\n"
"\n"
"def f():\n"
" pass\n"
) )
r = Repro("b", "t").add_typechecking_import("m.py", "from b import Y") r = Repro("b", "t").add_typechecking_import("m.py", "from b import Y")
_apply(r, tmp_path) _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: def test_add_typechecking_import_drops_a_lone_pass_placeholder(tmp_path: Path) -> None:
"""Populating a `pass`-only TYPE_CHECKING block replaces the placeholder.""" """Populating a `pass`-only TYPE_CHECKING block replaces the placeholder."""
(tmp_path / "m.py").write_text( (tmp_path / "m.py").write_text(
"from typing import TYPE_CHECKING\n" "from typing import TYPE_CHECKING\n\nif TYPE_CHECKING:\n pass\n\nx = 1\n"
"\n"
"if TYPE_CHECKING:\n"
" pass\n"
"\n"
"x = 1\n"
) )
r = Repro("b", "t").add_typechecking_import("m.py", "from b import Y") r = Repro("b", "t").add_typechecking_import("m.py", "from b import Y")
_apply(r, tmp_path) _apply(r, tmp_path)
.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_extract_symbols_to_new_module.py
+1 -5
View File
@@ -111,11 +111,7 @@ def test_extract_symbols_to_new_module_drops_relocated_assigns(tmp_path: Path) -
" return _FLAG\n" " return _FLAG\n"
) )
header = ( header = (
"from __future__ import annotations\n" "from __future__ import annotations\n\nimport os\n\n_FLAG = os.cpu_count()\n"
"\n"
"import os\n"
"\n"
"_FLAG = os.cpu_count()\n"
) )
r = Repro("b", "t").extract_symbols_to_new_module( r = Repro("b", "t").extract_symbols_to_new_module(
"src.py", "src.py",
.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_assign.py
+2 -14
View File
@@ -19,13 +19,7 @@ def test_move_assign_relocates_a_module_constant(tmp_path: Path) -> None:
_apply(r, tmp_path) _apply(r, tmp_path)
assert "LIMIT" not in (tmp_path / "src.py").read_text().split("def stay")[0] assert "LIMIT" not in (tmp_path / "src.py").read_text().split("def stay")[0]
assert (tmp_path / "dst.py").read_text() == ( assert (tmp_path / "dst.py").read_text() == (
"import sys\n" "import sys\n\nLIMIT = 480 # seconds\n\n\ndef keep():\n return 1\n"
"\n"
"LIMIT = 480 # seconds\n"
"\n"
"\n"
"def 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) _apply(r, tmp_path)
assert "LIMIT" not in (tmp_path / "src.py").read_text().split("def stay")[0] assert "LIMIT" not in (tmp_path / "src.py").read_text().split("def stay")[0]
assert (tmp_path / "dst.py").read_text() == ( assert (tmp_path / "dst.py").read_text() == (
"import sys\n" "import sys\n\nLIMIT: int = 480\n\n\ndef keep():\n return 1\n"
"\n"
"LIMIT: int = 480\n"
"\n"
"\n"
"def keep():\n"
" return 1\n"
) )
.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_symbol.py
+1 -7
View File
@@ -260,13 +260,7 @@ def test_move_symbol_dedent_leaves_string_literal_interior_lines(
) )
_apply(r, tmp_path) _apply(r, tmp_path)
assert (tmp_path / "dst.py").read_text() == ( assert (tmp_path / "dst.py").read_text() == (
"import os\n" "import os\n\ndef foo(self):\n s = '''raw\n partial\n'''\n return s\n"
"\n"
"def foo(self):\n"
" s = '''raw\n"
" partial\n"
"'''\n"
" return s\n"
) )
.claude/skills/sglang-prod-incident-triage/scripts/incident_artifact_tool.py
+4 -3
View File
@@ -101,8 +101,7 @@ def format_summary_line(filename: str, result: Dict[str, Any]) -> str:
if result.get("ok"): if result.get("ok"):
return f"{filename}: ok" return f"{filename}: ok"
return ( return (
f"{filename}: failed status={result.get('status')} " f"{filename}: failed status={result.get('status')} error={result.get('error')}"
f"error={result.get('error')}"
) )
@@ -617,7 +616,9 @@ def summarize_dump_file(path: Path, max_requests: int, preview_chars: int) -> st
time_span = ( time_span = (
max(timestamps) - min(timestamps) max(timestamps) - min(timestamps)
if len(timestamps) >= 2 if len(timestamps) >= 2
else 0.0 if len(timestamps) == 1 else None else 0.0
if len(timestamps) == 1
else None
) )
lines = [ lines = [
.pre-commit-config.yaml
+1 -4
View File
@@ -48,10 +48,7 @@ repos:
python/sglang/srt/grpc/.*_pb2\.pyi$| python/sglang/srt/grpc/.*_pb2\.pyi$|
python/sglang/srt/grpc/.*_pb2_grpc\.pyi$| python/sglang/srt/grpc/.*_pb2_grpc\.pyi$|
)$ )$
- repo: https://github.com/psf/black - id: ruff-format
rev: 26.1.0
hooks:
- id: black-jupyter
exclude: '^python/sglang/srt/grpc/.*_pb2\.py$|^python/sglang/srt/grpc/.*_pb2_grpc\.py$|^python/sglang/srt/grpc/.*_pb2\.pyi$|^python/sglang/srt/grpc/.*_pb2_grpc\.pyi$' 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 - repo: https://github.com/codespell-project/codespell
rev: v2.4.1 rev: v2.4.1
3rdparty/amd/tuning/benchmark_moe_rocm.py
Vendored
+4 -2
View File
@@ -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) configs = union_of_list_of_dicts(prune_configs_1, prune_configs_2)
print(f"{bs=} || {len(full_configs)=} | {len(prune_configs_1)=} | \ print(
{len(prune_configs_2)=} | {len(configs)=}") f"{bs=} || {len(full_configs)=} | {len(prune_configs_1)=} | \
{len(prune_configs_2)=} | {len(configs)=}"
)
best_config = None best_config = None
best_time_us = 1e20 best_time_us = 1e20
benchmark/asr/bench_sglang.py
+1 -1
View File
@@ -343,7 +343,7 @@ def run_evaluation(args):
print("\n" + "=" * 20 + " Sample Predictions " + "=" * 20) print("\n" + "=" * 20 + " Sample Predictions " + "=" * 20)
num_to_show = min(args.print_n, len(results)) num_to_show = min(args.print_n, len(results))
for i in range(num_to_show): for i in range(num_to_show):
print(f"Sample {i+1}:") print(f"Sample {i + 1}:")
print(f" REF: {references[i]}") print(f" REF: {references[i]}")
print(f" PRED: {predictions[i]}") print(f" PRED: {predictions[i]}")
print("-" * 40) print("-" * 40)
benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py
+1 -3
View File
@@ -243,9 +243,7 @@ def run(task, fi, tri, device, dtype, args):
) # noqa: E731 ) # noqa: E731
else: else:
inp = make_decode_inputs(B, H, HV, K, V, pool, device, dtype) inp = make_decode_inputs(B, H, HV, K, V, pool, device, dtype)
corr = lambda kern: call_decode( corr = lambda kern: call_decode(kern, inp, inp["ssm"].clone()) # noqa: E731
kern, inp, inp["ssm"].clone()
) # noqa: E731
ssm_t = inp["ssm"].clone() ssm_t = inp["ssm"].clone()
timed = lambda kern: call_decode(kern, inp, ssm_t) # noqa: E731 timed = lambda kern: call_decode(kern, inp, ssm_t) # noqa: E731
benchmark/gsm8k/bench_sglang.py
+3 -3
View File
@@ -53,9 +53,9 @@ def main(args):
if args.enable_thinking: if args.enable_thinking:
from transformers import AutoTokenizer from transformers import AutoTokenizer
assert ( assert args.tokenizer_path is not None, (
args.tokenizer_path is not None "--tokenizer-path is required when --enable-thinking is set"
), "--tokenizer-path is required when --enable-thinking is set" )
tokenizer = AutoTokenizer.from_pretrained( tokenizer = AutoTokenizer.from_pretrained(
args.tokenizer_path, trust_remote_code=True args.tokenizer_path, trust_remote_code=True
) )
benchmark/hf3fs/bench_client.py
+4 -4
View File
@@ -14,11 +14,11 @@ def print_stats(x: List[int]):
x = sorted(x) x = sorted(x)
lenx = len(x) lenx = len(x)
print( print(
f"mean = {sum(x)/len(x):.2f}, " f"mean = {sum(x) / len(x):.2f}, "
f"min = {min(x):.2f}, " f"min = {min(x):.2f}, "
f"p25 = {x[int(lenx*0.25)]:.2f}, " f"p25 = {x[int(lenx * 0.25)]:.2f}, "
f"p50 = {x[int(lenx*0.5)]:.2f}, " f"p50 = {x[int(lenx * 0.5)]:.2f}, "
f"p75 = {x[int(lenx*0.75)]:.2f}, " f"p75 = {x[int(lenx * 0.75)]:.2f}, "
f"max = {max(x):.2f}" f"max = {max(x):.2f}"
) )
benchmark/hf3fs/bench_storage.py
+4 -4
View File
@@ -18,11 +18,11 @@ def print_stats(x: List[int]):
x = sorted(x) x = sorted(x)
lenx = len(x) lenx = len(x)
print( print(
f"mean = {sum(x)/len(x):.2f}, " f"mean = {sum(x) / len(x):.2f}, "
f"min = {min(x):.2f}, " f"min = {min(x):.2f}, "
f"p25 = {x[int(lenx*0.25)]:.2f}, " f"p25 = {x[int(lenx * 0.25)]:.2f}, "
f"p50 = {x[int(lenx*0.5)]:.2f}, " f"p50 = {x[int(lenx * 0.5)]:.2f}, "
f"p75 = {x[int(lenx*0.75)]:.2f}, " f"p75 = {x[int(lenx * 0.75)]:.2f}, "
f"max = {max(x):.2f}" f"max = {max(x):.2f}"
) )
benchmark/hf3fs/bench_zerocopy.py
+2 -2
View File
@@ -109,7 +109,7 @@ elif hicache_mem_layout == "layer_first":
for operation in operations: for operation in operations:
cache_controller.generic_page_backup(operation, batch_size=128) cache_controller.generic_page_backup(operation, batch_size=128)
tok = time.monotonic() tok = time.monotonic()
print(f"{tok-tik:.6f} s") print(f"{tok - tik:.6f} s")
operations = [ operations = [
PrefetchOperation( PrefetchOperation(
@@ -137,4 +137,4 @@ elif hicache_mem_layout == "layer_first":
for operation in operations: for operation in operations:
cache_controller.generic_page_transfer(operation, batch_size=128) cache_controller.generic_page_transfer(operation, batch_size=128)
tok = time.monotonic() tok = time.monotonic()
print(f"{tok-tik:.6f} s") print(f"{tok - tik:.6f} s")
benchmark/hicache/bench_mix.py
+3 -3
View File
@@ -457,7 +457,7 @@ class WorkloadGenerator:
try: try:
user_data, response = self.response_queue.get(timeout=10) user_data, response = self.response_queue.get(timeout=10)
logger.info( 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: if not response.success:
raise ValueError(f"Request failed with error: {response.error}") raise ValueError(f"Request failed with error: {response.error}")
@@ -540,10 +540,10 @@ class WorkloadGenerator:
output_stats = self.user_generator.output_stats output_stats = self.user_generator.output_stats
print(f"round_ratios: {user_stats}") print(f"round_ratios: {user_stats}")
print( 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( 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 return performance_data
benchmark/hicache/bench_serving.py
+3 -3
View File
@@ -75,9 +75,9 @@ async def async_request_openai_completions(
pbar: Optional[tqdm] = None, pbar: Optional[tqdm] = None,
) -> RequestFuncOutput: ) -> RequestFuncOutput:
api_url = request_func_input.api_url api_url = request_func_input.api_url
assert api_url.endswith( assert api_url.endswith("completions"), (
"completions" "OpenAI Completions API URL must end with 'completions'."
), "OpenAI Completions API URL must end with 'completions'." )
async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
payload = { payload = {
benchmark/hicache/nextqa.py
+1 -1
View File
@@ -120,7 +120,7 @@ class NExTQALoader(VideoLoader):
video = Video(video_path, num_frames) video = Video(video_path, num_frames)
prompt = entry["question"] + "?" prompt = entry["question"] + "?"
if self.task == "MC": # add choices 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) return VideoPrompt(video_path, num_frames, prompt)
def __iter__(self): def __iter__(self):
benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py
+3 -3
View File
@@ -149,9 +149,9 @@ def _check_correctness():
cos = torch.nn.functional.cosine_similarity( cos = torch.nn.functional.cosine_similarity(
(q.float() * scale).flatten(), ref_deq.flatten(), dim=0 (q.float() * scale).flatten(), ref_deq.flatten(), dim=0
).item() ).item()
assert ( assert cos > 0.99, (
cos > 0.99 f"{provider} h={hidden_size} residual={add_residual} cos={cos:.4f}"
), f"{provider} h={hidden_size} residual={add_residual} cos={cos:.4f}" )
print("correctness check passed (all fused providers vs unfused within FP8)") print("correctness check passed (all fused providers vs unfused within FP8)")
benchmark/kernels/deepep/deepep_utils.py
+3 -3
View File
@@ -191,9 +191,9 @@ def bench_kineto(
kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names
assert all([isinstance(name, str) for name in kernel_names]) assert all([isinstance(name, str) for name in kernel_names])
for name in kernel_names: for name in kernel_names:
assert ( assert sum([name in line for line in prof_lines]) == 1, (
sum([name in line for line in prof_lines]) == 1 f"Errors of the kernel {name} in the profiling table"
), f"Errors of the kernel {name} in the profiling table" )
# Save chrome traces # Save chrome traces
if trace_path is not None: if trace_path is not None:
benchmark/kernels/deepep/tuning_deepep.py
+21 -15
View File
@@ -155,7 +155,7 @@ def test_main(
for with_topk in (False, True): for with_topk in (False, True):
if local_rank == 0: if local_rank == 0:
print( 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, flush=True,
end="", end="",
) )
@@ -198,9 +198,9 @@ def test_main(
# Checks # Checks
recv_gbl_rank_prefix_sum = handle[-4] recv_gbl_rank_prefix_sum = handle[-4]
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size( assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(0), (
0 f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}" )
assert ( assert (
gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist() gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
== recv_num_tokens_per_expert_list == recv_num_tokens_per_expert_list
@@ -325,11 +325,14 @@ def test_main(
tune_args = {"x": current_x, "handle": handle, "config": config} tune_args = {"x": current_x, "handle": handle, "config": config}
t = bench(lambda: buffer.dispatch(**tune_args))[0] t = bench(lambda: buffer.dispatch(**tune_args))[0]
if t < best_time: if t < best_time:
best_time, best_results = t, ( best_time, best_results = (
num_sms, t,
nvl_chunk_size, (
rdma_chunk_size, num_sms,
config_kwargs, nvl_chunk_size,
rdma_chunk_size,
config_kwargs,
),
) )
if local_rank == 0: if local_rank == 0:
print( print(
@@ -338,7 +341,7 @@ def test_main(
) )
if local_rank == 0: if local_rank == 0:
print( 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, flush=True,
) )
print("", flush=True) print("", flush=True)
@@ -399,11 +402,14 @@ def test_main(
flush=True, flush=True,
) )
if t < best_time: if t < best_time:
best_time, best_results = t, ( best_time, best_results = (
num_sms, t,
nvl_chunk_size, (
rdma_chunk_size, num_sms,
config_kwargs, nvl_chunk_size,
rdma_chunk_size,
config_kwargs,
),
) )
if local_rank == 0: if local_rank == 0:
benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py
+1 -2
View File
@@ -59,7 +59,6 @@ def tl_gemm(
bx, bx,
by, by,
): ):
A_shared = T.alloc_shared(A_shared_shape, in_dtype) A_shared = T.alloc_shared(A_shared_shape, in_dtype)
B_shared = T.alloc_shared(B_shared_shape, in_dtype) B_shared = T.alloc_shared(B_shared_shape, in_dtype)
C_shared = T.alloc_shared(C_shared_shape, out_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 tflops = flops / (ms * 1e-3) / 1e12
# Print shape-specific results with TFLOPS # 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 ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms
return benchmark return benchmark
benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py
+1 -1
View File
@@ -224,7 +224,7 @@ def _benchmark(m, n, k, tp_size, provider):
tflops = flops / (ms * 1e-3) / 1e12 tflops = flops / (ms * 1e-3) / 1e12
# Print shape-specific results with TFLOPS # 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 return ms, max_ms, min_ms
benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py
+1 -1
View File
@@ -435,7 +435,7 @@ def get_benchmark(tp_size):
flops = 2 * m * n * k # multiply-adds flops = 2 * m * n * k # multiply-adds
tflops = flops / (ms * 1e-3) / 1e12 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 ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms
return benchmark return benchmark
benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py
+1 -2
View File
@@ -243,8 +243,7 @@ def main():
else: else:
speedup = f"{legacy_us / us:.2f}x" speedup = f"{legacy_us / us:.2f}x"
print( print(
f"{s_q:>5} {kv_len:>7} {name:>10} {us:>9.1f} " f"{s_q:>5} {kv_len:>7} {name:>10} {us:>9.1f} {tbps:>9.3f} {speedup:>8}"
f"{tbps:>9.3f} {speedup:>8}"
) )
print() print()
benchmark/kernels/elementwise/benchmark_concat_mla.py
+1 -3
View File
@@ -143,9 +143,7 @@ output_exp = execute_and_get_output(fn_cuda, data)
if not torch.all(output_ref == output_exp): if not torch.all(output_ref == output_exp):
abs_delta = torch.abs(output_ref - output_exp) abs_delta = torch.abs(output_ref - output_exp)
raise AssertionError( raise AssertionError(
f"{output_ref=} {output_exp=} " f"{output_ref=} {output_exp=} {abs_delta=} {torch.argwhere(abs_delta != 0.0)=} "
f"{abs_delta=} "
f"{torch.argwhere(abs_delta != 0.0)=} "
) )
benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py
+1 -3
View File
@@ -535,7 +535,6 @@ class BestConfigTrace:
class BenchmarkWorker: class BenchmarkWorker:
def __init__(self, seed: int, server_args: ServerArgs) -> None: def __init__(self, seed: int, server_args: ServerArgs) -> None:
torch.set_default_device("cuda") torch.set_default_device("cuda")
torch.cuda.manual_seed_all(0) 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] down_use_tma_map[block_m] = time_cost_all[2] > time_cost_all[3]
print( print(
f"Round 1 done. Down TMA decisions per BLOCK_SIZE_M: " f"Round 1 done. Down TMA decisions per BLOCK_SIZE_M: {down_use_tma_map}"
f"{down_use_tma_map}"
) )
# === Round 2: Up with c_sorted from round 1 === # === Round 2: Up with c_sorted from round 1 ===
benchmark/kernels/lora_csgmv/tune_lora_csgmv.py
+8 -8
View File
@@ -470,9 +470,9 @@ def _tune_shrink(
device: torch.device, device: torch.device,
) -> tuple: ) -> tuple:
"""Tune shrink kernel for one layer type. Returns (best_configs, results).""" """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"Tuning SHRINK — {label} (K={K}, N={N}, slices={num_slices})")
print(f"{'='*80}") print(f"{'=' * 80}")
search = get_shrink_search_space() search = get_shrink_search_space()
print(f"Search space: {len(search)} configs") print(f"Search space: {len(search)} configs")
@@ -508,7 +508,7 @@ def _tune_shrink(
best_config = config best_config = config
if (i + 1) % 20 == 0: if (i + 1) % 20 == 0:
print( 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) best_configs[chunk_size] = sort_config(best_config)
@@ -533,9 +533,9 @@ def _tune_expand(
device: torch.device, device: torch.device,
) -> tuple: ) -> tuple:
"""Tune expand kernel for one layer type. Returns (best_configs, results).""" """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"Tuning EXPAND — {label} (output_dim={output_dim}, slices={num_slices})")
print(f"{'='*80}") print(f"{'=' * 80}")
search = get_expand_search_space() search = get_expand_search_space()
print(f"Search space: {len(search)} configs") print(f"Search space: {len(search)} configs")
@@ -584,7 +584,7 @@ def _tune_expand(
best_config = config best_config = config
if (i + 1) % 50 == 0: if (i + 1) % 50 == 0:
print( 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) best_configs[chunk_size] = sort_config(best_config)
@@ -673,9 +673,9 @@ def main(args: argparse.Namespace):
) )
# --- Summary --- # --- Summary ---
print(f"\n{'='*80}") print(f"\n{'=' * 80}")
print(f"SUMMARY") print(f"SUMMARY")
print(f"{'='*80}") print(f"{'=' * 80}")
print( print(
f"\n{'layer':<10} {'kernel':<8} {'K/dim':>6} {'chunk':>6}" f"\n{'layer':<10} {'kernel':<8} {'K/dim':>6} {'chunk':>6}"
f" {'baseline':>10} {'tuned':>10} {'speedup':>8} config" f" {'baseline':>10} {'tuned':>10} {'speedup':>8} config"
benchmark/lean_kernel_sweep.py
+5 -3
View File
@@ -137,19 +137,21 @@ def main():
# b32 x 128K on the 8-KV-head config exceeds the microbench's single # 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. # contiguous KV tensor (faults the GPU); real serving uses a paged pool.
if H_KV == 8 and B == 32 and S == 131072: 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,") rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,skip,")
continue continue
try: try:
std, lean, cos, gate = run(H_Q, H_KV, B, S) std, lean, cos, gate = run(H_Q, H_KV, B, S)
except torch.cuda.OutOfMemoryError: except torch.cuda.OutOfMemoryError:
torch.cuda.empty_cache() 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,") rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,OOM,")
continue continue
sp = std / lean sp = std / lean
print( 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( rows.append(
f"{name},{H_Q},{H_KV},{B},{S},{std:.4f},{lean:.4f},{sp:.4f},{cos:.4f},{int(gate)}" f"{name},{H_Q},{H_KV},{B},{S},{std:.4f},{lean:.4f},{sp:.4f},{cos:.4f},{int(gate)}"
benchmark/mmlu/bench_sglang.py
+1 -1
View File
@@ -163,7 +163,7 @@ def main(args):
pt = 0 pt = 0
for subject, num_qs in zip(subjects[: args.nsub], num_questions): for subject, num_qs in zip(subjects[: args.nsub], num_questions):
print( 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 pt += num_qs
assert pt == len(cors) assert pt == len(cors)
benchmark/ocr/bench_sglang.py
+1 -1
View File
@@ -502,7 +502,7 @@ async def process_sample(
} }
) )
print( 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, flush=True,
) )
# Print OCR output (truncate long outputs) # Print OCR output (truncate long outputs)
benchmark/ocr/generate_report.py
+4 -4
View File
@@ -159,12 +159,12 @@ def _render_sample(sample: dict, failures_only: bool) -> str:
if failures_only and passed == total and not error: if failures_only and passed == total and not error:
return "" 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" header_cls = "fail" if (error or passed < total) else "pass"
parts = [f'<div class="sample">'] parts = [f'<div class="sample">']
parts.append( parts.append(
f'<details {"open" if (error or passed < total) else ""}>' f"<details {'open' if (error or passed < total) else ''}>"
f'<summary class="sample-header {header_cls}">' f'<summary class="sample-header {header_cls}">'
f"<span>📄 {html.escape(pdf)} &nbsp;·&nbsp; page {page}</span>" f"<span>📄 {html.escape(pdf)} &nbsp;·&nbsp; page {page}</span>"
f"<span>" f"<span>"
@@ -218,7 +218,7 @@ def _render_sample(sample: dict, failures_only: bool) -> str:
f'<div class="rendered">{_latex_to_display(latex)}</div>' f'<div class="rendered">{_latex_to_display(latex)}</div>'
) )
elif ttype in ("present", "absent", "text_presence", "text_absence"): elif ttype in ("present", "absent", "text_presence", "text_absence"):
parts.append(f'<pre>{html.escape(ti.get("text", ""))}</pre>') parts.append(f"<pre>{html.escape(ti.get('text', ''))}</pre>")
elif ttype in ("order", "natural_reading_order"): elif ttype in ("order", "natural_reading_order"):
before = ti.get("before", "") before = ti.get("before", "")
after = ti.get("after", "") after = ti.get("after", "")
@@ -236,7 +236,7 @@ def _render_sample(sample: dict, failures_only: bool) -> str:
parts.append(f'<div class="rendered">{m}</div>') parts.append(f'<div class="rendered">{m}</div>')
if len(matches) > 6: if len(matches) > 6:
parts.append( parts.append(
f'<p style="color:#888;font-size:0.8em">… and {len(matches)-6} more</p>' f'<p style="color:#888;font-size:0.8em">… and {len(matches) - 6} more</p>'
) )
else: else:
parts.append( parts.append(
benchmark/prefill_only/util.py
+3 -4
View File
@@ -383,14 +383,14 @@ async def send_warmup_requests(
http_url, data=request_json, headers=headers http_url, data=request_json, headers=headers
) as resp: ) as resp:
if resp.status == 200: 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: else:
print( 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: 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") print("HTTP warmup requests completed")
@@ -745,7 +745,6 @@ async def run_generic_benchmark(
async with aiohttp.ClientSession( async with aiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(total=300) timeout=aiohttp.ClientTimeout(total=300)
) as session: ) as session:
# Send START_PROFILE if profiling is enabled # Send START_PROFILE if profiling is enabled
if config.profile: if config.profile:
await send_profile_request("START_PROFILE", http_url, session=session) await send_profile_request("START_PROFILE", http_url, session=session)
benchmark/scheduler/bench_token_storage.py
+6 -4
View File
@@ -254,7 +254,7 @@ def print_breakdown(title: str, results: dict[str, dict[str, float]]) -> None:
def microbench_torch_tensor_paths( def microbench_torch_tensor_paths(
sizes: tuple[int, ...] = (1_000, 10_000, 100_000) sizes: tuple[int, ...] = (1_000, 10_000, 100_000),
) -> None: ) -> None:
"""Compare three CPU-buffer -> pinned cuda tensor paths. """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", "(C) from_numpy(frombuf(array('q'))).pin() -> cuda",
lambda x: torch.from_numpy(np.frombuffer(x, dtype=np.int64)) lambda x: (
.pin_memory() torch.from_numpy(np.frombuffer(x, dtype=np.int64))
.to("cuda", non_blocking=True), .pin_memory()
.to("cuda", non_blocking=True)
),
), ),
]: ]:
cells = [] cells = []
docs/demo/deepseek_v4_flash.ipynb
+2 -2
View File
@@ -347,11 +347,11 @@
" try:\n", " try:\n",
" paper[\"full_text\"] = download_and_extract(paper)\n", " paper[\"full_text\"] = download_and_extract(paper)\n",
" print(\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", " )\n",
" except Exception as e:\n", " except Exception as e:\n",
" paper[\"full_text\"] = None\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}\")"
] ]
}, },
{ {
examples/frontend_language/usage/llava_video/srt_example_llava_v.py
-1
View File
@@ -124,7 +124,6 @@ def batch(video_dir, save_dir, cur_chunk, num_chunks, num_frames=16, batch_size=
if __name__ == "__main__": if __name__ == "__main__":
url = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4" url = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4"
cache_dir = os.path.expanduser("~/.cache") cache_dir = os.path.expanduser("~/.cache")
examples/frontend_language/usage/readme_examples.py
+1 -1
View File
@@ -31,7 +31,7 @@ def tip_suggestion(s):
forks = s.fork(2) forks = s.fork(2)
for i, f in enumerate(forks): 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") f += sgl.gen(f"detailed_tip", max_tokens=256, stop="\n\n")
s += "Tip 1:" + forks[0]["detailed_tip"] + "\n" s += "Tip 1:" + forks[0]["detailed_tip"] + "\n"
examples/profiler/nsys_profile_tools/gputrc2graph.py
+4 -4
View File
@@ -86,7 +86,7 @@ class GPUTrace2Graph:
# Update current_end for overlapping intervals # Update current_end for overlapping intervals
for i in range(1, len(df)): for i in range(1, len(df)):
if i % display_units == 0: 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 starts[i] <= current_end:
if ends[i] > current_end: if ends[i] > current_end:
# Partial overlap # Partial overlap
@@ -182,9 +182,9 @@ class GPUTrace2Graph:
def is_valid_file(self, base_file): def is_valid_file(self, base_file):
"""asserts if base_file is non-existent or is empty""" """asserts if base_file is non-existent or is empty"""
assert ( assert os.path.isfile(base_file) and os.path.getsize(base_file) > 0, (
os.path.isfile(base_file) and os.path.getsize(base_file) > 0 f"{base_file} doesn't exist or is empty"
), f"{base_file} doesn't exist or is empty" )
def should_gen_file(self, new_file, base_file): def should_gen_file(self, new_file, base_file):
"""figure out if new file should be generated from base_file""" """figure out if new file should be generated from base_file"""
examples/runtime/engine/fastapi_engine_inference.py
+1 -1
View File
@@ -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.""" """Sends generation requests to the running server for a list of prompts."""
# Iterate through prompts and send requests # Iterate through prompts and send requests
for i, prompt in enumerate(prompts): 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 = { payload = {
"prompt": prompt, "prompt": prompt,
"max_new_tokens": max_new_tokens, "max_new_tokens": max_new_tokens,
examples/runtime/engine/offline_batch_inference_qwen_1m.py
+2 -5
View File
@@ -17,8 +17,7 @@ def load_prompt() -> str:
# https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/1m.txt # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/1m.txt
with urlopen( with urlopen(
"https://qianwen-res.oss-cn-beijing.aliyuncs.com" "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/64k.txt",
"/Qwen2.5-1M/test-data/64k.txt",
timeout=5, timeout=5,
) as response: ) as response:
prompt = response.read().decode("utf-8") prompt = response.read().decode("utf-8")
@@ -41,9 +40,7 @@ def process_requests(llm: sgl.Engine, prompts: list[str]) -> None:
for output in outputs: for output in outputs:
prompt_token_ids = output["meta_info"]["prompt_tokens"] prompt_token_ids = output["meta_info"]["prompt_tokens"]
generated_text = output["text"] generated_text = output["text"]
print( print(f"Prompt length: {prompt_token_ids}, Generated text: {generated_text!r}")
f"Prompt length: {prompt_token_ids}, " f"Generated text: {generated_text!r}"
)
# Create an LLM. # Create an LLM.
examples/runtime/qwen3_vl_reranker.py
+3 -3
View File
@@ -44,7 +44,7 @@ def rerank_text_only():
print("Results (sorted by relevance):") print("Results (sorted by relevance):")
for i, result in enumerate(results): 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() print()
@@ -99,7 +99,7 @@ def rerank_with_images():
print("Results (sorted by relevance):") print("Results (sorted by relevance):")
for i, result in enumerate(results): 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() print()
@@ -149,7 +149,7 @@ def rerank_multimodal_query():
print("Results (sorted by relevance):") print("Results (sorted by relevance):")
for i, result in enumerate(results): 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() print()
examples/usage/modelopt_quantize_and_export.py
+1 -1
View File
@@ -213,7 +213,7 @@ def deploy_exported_model(
outputs = llm.generate(prompts, sampling_params) outputs = llm.generate(prompts, sampling_params)
for i, output in enumerate(outputs): 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(f"Output: {output['text']}")
print() print()
examples/usage/reasoning_aware_compression/rac_collect_traces.py
+1 -1
View File
@@ -335,7 +335,7 @@ def report(manifest: TraceManifest) -> None:
print(f" decode tokens : {manifest.num_decode_tokens}") print(f" decode tokens : {manifest.num_decode_tokens}")
print(f" total tokens : {total}") print(f" total tokens : {total}")
print(f" decode share : {decode_share:.1%}") 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": if manifest.calibration_mode == "rac":
print( print(
"\nThe decode share is the activation mass that prompt-only " "\nThe decode share is the activation mass that prompt-only "
examples/usage/reasoning_aware_compression/rac_serve_and_eval.py
+1 -1
View File
@@ -209,7 +209,7 @@ def report(results: List[EvalResult]) -> None:
f"{result.model_path:<{width}} " f"{result.model_path:<{width}} "
f"{result.accuracy:>7.3f} " f"{result.accuracy:>7.3f} "
f"{result.mean_completion_tokens:>16.0f} " 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: if len(results) > 1:
experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py
+12 -12
View File
@@ -194,9 +194,9 @@ def _direct_warm(worker_url: str, model_id: str, prefix: str) -> None:
}, },
timeout=60.0, timeout=60.0,
) )
assert ( assert r.status_code == 200, (
r.status_code == 200 f"direct warm to {worker_url} failed: HTTP {r.status_code} {r.text!r}"
), 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: 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) after = _success_counts_by_worker(router_url)
deltas = {w: after.get(w, 0) - before.get(w, 0) for w in set(after) | set(before)} 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] winners = [w for w, d in deltas.items() if d > 0]
assert ( assert len(winners) == 1, (
len(winners) == 1 f"expected exactly one worker delta on {router_url}, got {deltas}"
), f"expected exactly one worker delta on {router_url}, got {deltas}" )
return winners[0] return winners[0]
@@ -309,15 +309,15 @@ def test_routers_route_by_prefix_content(
landed = _route_through( landed = _route_through(
router.base_url, spec["model"], PREFIX_X router.base_url, spec["model"], PREFIX_X
) )
assert ( assert landed == worker_x.url, (
landed == worker_x.url f"router {label}: PREFIX_X must route to {worker_x.url}; landed on {landed}"
), f"router {label}: PREFIX_X must route to {worker_x.url}; landed on {landed}" )
landed = _route_through( landed = _route_through(
router.base_url, spec["model"], PREFIX_Y router.base_url, spec["model"], PREFIX_Y
) )
assert ( assert landed == worker_y.url, (
landed == worker_y.url f"router {label}: PREFIX_Y must route to {worker_y.url}; landed on {landed}"
), f"router {label}: PREFIX_Y must route to {worker_y.url}; landed on {landed}" )
except Exception: except Exception:
_dump_logs(logs) _dump_logs(logs)
raise raise
experimental/sgl-router/tests/e2e/chat_completions/test_validation.py
+6 -6
View File
@@ -50,9 +50,9 @@ def test_chat_non_streaming_returns_assistant_message(
body = resp.json() body = resp.json()
choice = body["choices"][0] choice = body["choices"][0]
assert choice["message"]["role"] == "assistant" assert choice["message"]["role"] == "assistant"
assert choice["message"][ assert choice["message"]["content"], (
"content" f"empty assistant content: {choice!r}"
], f"empty assistant content: {choice!r}" )
assert choice.get("finish_reason"), choice assert choice.get("finish_reason"), choice
finally: finally:
gpu_allocator.release(gpu) gpu_allocator.release(gpu)
@@ -91,8 +91,8 @@ def test_chat_streaming_emits_sse_chunks_with_done(
if line.startswith("data:"): if line.startswith("data:"):
chunks.append(line.strip()) chunks.append(line.strip())
assert len(chunks) >= 2, f"expected >=2 SSE chunks, got: {chunks}" assert len(chunks) >= 2, f"expected >=2 SSE chunks, got: {chunks}"
assert any( assert any("[DONE]" in c for c in chunks), (
"[DONE]" in c for c in chunks f"no [DONE] terminator in stream: {chunks}"
), f"no [DONE] terminator in stream: {chunks}" )
finally: finally:
gpu_allocator.release(gpu) gpu_allocator.release(gpu)
experimental/sgl-router/tests/e2e/k8s_integration/test_discovery.py
+11 -9
View File
@@ -66,15 +66,17 @@ def test_router_discovers_multiple_workers(router_url):
# Scale down to 1 — router should still route after reconverging # Scale down to 1 — router should still route after reconverging
_scale_fake_worker(1) _scale_fake_worker(1)
_poll_until( _poll_until(
lambda: httpx.post( lambda: (
f"{router_url}/v1/chat/completions", httpx.post(
json={ f"{router_url}/v1/chat/completions",
"model": "tiny", json={
"messages": [{"role": "user", "content": "post-scale-down"}], "model": "tiny",
}, "messages": [{"role": "user", "content": "post-scale-down"}],
timeout=10.0, },
).status_code timeout=10.0,
== 200, ).status_code
== 200
),
"router routes after scale-down to 1", "router routes after scale-down to 1",
timeout=60, timeout=60,
interval=3, interval=3,
experimental/sgl-router/tests/e2e/test_chat_smoke.py
+6 -6
View File
@@ -16,9 +16,9 @@ def test_models(router: str) -> None:
assert resp.status_code == 200, resp.text assert resp.status_code == 200, resp.text
data = resp.json() data = resp.json()
ids = [m["id"] for m in data.get("data", [])] ids = [m["id"] for m in data.get("data", [])]
assert any( assert any(MODEL in mid for mid in ids), (
MODEL in mid for mid in ids f"Model {MODEL!r} not found in /v1/models response: {ids}"
), f"Model {MODEL!r} not found in /v1/models response: {ids}" )
def test_chat_non_streaming(router: str) -> None: def test_chat_non_streaming(router: str) -> None:
@@ -59,6 +59,6 @@ def test_chat_streaming(router: str) -> None:
chunks.append(line) chunks.append(line)
assert len(chunks) >= 2, f"Expected >=2 SSE chunks, got {len(chunks)}: {chunks}" assert len(chunks) >= 2, f"Expected >=2 SSE chunks, got {len(chunks)}: {chunks}"
assert any( assert any("[DONE]" in c for c in chunks), (
"[DONE]" in c for c in chunks f"No [DONE] chunk found in SSE stream: {chunks}"
), f"No [DONE] chunk found in SSE stream: {chunks}" )
experimental/sgl-router/tests/e2e/test_tokenize_smoke.py
+6 -6
View File
@@ -20,9 +20,9 @@ def test_tokenize_round_trip(router: str) -> None:
) )
assert tok_resp.status_code == 200, tok_resp.text assert tok_resp.status_code == 200, tok_resp.text
tokens = tok_resp.json()["tokens"] tokens = tok_resp.json()["tokens"]
assert ( assert isinstance(tokens, list) and len(tokens) > 0, (
isinstance(tokens, list) and len(tokens) > 0 f"Expected non-empty token list, got: {tokens}"
), f"Expected non-empty token list, got: {tokens}" )
# Detokenize # Detokenize
detok_resp = httpx.post( 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 assert detok_resp.status_code == 200, detok_resp.text
recovered = detok_resp.json()["text"] recovered = detok_resp.json()["text"]
assert ( assert TEXT in recovered or recovered in TEXT, (
TEXT in recovered or recovered in TEXT f"Round-trip mismatch: original={TEXT!r}, recovered={recovered!r}"
), f"Round-trip mismatch: original={TEXT!r}, recovered={recovered!r}" )
experimental/sgl-router/tests/scripts/generate_parity_fixtures.py
+1 -1
View File
@@ -85,7 +85,7 @@ def load_tokenizer_with_fallback(primary, fallback, slug):
raise raise
continue continue
raise RuntimeError( raise RuntimeError(
f"No accessible tokenizer for slug={slug} " f"(tried: {primary}, {fallback})" f"No accessible tokenizer for slug={slug} (tried: {primary}, {fallback})"
) )
python/sglang/benchmark/dspark_sps_profiler.py
+1 -2
View File
@@ -687,8 +687,7 @@ def run_one_round(
rank_rows = fetch_rank_rows(base_url=context.base_url) rank_rows = fetch_rank_rows(base_url=context.base_url)
if len(rank_rows) != len(watermarks): if len(rank_rows) != len(watermarks):
raise RuntimeError( raise RuntimeError(
f"DP rank count changed mid-profile: {len(watermarks)} -> " f"DP rank count changed mid-profile: {len(watermarks)} -> {len(rank_rows)}."
f"{len(rank_rows)}."
) )
new_rank_rows = [ new_rank_rows = [
[row for row in rows if row.forward_ct > watermark] [row for row in rows if row.forward_ct > watermark]
python/sglang/benchmark/offline_throughput.py
+5 -5
View File
@@ -136,13 +136,13 @@ class BenchArgs:
"--gsp-system-prompt-len", "--gsp-system-prompt-len",
type=int, type=int,
default=BenchArgs.gsp_system_prompt_len, 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( parser.add_argument(
"--gsp-question-len", "--gsp-question-len",
type=int, type=int,
default=BenchArgs.gsp_question_len, 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( parser.add_argument(
"--gsp-output-len", "--gsp-output-len",
@@ -259,9 +259,9 @@ def throughput_test_once(
] ]
if profile: if profile:
assert ( assert "SGLANG_TORCH_PROFILER_DIR" in os.environ, (
"SGLANG_TORCH_PROFILER_DIR" in os.environ "Please set SGLANG_TORCH_PROFILER_DIR."
), "Please set SGLANG_TORCH_PROFILER_DIR." )
os.makedirs(os.environ["SGLANG_TORCH_PROFILER_DIR"], exist_ok=True) os.makedirs(os.environ["SGLANG_TORCH_PROFILER_DIR"], exist_ok=True)
known_files = None known_files = None
backend.start_profile( backend.start_profile(
python/sglang/benchmark/one_batch_server.py
+8 -8
View File
@@ -486,7 +486,7 @@ def _warmup_cache(
return return
print( 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)" f"({cached_token_len} tokens per request)"
) )
# Create prefix input_ids for cache warming # 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}." f"\nInput lens: {bench_args.input_len}. Output lens: {bench_args.output_len}."
) )
if bench_args.cache_hit_rate > 0.0: 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" summary += "\n"
if is_blackwell(): if is_blackwell():
@@ -1241,9 +1241,9 @@ def run_benchmark_internal(
skip_max_running_requests_threshold = float("inf") skip_max_running_requests_threshold = float("inf")
skip_token_capacity_threshold = float("inf") skip_token_capacity_threshold = float("inf")
else: else:
assert ( assert max_running_requests_per_dp > 0, (
max_running_requests_per_dp > 0 f"effective_max_running_requests_per_dp is not set, {max_running_requests_per_dp=}"
), 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 skip_max_running_requests_threshold = max_running_requests_per_dp * dp_size
print(f"{max_running_requests_per_dp=}") print(f"{max_running_requests_per_dp=}")
@@ -1288,9 +1288,9 @@ def run_benchmark_internal(
"--lora-request-distribution=distinct/skewed requires more than " "--lora-request-distribution=distinct/skewed requires more than "
"one adapter via --lora-name." "one adapter via --lora-name."
) )
assert ( assert bench_args.lora_zipf_alpha > 1, (
bench_args.lora_zipf_alpha > 1 f"--lora-zipf-alpha must be > 1, got {bench_args.lora_zipf_alpha}"
), f"--lora-zipf-alpha must be > 1, got {bench_args.lora_zipf_alpha}" )
if bench_args.apply_chat_template and not ( if bench_args.apply_chat_template and not (
bench_args.fixed_prompt_file or bench_args.dataset_name in REPLAY_TEXT_DATASETS bench_args.fixed_prompt_file or bench_args.dataset_name in REPLAY_TEXT_DATASETS
python/sglang/benchmark/serving.py
+26 -26
View File
@@ -261,9 +261,9 @@ async def async_request_openai_completions(
pbar: Optional[tqdm] = None, pbar: Optional[tqdm] = None,
) -> RequestFuncOutput: ) -> RequestFuncOutput:
api_url = request_func_input.api_url api_url = request_func_input.api_url
assert api_url.endswith( assert api_url.endswith("completions"), (
"completions" "OpenAI Completions API URL must end with 'completions'."
), "OpenAI Completions API URL must end with 'completions'." )
prompt = request_func_input.prompt prompt = request_func_input.prompt
@@ -392,9 +392,9 @@ async def async_request_openai_chat_completions(
latency, TTFT, ITL, and success status. latency, TTFT, ITL, and success status.
""" """
api_url = request_func_input.api_url api_url = request_func_input.api_url
assert api_url.endswith( assert api_url.endswith("chat/completions"), (
"chat/completions" "OpenAI Chat Completions API URL must end with 'chat/completions'."
), "OpenAI Chat Completions API URL must end with 'chat/completions'." )
# TODO put it to other functions when `pbar` logic is refactored # TODO put it to other functions when `pbar` logic is refactored
if getattr(args, "print_requests", False): 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: def wrap_multi_turn_request_func(request_func: Callable, backend: str) -> Callable:
assert ( assert backend in MULTI_TURN_BACKENDS, (
backend in MULTI_TURN_BACKENDS f"Multi-turn only supports chat backends: {MULTI_TURN_BACKENDS}, got {backend}"
), f"Multi-turn only supports chat backends: {MULTI_TURN_BACKENDS}, got {backend}" )
async def f( async def f(
request_func_input: RequestFuncInput, request_func_input: RequestFuncInput,
@@ -1534,9 +1534,9 @@ async def benchmark(
lora_name = lora_names[lora_idx] lora_name = lora_names[lora_idx]
lora_idx = (lora_idx + 1) % len(lora_names) lora_idx = (lora_idx + 1) % len(lora_names)
else: else:
assert ( assert lora_request_distribution == "skewed", (
lora_request_distribution == "skewed" f"Unexpected lora_request_distribution: {lora_request_distribution}. Expected 'skewed'."
), f"Unexpected lora_request_distribution: {lora_request_distribution}. Expected 'skewed'." )
lora_name = np.random.choice(lora_names, p=lora_probs) lora_name = np.random.choice(lora_names, p=lora_probs)
else: else:
@@ -2000,9 +2000,9 @@ def run_benchmark(args_: argparse.Namespace):
extra_request_body["bootstrap_room"] = 0 extra_request_body["bootstrap_room"] = 0
if args.tokenize_prompt: if args.tokenize_prompt:
assert ( assert args.backend == "sglang", (
args.backend == "sglang" "`--tokenize-prompt` only compatible with `--backend sglang` currently"
), "`--tokenize-prompt` only compatible with `--backend sglang` currently" )
# Set url # Set url
if args.port is None: if args.port is None:
@@ -2079,18 +2079,18 @@ def run_benchmark(args_: argparse.Namespace):
if args.dataset_name in ["image", "mmmu"]: if args.dataset_name in ["image", "mmmu"]:
args.apply_chat_template = True args.apply_chat_template = True
assert ( assert not args.tokenize_prompt, (
not args.tokenize_prompt "`--tokenize-prompt` not compatible with image dataset"
), "`--tokenize-prompt` not compatible with image dataset" )
if args.lora_request_distribution in ["distinct", "skewed"]: if args.lora_request_distribution in ["distinct", "skewed"]:
assert ( assert args.lora_name is not None and len(args.lora_name) > 1, (
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."
), "More than 1 LoRA adapter must be specified via --lora-name to use 'distinct' or 'skewed' request distribution." )
assert ( assert args.lora_zipf_alpha > 1, (
args.lora_zipf_alpha > 1 f"Got invalid value for --lora-zipf-alpha of {args.lora_zipf_alpha}. It must be greater than 1."
), f"Got invalid value for --lora-zipf-alpha of {args.lora_zipf_alpha}. It must be greater than 1." )
print(f"{args}\n") print(f"{args}\n")
@@ -2364,13 +2364,13 @@ def cli_main():
"--image-format", "--image-format",
type=str, type=str,
default="jpeg", 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( parser.add_argument(
"--image-content", "--image-content",
type=str, type=str,
default="random", 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( parser.add_argument(
"--request-rate", "--request-rate",
python/sglang/cli/killall.py
+1 -2
View File
@@ -315,8 +315,7 @@ def _print_diagnostics(unkillable_pids):
print(f" {line}") print(f" {line}")
else: else:
print( print(
"\n[killall] Diagnostic — no sglang/python/gpu processes " "\n[killall] Diagnostic — no sglang/python/gpu processes in this container"
"in this container"
) )
python/sglang/cli/serve.py
+1 -1
View File
@@ -195,7 +195,7 @@ def serve(args, extra_argv):
else: else:
registered = registry.get(backend_name) registered = registry.get(backend_name)
logger.info( logger.info(
"Dispatch override enabled: --model-type=%s " "(skip auto detection)", "Dispatch override enabled: --model-type=%s (skip auto detection)",
backend_name, backend_name,
) )
python/sglang/cli/utils.py
+1 -2
View File
@@ -126,8 +126,7 @@ def get_model_path(extra_argv):
) )
else: else:
raise Exception( raise Exception(
"Error: --model-path is required. " "Error: --model-path is required. Please provide the path to the model."
"Please provide the path to the model."
) )
return model_path return model_path
python/sglang/kernels/aot/benchmark/bench_amd_deterministic_allreduce.py
+11 -11
View File
@@ -245,7 +245,7 @@ def worker(world_size, rank, port, results_queue):
if input_size_bytes > custom_ar.max_size: if input_size_bytes > custom_ar.max_size:
if rank == 0: if rank == 0:
print( 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 deterministic_kernel_available = False
else: else:
@@ -412,17 +412,17 @@ def worker(world_size, rank, port, results_queue):
} }
print( 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( 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: if custom_ar is not None and lat_custom_ar_median is not None:
overhead_custom = ( overhead_custom = (
(lat_custom_ar_median - lat_ar_median) / lat_ar_median (lat_custom_ar_median - lat_ar_median) / lat_ar_median
) * 100 ) * 100
print( 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: if lat_deterministic_kernel_median is not None:
overhead_kernel = ( overhead_kernel = (
@@ -433,7 +433,7 @@ def worker(world_size, rank, port, results_queue):
/ lat_rs_ag_median / lat_rs_ag_median
) * 100 ) * 100
print( 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: if lat_optimized_rs_ag_median is not None:
overhead_opt = ( 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 (lat_rs_ag_median - lat_optimized_rs_ag_median) / lat_rs_ag_median
) * 100 ) * 100
print( 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}%") print(f" RS+AG Overhead: {overhead_rs_ag:+.1f}%")
@@ -515,8 +515,8 @@ def main():
ar_det_str = "✓" if r["all_reduce"]["deterministic"] else "✗" ar_det_str = "✓" if r["all_reduce"]["deterministic"] else "✗"
rs_ag_det_str = "✓" if r["rs_ag"]["deterministic"] else "✗" rs_ag_det_str = "✓" if r["rs_ag"]["deterministic"] else "✗"
line = ( line = (
f"{bs:<8} {r['all_reduce']['latency_median']*1000:<12.3f} {ar_det_str:<8} " 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['rs_ag']['latency_median'] * 1000:<15.3f} {rs_ag_det_str:<10} "
f"{r['overhead_rs_ag_pct']:<12.1f}" f"{r['overhead_rs_ag_pct']:<12.1f}"
) )
if r.get("custom_ar") is not None: if r.get("custom_ar") is not None:
@@ -526,7 +526,7 @@ def main():
(custom_ar["latency_median"] - r["all_reduce"]["latency_median"]) (custom_ar["latency_median"] - r["all_reduce"]["latency_median"])
/ r["all_reduce"]["latency_median"] / r["all_reduce"]["latency_median"]
) * 100 ) * 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: if r.get("deterministic_kernel") is not None:
det_kernel = r["deterministic_kernel"] det_kernel = r["deterministic_kernel"]
det_kernel_det_str = "✓" if det_kernel["deterministic"] else "✗" 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"] - det_kernel["latency_median"])
/ r["rs_ag"]["latency_median"] / r["rs_ag"]["latency_median"]
) * 100 ) * 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: if r.get("optimized_rs_ag") is not None:
opt_rs_ag = r["optimized_rs_ag"] opt_rs_ag = r["optimized_rs_ag"]
opt_rs_ag_det_str = "✓" if opt_rs_ag["deterministic"] else "✗" 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"] - opt_rs_ag["latency_median"])
/ r["rs_ag"]["latency_median"] / r["rs_ag"]["latency_median"]
) * 100 ) * 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(line)
print("=" * 80) print("=" * 80)
python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py
+2 -4
View File
@@ -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() q_size = qn.numel() * qn.element_size() + qr.numel() * qr.element_size()
gbps = ( gbps = lambda ms: (
lambda ms: ( (q_size + q_size * dv / d + kv_cache.numel() * kv_cache.element_size())
q_size + q_size * dv / d + kv_cache.numel() * kv_cache.element_size()
)
* 1e-9 * 1e-9
/ (ms * 1e-3) / (ms * 1e-3)
) )
python/sglang/kernels/aot/benchmark/bench_fp4_gemm.py
+6 -6
View File
@@ -368,9 +368,9 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
res_fi, res_fi,
backend="cudnn", backend="cudnn",
) )
assert torch.allclose( assert torch.allclose(res_fi, res_cutlass, atol=1e-3, rtol=1e-3), (
res_fi, res_cutlass, atol=1e-3, rtol=1e-3 "cudnn fp4 doesn't match cutlass fp4"
), "cudnn fp4 doesn't match cutlass fp4" )
mm_fp4( mm_fp4(
a_fp4, a_fp4,
b_fp4_T, b_fp4_T,
@@ -381,9 +381,9 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
res_fi, res_fi,
backend="trtllm", backend="trtllm",
) )
assert torch.allclose( assert torch.allclose(res_fi, res_cutlass, atol=1e-3, rtol=1e-3), (
res_fi, res_cutlass, atol=1e-3, rtol=1e-3 "trtllm fp4 doesn't match cutlass fp4"
), "trtllm fp4 doesn't match cutlass fp4" )
if csv_file: if csv_file:
with open(csv_file, "a", newline="") as f: with open(csv_file, "a", newline="") as f:
python/sglang/kernels/aot/benchmark/bench_int8_gemm.py
+2 -2
View File
@@ -127,8 +127,8 @@ def benchmark(batch_size, provider, N, K):
lambda: vllm_scaled_mm(a, b, scale_a, scale_b, torch.float16, bias), lambda: vllm_scaled_mm(a, b, scale_a, scale_b, torch.float16, bias),
quantiles=quantiles, quantiles=quantiles,
) )
gbps = ( gbps = lambda ms: (
lambda ms: ( (
(2 * M * N * K - M * N) * a.element_size() (2 * M * N * K - M * N) * a.element_size()
+ (3 * M * N) * scale_a.element_size() + (3 * M * N) * scale_a.element_size()
) )
python/sglang/kernels/aot/python/sgl_kernel/attention.py
+9 -9
View File
@@ -38,9 +38,9 @@ def cutlass_mla_decode(
) -> torch.Tensor: ) -> torch.Tensor:
assert q_nope.ndim == 3, f"q_nope must be a 3D tensor, but got {q_nope.ndim}" 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 q_pe.ndim == 3, f"q_pe must be a 3D tensor, but got {q_pe.ndim}"
assert ( assert kv_c_and_k_pe_cache.ndim == 3, (
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}"
), 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, H, D_q_nope = q_nope.shape
B_q_2, H_2, D_q_pe = q_pe.shape B_q_2, H_2, D_q_pe = q_pe.shape
@@ -77,12 +77,12 @@ def cutlass_mla_decode(
torch.bfloat16, torch.bfloat16,
), f"q_nope.dtype needs to be fp16 or bf16 but got {q_nope.dtype}." ), 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 q_nope.dtype == q_pe.dtype == kv_c_and_k_pe_cache.dtype
assert ( assert seq_lens.dtype == torch.int32, (
seq_lens.dtype == torch.int32 f"seq_lens.dtype needs to be int32 but got {seq_lens.dtype}."
), f"seq_lens.dtype needs to be int32 but got {seq_lens.dtype}." )
assert ( assert page_table.dtype == torch.int32, (
page_table.dtype == torch.int32 f"page_table.dtype needs to be int32 but got {page_table.dtype}."
), f"page_table.dtype needs to be int32 but got {page_table.dtype}." )
out = q_nope.new_empty((B_q, MAX_HEADS, D_latent)) out = q_nope.new_empty((B_q, MAX_HEADS, D_latent))
python/sglang/kernels/aot/python/sgl_kernel/elementwise.py
+6 -6
View File
@@ -247,12 +247,12 @@ def gemma_fused_add_rmsnorm(
def _check_shape(input: torch.Tensor, output: torch.Tensor) -> None: def _check_shape(input: torch.Tensor, output: torch.Tensor) -> None:
assert input.ndim == output.ndim, f"{input.ndim} != {output.ndim}" assert input.ndim == output.ndim, f"{input.ndim} != {output.ndim}"
assert ( assert input.shape[:-1] == output.shape[:-1], (
input.shape[:-1] == output.shape[:-1] f"{input.shape[:-1]} != {output.shape[:-1]}"
), f"{input.shape[:-1]} != {output.shape[:-1]}" )
assert ( assert input.shape[-1] == 2 * output.shape[-1], (
input.shape[-1] == 2 * output.shape[-1] f"{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: def silu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
python/sglang/kernels/aot/python/sgl_kernel/flash_mla.py
+6 -6
View File
@@ -159,9 +159,9 @@ def flash_mla_with_kvcache(
assert extra_topk_length is None assert extra_topk_length is None
if indices is not None: if indices is not None:
assert causal == False, "causal must be `false` if sparse attention is enabled." assert causal == False, "causal must be `false` if sparse attention is enabled."
assert (descale_q is None) == ( assert (descale_q is None) == (descale_k is None), (
descale_k is None "descale_q and descale_k should be both None or both not None"
), "descale_q and descale_k should be both None or both not None" )
if indices is None and q.element_size() == 1: if indices is None and q.element_size() == 1:
out, softmax_lse = torch.ops.sgl_kernel.fwd_kvcache_mla_fp8.default( 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.causal == causal, helper_msg
assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, 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.topk == topk, helper_msg
assert ( assert sched_meta.config.extra_page_block_size == extra_page_block_size, (
sched_meta.config.extra_page_block_size == extra_page_block_size helper_msg
), helper_msg )
assert sched_meta.config.extra_topk == extra_topk, helper_msg assert sched_meta.config.extra_topk == extra_topk, helper_msg
if topk is not None: if topk is not None:
python/sglang/kernels/aot/python/sgl_kernel/metal.py
+2 -2
View File
@@ -76,11 +76,11 @@ def rope_pool_fused(
if q_shape != (q_shape[0], num_qo_heads, head_dim): if q_shape != (q_shape[0], num_qo_heads, head_dim):
raise ValueError( 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): if k_shape != (q_shape[0], num_kv_heads, head_dim):
raise ValueError( 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: if v_shape != k_shape:
raise ValueError(f"v shape must match k shape, got {v.shape} vs {k.shape}") raise ValueError(f"v shape must match k shape, got {v.shape} vs {k.shape}")
python/sglang/kernels/aot/python/sgl_kernel/musa.py
+6 -6
View File
@@ -86,9 +86,9 @@ def musa_fused_gemv(
out_shape = x.shape[:-1] + ( out_shape = x.shape[:-1] + (
qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2, qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2,
) )
assert not ( assert not (use_swigelu and use_rms_norm), (
use_swigelu and use_rms_norm "gemv only fused one activation (swigelu or rms_norm)!"
), "gemv only fused one activation (swigelu or rms_norm)!" )
if use_rms_norm: if use_rms_norm:
if gamma is None: if gamma is None:
@@ -113,9 +113,9 @@ def musa_fused_gemv(
return output return output
# w4a16 gemv # w4a16 gemv
elif qweight_scales is not None: elif qweight_scales is not None:
assert ( assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
x.dtype == torch.bfloat16 or x.dtype == torch.float16 "W4A16 gemv only support bfloat16 or float16!"
), "W4A16 gemv only support bfloat16 or float16!" )
use_int4_w4a16 = True use_int4_w4a16 = True
out_shape = x.shape[:-1] + ( out_shape = x.shape[:-1] + (
qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2, qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2,
python/sglang/kernels/aot/python/sgl_kernel/scalar_type.py
+15 -15
View File
@@ -70,9 +70,9 @@ class ScalarType:
""" """
def _floating_point_max_int(self) -> int: def _floating_point_max_int(self) -> int:
assert ( assert self.mantissa <= 52 and self.exponent <= 11, (
self.mantissa <= 52 and self.exponent <= 11 f"Cannot represent max/min as a double for type {self.__str__()}"
), f"Cannot represent max/min as a double for type {self.__str__()}" )
max_mantissa = (1 << self.mantissa) - 1 max_mantissa = (1 << self.mantissa) - 1
if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN: if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN:
@@ -80,9 +80,9 @@ class ScalarType:
max_exponent = (1 << self.exponent) - 2 max_exponent = (1 << self.exponent) - 2
if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN or self.nan_repr == NanRepr.NONE: if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN or self.nan_repr == NanRepr.NONE:
assert ( assert self.exponent < 11, (
self.exponent < 11 f"Cannot represent max/min as a double for type {self.__str__()}"
), f"Cannot represent max/min as a double for type {self.__str__()}" )
max_exponent = max_exponent + 1 max_exponent = max_exponent + 1
# adjust the exponent to match that of a double # adjust the exponent to match that of a double
@@ -109,25 +109,25 @@ class ScalarType:
if self.is_floating_point(): if self.is_floating_point():
return self._floating_point_max() return self._floating_point_max()
else: else:
assert ( assert self.size_bits < 64 or self.size_bits == 64 and self.is_signed(), (
self.size_bits < 64 or self.size_bits == 64 and self.is_signed() "Cannot represent max as an int"
), "Cannot represent max as an int" )
return (1 << self.mantissa) - 1 return (1 << self.mantissa) - 1
def _raw_min(self) -> Union[int, float]: def _raw_min(self) -> Union[int, float]:
if self.is_floating_point(): if self.is_floating_point():
assert ( assert self.is_signed(), (
self.is_signed() "We currently assume all floating point types are signed"
), "We currently assume all floating point types are signed" )
sign_bit_double = 1 << 63 sign_bit_double = 1 << 63
max_raw = self._floating_point_max_int() max_raw = self._floating_point_max_int()
min_raw = max_raw | sign_bit_double min_raw = max_raw | sign_bit_double
return struct.unpack("!d", struct.pack("!Q", min_raw))[0] return struct.unpack("!d", struct.pack("!Q", min_raw))[0]
else: else:
assert ( assert not self.is_signed() or self.size_bits <= 64, (
not self.is_signed() or self.size_bits <= 64 "Cannot represent min as a int64_t"
), "Cannot represent min as a int64_t" )
if self.is_signed(): if self.is_signed():
return -(1 << (self.size_bits - 1)) return -(1 << (self.size_bits - 1))
python/sglang/kernels/aot/python/sgl_kernel/test_utils.py
+3 -3
View File
@@ -94,9 +94,9 @@ def _compute_imbalanced_split(
def assert_all_close_or_tiny_diff(a: torch.Tensor, b: torch.Tensor): def assert_all_close_or_tiny_diff(a: torch.Tensor, b: torch.Tensor):
assert (a.shape == b.shape) and ( assert (a.shape == b.shape) and (a.dtype == b.dtype), (
a.dtype == b.dtype f"{a.shape=} {b.shape=} {a.dtype=} {b.dtype=}"
), f"{a.shape=} {b.shape=} {a.dtype=} {b.dtype=}" )
numel = a.numel() numel = a.numel()
if a.dtype == torch.float8_e4m3fn: if a.dtype == torch.float8_e4m3fn:
python/sglang/kernels/aot/python/sgl_kernel/testing/rotary_embedding.py
+6 -6
View File
@@ -112,9 +112,9 @@ class RotaryEmbedding(torch.nn.Module):
fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
) -> Tuple[torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor]:
"""A PyTorch-native implementation of forward().""" """A PyTorch-native implementation of forward()."""
assert ( assert fused_set_kv_buffer_arg is None, (
fused_set_kv_buffer_arg is None "fused_set_kv_buffer_arg is not supported for native implementation"
), "fused_set_kv_buffer_arg is not supported for native implementation" )
if offsets is not None: if offsets is not None:
positions = positions + offsets positions = positions + offsets
@@ -182,9 +182,9 @@ class SglKernelRotaryEmbedding(RotaryEmbedding):
offsets: Optional[torch.Tensor] = None, offsets: Optional[torch.Tensor] = None,
fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
) -> Tuple[torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor]:
assert ( assert fused_set_kv_buffer_arg is None, (
fused_set_kv_buffer_arg is None "fused_set_kv_buffer_arg is not supported for sgl-kernel implementation"
), "fused_set_kv_buffer_arg is not supported for sgl-kernel implementation" )
if self.cos_sin_cache.dtype != query.dtype: if self.cos_sin_cache.dtype != query.dtype:
self.cos_sin_cache = self.cos_sin_cache.to(query.dtype) self.cos_sin_cache = self.cos_sin_cache.to(query.dtype)
torch.ops.sgl_kernel.rotary_embedding( torch.ops.sgl_kernel.rotary_embedding(
python/sglang/kernels/aot/python/sgl_kernel/top_k.py
+9 -9
View File
@@ -33,9 +33,9 @@ def fast_topk_v2(
Returns: Returns:
The topk indices tensor of shape (B, topk) The topk indices tensor of shape (B, topk)
""" """
assert ( assert topk == 2048, (
topk == 2048 "fast_topk_v2 is only optimized for deepseek v3.2 model, where topk=2048"
), "fast_topk_v2 is only optimized for deepseek v3.2 model, where topk=2048" )
assert score.dim() == 2 assert score.dim() == 2
topk_indices = score.new_empty((score.size(0), topk), dtype=torch.int32) topk_indices = score.new_empty((score.size(0), topk), dtype=torch.int32)
torch.ops.sgl_kernel.fast_topk(score, topk_indices, lengths, row_starts) torch.ops.sgl_kernel.fast_topk(score, topk_indices, lengths, row_starts)
@@ -68,9 +68,9 @@ def fast_topk_transform_fused(
Returns: Returns:
The topk indices tensor of shape (B, topk) The topk indices tensor of shape (B, topk)
""" """
assert ( assert topk == 2048, (
topk == 2048 "fast_topk_transform_fused is only optimized for deepseek v3.2 model, where topk=2048"
), "fast_topk_transform_fused is only optimized for deepseek v3.2 model, where topk=2048" )
assert score.dim() == 2 assert score.dim() == 2
src_page_table = page_table_size_1 src_page_table = page_table_size_1
dst_page_table = score.new_empty((score.shape[0], topk), dtype=torch.int32) dst_page_table = score.new_empty((score.shape[0], topk), dtype=torch.int32)
@@ -138,9 +138,9 @@ def fast_topk_transform_ragged_fused(
Returns: Returns:
The topk indices tensor of shape (B, topk) The topk indices tensor of shape (B, topk)
""" """
assert ( assert topk == 2048, (
topk == 2048 "fast_topk_transform_ragged_fused is only optimized for deepseek v3.2 model, where topk=2048"
), "fast_topk_transform_ragged_fused is only optimized for deepseek v3.2 model, where topk=2048" )
assert score.dim() == 2 assert score.dim() == 2
topk_indices_ragged = score.new_empty((score.shape[0], topk), dtype=torch.int32) topk_indices_ragged = score.new_empty((score.shape[0], topk), dtype=torch.int32)
torch.ops.sgl_kernel.fast_topk_transform_ragged_fused( torch.ops.sgl_kernel.fast_topk_transform_ragged_fused(
python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py
+9 -9
View File
@@ -116,15 +116,15 @@ def test_tree_speculative_sampling_target_only(
deterministic=True, deterministic=True,
) )
assert ( assert predicts.tolist() == expected_predicts, (
predicts.tolist() == expected_predicts f"Predicts mismatch for thresholds ({threshold_single}, {threshold_acc})"
), f"Predicts mismatch for thresholds ({threshold_single}, {threshold_acc})" )
assert ( assert accept_index.tolist() == expected_accept_index, (
accept_index.tolist() == expected_accept_index f"Accept index mismatch for thresholds ({threshold_single}, {threshold_acc})"
), f"Accept index mismatch for thresholds ({threshold_single}, {threshold_acc})" )
assert ( assert accept_token_num.tolist() == expected_accept_token_num, (
accept_token_num.tolist() == expected_accept_token_num f"Accept token num mismatch for thresholds ({threshold_single}, {threshold_acc})"
), f"Accept token num mismatch for thresholds ({threshold_single}, {threshold_acc})" )
if __name__ == "__main__": if __name__ == "__main__":
python/sglang/kernels/aot/tests/test_custom_allreduce.py
+3 -3
View File
@@ -92,9 +92,9 @@ def multi_process_parallel(
for i in range(world_size): for i in range(world_size):
procs[i].join() procs[i].join()
assert ( assert procs[i].exitcode == 0, (
procs[i].exitcode == 0 f"Process {i} failed with exit code {procs[i].exitcode}"
), f"Process {i} failed with exit code {procs[i].exitcode}" )
class TestCustomAllReduce(unittest.TestCase): class TestCustomAllReduce(unittest.TestCase):
python/sglang/kernels/aot/tests/test_flash_attn_sparse.py
+8 -6
View File
@@ -251,12 +251,14 @@ def test_sparse_attention(
ref_out, ref_lse = ref_attn(q, k, v) ref_out, ref_lse = ref_attn(q, k, v)
torch.testing.assert_close( (
out, ref_out, atol=2e-2, rtol=1e-2 torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=1e-2),
), f"{torch.max(torch.abs(out - ref_out))}" 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(lse, ref_lse, atol=2e-2, rtol=1e-2),
f"{torch.max(torch.abs(lse - ref_lse))}",
)
# sparse attention utils # sparse attention utils
python/sglang/kernels/aot/tests/test_flashmla.py
+1 -3
View File
@@ -198,9 +198,7 @@ def reference_torch_prefill(
kvs = torch.index_select( kvs = torch.index_select(
kvs, 0, indices.masked_fill(invalid_indices_mask, 0).flatten() kvs, 0, indices.masked_fill(invalid_indices_mask, 0).flatten()
).view( ).view(s_q, topk, 576) # [s_q, topk, d_qk]
s_q, topk, 576
) # [s_q, topk, d_qk]
attn_score = qs @ kvs.transpose(1, 2) # [s_q, h_q, topk] 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.masked_fill_(invalid_indices_mask.unsqueeze(1), float("-inf"))
attn_score *= sm_scale * math.log2(math.e) attn_score *= sm_scale * math.log2(math.e)
python/sglang/kernels/aot/tests/test_fused_qk_norm_rope.py
+3 -3
View File
@@ -76,9 +76,9 @@ def torch_ref_rms_norm_rope(
v_size = num_heads_v * head_dim v_size = num_heads_v * head_dim
# Verify dimensions match # Verify dimensions match
assert ( assert hidden_size == q_size + k_size + v_size, (
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}"
), 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 # Split the tensor into Q, K, V parts
q = qkv[:, :q_size] q = qkv[:, :q_size]
python/sglang/kernels/aot/tests/test_moe_topk_sigmoid.py
+24 -24
View File
@@ -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) 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 # Verify the top-k weights and indices match the torch native ones
assert torch.allclose( assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}"
), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" )
assert torch.allclose( assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), (
topk_indices_ref.int(), topk_indices, atol=0, rtol=0 f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}"
), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" )
@pytest.mark.parametrize( @pytest.mark.parametrize(
@@ -87,13 +87,13 @@ def test_topk_sigmoid_dtype_regression(num_tokens, num_experts, topk, dtype):
gating_output.float(), gating_output.float(),
) )
assert torch.allclose( assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
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}"
), f"Weights mismatch: SGLang old interface={topk_weights_ref} vs SGLang new interface={topk_weights}" )
assert torch.allclose( assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), (
topk_indices_ref.int(), topk_indices, atol=0, rtol=0 f"Indices mismatch: SGLang old interface={topk_indices_ref}, SGLang new interface={topk_indices}"
), f"Indices mismatch: SGLang old interface={topk_indices_ref}, SGLang new interface={topk_indices}" )
@pytest.mark.parametrize( @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) topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True)
assert torch.allclose( assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
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}"
), f"Weights mismatch: SGLang w/o fused renormalize={topk_weights_ref} vs SGLang w/ fused renormalize={topk_weights}" )
assert torch.allclose( assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), (
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}"
), f"Indices mismatch: SGLang w/o fused renormalize={topk_indices_ref}, SGLang w/ fused renormalize={topk_indices}" )
@pytest.mark.parametrize( @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) 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 # Verify the top-k weights and indices match the torch native ones
assert torch.allclose( assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}"
), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" )
assert torch.allclose( assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), (
topk_indices_ref.int(), topk_indices, atol=0, rtol=0 f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}"
), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" )
if __name__ == "__main__": if __name__ == "__main__":
python/sglang/kernels/aot/tests/test_moe_topk_softmax.py
+24 -24
View File
@@ -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) 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 # Verify the top-k weights and indices match the torch native ones
assert torch.allclose( assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}"
), f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}" )
assert compare_topk_values( assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), (
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}"
), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" )
@pytest.mark.parametrize( @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) 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 # Verify the top-k weights and indices match the torch native ones
assert torch.allclose( assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}"
), f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}" )
assert compare_topk_values( assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), (
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}"
), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" )
@pytest.mark.parametrize( @pytest.mark.parametrize(
@@ -122,13 +122,13 @@ def test_topk_softmax_dtype_regression(num_tokens, num_experts, topk, dtype):
gating_output.float(), gating_output.float(),
) )
assert torch.allclose( assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
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}"
), f"Weights mismatch: SGLang old interface={topk_indices_ref} vs SGLang new interface={topk_weights}" )
assert compare_topk_values( assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), (
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}"
), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" )
@pytest.mark.parametrize( @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) topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True)
assert torch.allclose( assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
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}"
), f"Weights mismatch: SGLang w/o fused renormalize={topk_indices_ref} vs SGLang w/ fused renormalize={topk_weights}" )
assert compare_topk_values( assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), (
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}"
), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" )
if __name__ == "__main__": if __name__ == "__main__":
python/sglang/kernels/jit/__main__.py
+3 -3
View File
@@ -72,9 +72,9 @@ def generate_clangd():
arch = make_jit_cuda_arch(int(major), int(minor)) arch = make_jit_cuda_arch(int(major), int(minor))
else: else:
arch = get_jit_cuda_arch() arch = get_jit_cuda_arch()
assert ( assert arch.major > 0, (
arch.major > 0 "Cannot detect CUDA architecture, please specify --cuda-target explicitly."
), "Cannot detect CUDA architecture, please specify --cuda-target explicitly." )
compile_flags = [ compile_flags = [
"-xcuda", "-xcuda",
python/sglang/kernels/jit/benchmark/marker.py
+10 -11
View File
@@ -253,9 +253,9 @@ class Benchmark(Generic[F]):
f"parametrize name {name!r} is not a parameter of " f"parametrize name {name!r} is not a parameter of "
f"{self._fn.__name__}; available: {list(self._fn_params)}" f"{self._fn.__name__}; available: {list(self._fn_params)}"
) )
assert ( assert name not in self._seen_args, (
name not in self._seen_args f"parametrize name {name!r} is already used"
), f"parametrize name {name!r} is already used" )
self._seen_args.add(name) self._seen_args.add(name)
self._configs.insert(0, (names, vals)) 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 if p.default is inspect.Parameter.empty and p.kind in kinds
} - (set(flat_names) | {self._line_arg}) } - (set(flat_names) | {self._line_arg})
assert not missing, ( assert not missing, (
f"parameters not parametrized for {self._fn.__name__}: " f"parameters not parametrized for {self._fn.__name__}: {sorted(missing)}"
f"{sorted(missing)}"
) )
results, bandwidths, should_log_bw = self._collect_results() 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] return [(v,) for v in vs]
out: List[Tuple[Any, ...]] = [] out: List[Tuple[Any, ...]] = []
for v in vs: for v in vs:
assert isinstance( assert isinstance(v, (tuple, list)), (
v, (tuple, list) f"parametrize: multi-name values must be tuples, got {v!r}"
), f"parametrize: multi-name values must be tuples, got {v!r}" )
t = tuple(v) t = tuple(v)
assert ( assert len(t) == arity, (
len(t) == arity f"parametrize: each value must have length {arity}, got {t!r}"
), f"parametrize: each value must have length {arity}, got {t!r}" )
out.append(t) out.append(t)
return out return out
python/sglang/kernels/jit/utils/compile/loader.py
+1 -1
View File
@@ -121,7 +121,7 @@ def load_jit(
# Also the benign case where a concurrent GC unlinked the leaf # Also the benign case where a concurrent GC unlinked the leaf
# between the lookup and the load. # between the lookup and the load.
logger.warning( 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, spec.module_name,
e, e,
) )
python/sglang/kernels/kda_kernels/causal_conv3d_cat_pad_jit.py
+1 -1
View File
@@ -25,7 +25,7 @@ def _jit_causal_conv3d_cat_pad_module(dtype: torch.dtype) -> Module:
cuda_wrappers=[ cuda_wrappers=[
( (
"causal_conv3d_cat_pad", "causal_conv3d_cat_pad",
"causal_conv3d_cat_pad::" f"CausalConv3dCatPadKernel<{args}>::run", f"causal_conv3d_cat_pad::CausalConv3dCatPadKernel<{args}>::run",
) )
], ],
) )
python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py
+3 -3
View File
@@ -353,9 +353,9 @@ class _Qwen3xNvfp4Sm120Kernel:
self.occupancy, self.occupancy,
) )
assert ( assert self.epi_stage > 0, (
self.epi_stage > 0 "epi_stage <= 0, not enough shared memory. This configuration will be skipped."
), "epi_stage <= 0, not enough shared memory. This configuration will be skipped." )
( (
self.a_smem_layout_staged, self.a_smem_layout_staged,
python/sglang/kernels/kda_kernels/residual_gate_add_jit.py
+2 -2
View File
@@ -34,11 +34,11 @@ def _jit_residual_gate_add_module(dtype: torch.dtype) -> Module:
cuda_wrappers=[ cuda_wrappers=[
( (
"residual_gate_add", "residual_gate_add",
"residual_gate_add::" f"ResidualGateAddKernel<{args}>::run", f"residual_gate_add::ResidualGateAddKernel<{args}>::run",
), ),
( (
"residual_gate_add_transposed", "residual_gate_add_transposed",
"residual_gate_add::" f"ResidualGateAddKernel<{args}>::run_transposed", f"residual_gate_add::ResidualGateAddKernel<{args}>::run_transposed",
), ),
], ],
) )
python/sglang/kernels/ops/activation/activation.py
+3 -3
View File
@@ -157,9 +157,9 @@ def run_unary_activation(
Unlike :func:`run_activation`, there is no gate/up split — ``input`` and Unlike :func:`run_activation`, there is no gate/up split — ``input`` and
``out`` share the same shape. ``out`` share the same shape.
""" """
assert ( assert op_name in SUPPORTED_UNARY_ACTIVATIONS, (
op_name in SUPPORTED_UNARY_ACTIVATIONS f"Unsupported unary activation: {op_name}"
), f"Unsupported unary activation: {op_name}" )
if out is None: if out is None:
out = torch.empty_like(input) out = torch.empty_like(input)
_run_unary_activation_inplace(op_name, input, out) _run_unary_activation_inplace(op_name, input, out)
python/sglang/kernels/ops/activation/softcap.py
+3 -3
View File
@@ -101,9 +101,9 @@ def softcap_inplace_logits(full_logits, final_logit_softcapping):
row_stride = ncols row_stride = ncols
else: else:
assert full_logits.ndim == 2, "non-contiguous softcap requires 2D tensor" assert full_logits.ndim == 2, "non-contiguous softcap requires 2D tensor"
assert ( assert full_logits.stride(1) == 1, (
full_logits.stride(1) == 1 "non-contiguous softcap requires contiguous columns"
), "non-contiguous softcap requires contiguous columns" )
nrows, ncols = full_logits.shape nrows, ncols = full_logits.shape
row_stride = full_logits.stride(0) row_stride = full_logits.stride(0)
python/sglang/kernels/ops/attention/cutedsl_fp8_paged_mqa_logits.py
+6 -6
View File
@@ -221,12 +221,12 @@ class FP8MQALogitsKernel:
self.block_kv = block_kv self.block_kv = block_kv
self.phys_block_kv = phys_block_kv self.phys_block_kv = phys_block_kv
self.num_blocks_per_mma = block_kv // phys_block_kv self.num_blocks_per_mma = block_kv // phys_block_kv
assert ( assert block_kv % phys_block_kv == 0, (
block_kv % phys_block_kv == 0 f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}"
), f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}" )
assert ( assert self.num_blocks_per_mma <= 4, (
self.num_blocks_per_mma <= 4 f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 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.remove_kv_wait_in_epilogue = remove_kv_wait_in_epilogue
self.early_tmem_copy = early_tmem_copy self.early_tmem_copy = early_tmem_copy
self.smem_subpartition_opt = smem_subpartition_opt self.smem_subpartition_opt = smem_subpartition_opt
python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py
+57 -57
View File
@@ -3080,9 +3080,9 @@ def gated_delta_rule_mtp_wide_vec(
assert K_val == 128 and V_val == 128 assert K_val == 128 and V_val == 128
assert initial_state_source.dtype == torch.bfloat16 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 tile_v in (32, 64, 128), f"tile_v must be 32/64/128, got {tile_v}"
assert ( assert V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0, (
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"
), f"tile_v={tile_v} incompatible with 8 groups × ILP=4 layout" )
if cache_ring: if cache_ring:
assert replayssm_rawv is not None and replayssm_rawk is not None 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. # Validate recovery_steps for fused recovery+decode mode.
assert ( assert 0 <= recovery_steps <= T_val, (
0 <= recovery_steps <= T_val f"recovery_steps must be in [0, T={T_val}], got {recovery_steps}"
), f"recovery_steps must be in [0, T={T_val}], got {recovery_steps}" )
if recovery_steps > 0: if recovery_steps > 0:
assert ( assert not cache_intermediate_states, (
not cache_intermediate_states "recovery_steps > 0 is incompatible with intermediate state caching"
), "recovery_steps > 0 is incompatible with intermediate state caching" )
assert not disable_state_update, ( assert not disable_state_update, (
"recovery_steps > 0 requires state writeback " "recovery_steps > 0 requires state writeback "
"(disable_state_update=False); the boundary writeback at i_t=K-1 " "(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. # accepted_steps[i] is the per-request phase boundary.
per_request_accepted_steps = accepted_steps is not None per_request_accepted_steps = accepted_steps is not None
if per_request_accepted_steps: if per_request_accepted_steps:
assert accepted_steps.shape == ( assert accepted_steps.shape == (B_val,), (
B_val, f"accepted_steps must have shape [B={B_val}], got {accepted_steps.shape}"
), f"accepted_steps must have shape [B={B_val}], got {accepted_steps.shape}" )
assert ( assert accepted_steps.dtype == torch.int32, (
accepted_steps.dtype == torch.int32 f"accepted_steps must be int32, got {accepted_steps.dtype}"
), f"accepted_steps must be int32, got {accepted_steps.dtype}" )
assert accepted_steps.device == q.device assert accepted_steps.device == q.device
# FLA-style per-token pool scatter (vLLM API compat). When the public # 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. # this entry point hit the same fail-fast errors.
per_token_pool_scatter = ssm_state_indices is not None per_token_pool_scatter = ssm_state_indices is not None
if per_token_pool_scatter: if per_token_pool_scatter:
assert ( assert intermediate_states_buffer is None, (
intermediate_states_buffer is None "ssm_state_indices and intermediate_states_buffer are mutually exclusive"
), "ssm_state_indices and intermediate_states_buffer are mutually exclusive" )
assert ( assert not disable_state_update, (
not disable_state_update "ssm_state_indices requires state writes; disable_state_update must be False"
), "ssm_state_indices requires state writes; disable_state_update must be False" )
assert ( assert recovery_steps == 0, (
recovery_steps == 0 "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)"
), "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)" )
assert T_val >= 2, ( assert T_val >= 2, (
f"ssm_state_indices requires T >= 2 (got T={T_val}); " f"ssm_state_indices requires T >= 2 (got T={T_val}); "
f"for T=1 use output_state_indices" 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"ssm_state_indices must have shape [B={B_val}, T={T_val}], "
f"got {tuple(ssm_state_indices.shape)}" f"got {tuple(ssm_state_indices.shape)}"
) )
assert ( assert ssm_state_indices.dtype == torch.int32, (
ssm_state_indices.dtype == torch.int32 f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}"
), f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}" )
assert ssm_state_indices.device == q.device assert ssm_state_indices.device == q.device
phase_b_unroll = _select_wide_vec_phase_b_unroll( 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 K_val == 128 and V_val == 128
assert initial_state_source.dtype == torch.bfloat16 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 tile_v in (32, 64, 128), f"tile_v must be 32/64/128, got {tile_v}"
assert ( assert V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0, (
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"
), f"tile_v={tile_v} incompatible with 8 groups × ILP=4 layout" )
if scale is None: if scale is None:
scale = 1.0 / math.sqrt(K_val) 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"intermediate_states_buffer dim 0 ({buffer_size}) must equal "
f"batch size B={B}; the buffer is batch-scoped, not pool-scoped" f"batch size B={B}; the buffer is batch-scoped, not pool-scoped"
) )
assert ( assert cache_steps >= T, (
cache_steps >= T f"intermediate_states_buffer dim 1 ({cache_steps}) must be >= T={T}"
), f"intermediate_states_buffer dim 1 ({cache_steps}) must be >= T={T}" )
assert intermediate_states_buffer.dtype == torch.bfloat16 assert intermediate_states_buffer.dtype == torch.bfloat16
intermediate_states = intermediate_states_buffer.reshape( intermediate_states = intermediate_states_buffer.reshape(
B * cache_steps * HV, V, K B * cache_steps * HV, V, K
@@ -3860,28 +3860,28 @@ def gated_delta_rule_mtp(
# results/2026-06-03/FLA_SCATTER_MODE_PLAN.md. # results/2026-06-03/FLA_SCATTER_MODE_PLAN.md.
per_token_pool_scatter = ssm_state_indices is not None per_token_pool_scatter = ssm_state_indices is not None
if per_token_pool_scatter: if per_token_pool_scatter:
assert ( assert intermediate_states_buffer is None, (
intermediate_states_buffer is None "ssm_state_indices and intermediate_states_buffer are mutually exclusive"
), "ssm_state_indices and intermediate_states_buffer are mutually exclusive" )
assert ( assert not disable_state_update, (
not disable_state_update "ssm_state_indices requires state writes; disable_state_update must be False"
), "ssm_state_indices requires state writes; disable_state_update must be False" )
assert ( assert recovery_steps == 0, (
recovery_steps == 0 "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)"
), "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)" )
assert ( assert T >= 2, (
T >= 2 f"ssm_state_indices requires T >= 2 (got T={T}); for T=1 use output_state_indices"
), f"ssm_state_indices requires T >= 2 (got T={T}); for T=1 use output_state_indices" )
assert ssm_state_indices.shape == (B, T), ( assert ssm_state_indices.shape == (B, T), (
f"ssm_state_indices must have shape [B={B}, T={T}], " f"ssm_state_indices must have shape [B={B}, T={T}], "
f"got {tuple(ssm_state_indices.shape)}" f"got {tuple(ssm_state_indices.shape)}"
) )
assert ( assert ssm_state_indices.dtype == torch.int32, (
ssm_state_indices.dtype == torch.int32 f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}"
), f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}" )
assert ( assert ssm_state_indices.device == q.device, (
ssm_state_indices.device == q.device f"ssm_state_indices device {ssm_state_indices.device} != q 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 # 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 # 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 K opt-in (see gated_delta_rule_mtp_wide_vec for full rationale).
per_request_accepted_steps = accepted_steps is not None per_request_accepted_steps = accepted_steps is not None
if per_request_accepted_steps: if per_request_accepted_steps:
assert accepted_steps.shape == ( assert accepted_steps.shape == (B,), (
B, f"accepted_steps must have shape [B={B}], got {accepted_steps.shape}"
), f"accepted_steps must have shape [B={B}], got {accepted_steps.shape}" )
assert ( assert accepted_steps.dtype == torch.int32, (
accepted_steps.dtype == torch.int32 f"accepted_steps must be int32, got {accepted_steps.dtype}"
), f"accepted_steps must be int32, got {accepted_steps.dtype}" )
assert accepted_steps.device == q.device assert accepted_steps.device == q.device
# Contiguous pool -> sentinel keys + slot dim marked dynamic (pool-size # Contiguous pool -> sentinel keys + slot dim marked dynamic (pool-size
python/sglang/kernels/ops/attention/cutedsl_kda.py
+6 -6
View File
@@ -1427,12 +1427,12 @@ def cutedsl_fused_sigmoid_gating_kda_update(
N = initial_state_indices.shape[0] N = initial_state_indices.shape[0]
assert K == TILE_K, f"Current CuTe DSL KDA kernel requires K={TILE_K}, got {K}" assert K == TILE_K, f"Current CuTe DSL KDA kernel requires K={TILE_K}, got {K}"
assert ( assert V % TILE_V_SMALL == 0, (
V % TILE_V_SMALL == 0 f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0, got V={V}"
), f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0, got V={V}" )
assert ( assert V % TILE_V == 0, (
V % TILE_V == 0 f"Current CuTe DSL KDA kernel requires V % {TILE_V} == 0, got V={V}"
), 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, ( assert (V // TILE_V_SMALL) % NUM_BLOCKS_PER_STATE_SMALL == 0, (
"Small-batch KDA kernel requires num_v_tiles_small divisible by " "Small-batch KDA kernel requires num_v_tiles_small divisible by "
f"{NUM_BLOCKS_PER_STATE_SMALL}, got V={V}" f"{NUM_BLOCKS_PER_STATE_SMALL}, got V={V}"
python/sglang/kernels/ops/attention/decode_attention.py
-1
View File
@@ -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 # 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). # Triton compiler crashes in the Coalesce pass (max_output_tile_cnt is runtime-computed).
while iter < cta_end_tile_gid: while iter < cta_end_tile_gid:
tile_row_idx = iter // tiles_per_khead tile_row_idx = iter // tiles_per_khead
tile_idx = tile_row_idx * batch_size tile_idx = tile_row_idx * batch_size
tile_iter = tile_row_idx * tiles_per_khead tile_iter = tile_row_idx * tiles_per_khead
python/sglang/kernels/ops/attention/deepseek_v4_rope.py
+9 -9
View File
@@ -354,9 +354,9 @@ def apply_rotary_emb_triton(
grid = (batch_size, n_heads if is_3d else 1, num_blocks_dim) grid = (batch_size, n_heads if is_3d else 1, num_blocks_dim)
if positions is not None: if positions is not None:
assert positions.shape == ( assert positions.shape == (batch_size,), (
batch_size, f"positions shape {positions.shape} != ({batch_size},)"
), f"positions shape {positions.shape} != ({batch_size},)" )
apply_rotary_emb_triton_kernel[grid]( apply_rotary_emb_triton_kernel[grid](
x, x,
@@ -374,9 +374,9 @@ def apply_rotary_emb_triton(
BLOCK_SIZE=BLOCK_SIZE, BLOCK_SIZE=BLOCK_SIZE,
) )
else: else:
assert ( assert freqs_real.shape[0] == batch_size, (
freqs_real.shape[0] == batch_size f"freqs_cis batch size {freqs_real.shape[0]} != x batch size {batch_size}"
), f"freqs_cis batch size {freqs_real.shape[0]} != x batch size {batch_size}" )
apply_rotary_emb_triton_kernel[grid]( apply_rotary_emb_triton_kernel[grid](
x, x,
@@ -621,9 +621,9 @@ def fused_norm_rope_inplace_triton(
if weight is not None: if weight is not None:
assert weight.shape == (head_dim,) assert weight.shape == (head_dim,)
if positions is None: if positions is None:
assert ( assert freqs_real.shape[0] == M, (
freqs_real.shape[0] == M f"freqs_cis row count {freqs_real.shape[0]} != M={M}"
), f"freqs_cis row count {freqs_real.shape[0]} != M={M}" )
else: else:
assert positions.shape == (M,) and positions.dim() == 1 assert positions.shape == (M,) and positions.dim() == 1
python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py
+3 -3
View File
@@ -181,9 +181,9 @@ def dequantize_k_cache_paged(
output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache
""" """
dim_quant = quant_k_cache.shape[-1] dim_quant = quant_k_cache.shape[-1]
assert ( assert dim_quant == 656, (
dim_quant == 656 f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged"
), f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged" )
quant_k_cache = quant_k_cache.view((-1, dim_quant)) 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) # num_tokens can exceed kv_cache_size due to prefix sharing (multiple seqs share same KV slots)
python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py
+3 -3
View File
@@ -308,9 +308,9 @@ def _set_k_and_s_triton(
assert scale_dim == 1 assert scale_dim == 1
if _is_hip: if _is_hip:
if _use_aiter_preshuffle: if _use_aiter_preshuffle:
assert ( assert page_size % 16 == 0, (
page_size % 16 == 0 f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}" )
else: else:
assert page_size == 64 assert page_size == 64
python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py
+33 -46
View File
@@ -155,9 +155,9 @@ def act_quant(
- A tensor of scaling factors with dtype `torch.float32`. - A tensor of scaling factors with dtype `torch.float32`.
""" """
assert x.is_contiguous(), "Input tensor must be contiguous" assert x.is_contiguous(), "Input tensor must be contiguous"
assert ( assert x.size(-1) % block_size == 0, (
x.size(-1) % block_size == 0 f"Last dimension size must be divisible by block_size (block_size={block_size})"
), f"Last dimension size must be divisible by block_size (block_size={block_size})" )
N = x.size(-1) N = x.size(-1)
if _is_fp8_fnuz: if _is_fp8_fnuz:
y = torch.empty_like(x, dtype=torch.float8_e4m3fnuz) y = torch.empty_like(x, dtype=torch.float8_e4m3fnuz)
@@ -272,16 +272,16 @@ def sparse_attention_fwd_kernel_v1(
num_stages=2, num_stages=2,
threads=256, threads=256,
): ):
assert dim == tilelang.math.next_power_of_2( assert dim == tilelang.math.next_power_of_2(dim), (
dim f"haven't check padding correctness yet, dim={dim}"
), f"haven't check padding correctness yet, dim={dim}" )
assert tail_dim == tilelang.math.next_power_of_2( assert tail_dim == tilelang.math.next_power_of_2(tail_dim), (
tail_dim f"haven't check padding correctness yet, dim={tail_dim}"
), f"haven't check padding correctness yet, dim={tail_dim}" )
assert is_causal == True, "non-casual is not supported" assert is_causal == True, "non-casual is not supported"
assert ( assert topk % block_I == 0, (
topk % block_I == 0 "otherwise will load some index=0 thus causing wrong kv to be loaded"
), "otherwise will load some index=0 thus causing wrong kv to be loaded" )
if sm_scale is None: if sm_scale is None:
sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e)
else: else:
@@ -361,7 +361,6 @@ def sparse_attention_fwd_kernel_v1(
T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) 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 i_i in T.Pipelined(NI, num_stages=num_stages):
for bi_i in T.Parallel(BI): for bi_i in T.Parallel(BI):
mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] >= 0 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, sm_scale: Optional[float] = None,
block_I: int = 64, block_I: int = 64,
): ):
assert dim == tilelang.math.next_power_of_2( assert dim == tilelang.math.next_power_of_2(dim), (
dim f"haven't check padding correctness yet, dim={dim}"
), f"haven't check padding correctness yet, dim={dim}" )
assert tail_dim == tilelang.math.next_power_of_2( assert tail_dim == tilelang.math.next_power_of_2(tail_dim), (
tail_dim f"haven't check padding correctness yet, dim={tail_dim}"
), f"haven't check padding correctness yet, dim={tail_dim}" )
assert ( assert topk % block_I == 0, (
topk % block_I == 0 "otherwise will load some index=0 thus causing wrong kv to be loaded"
), "otherwise will load some index=0 thus causing wrong kv to be loaded" )
if sm_scale is None: if sm_scale is None:
sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e)
else: else:
@@ -1078,9 +1077,9 @@ def sparse_mla_fwd_decode_partial_fp8(
threads=256, threads=256,
): ):
assert d_v == 512, f"only support d_v=512" assert d_v == 512, f"only support d_v=512"
assert ( assert topk % block_I == 0, (
topk % block_I == 0 "otherwise will load some index=0 thus causing wrong kv to be loaded"
), "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 # 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. # 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 h_per_block = 16
# Match bf16 partial behavior: keep fixed 16-head tiles and use # Match bf16 partial behavior: keep fixed 16-head tiles and use
# sliced T.copy on H0:H1 for tail handling. # sliced T.copy on H0:H1 for tail handling.
assert ( assert num_heads <= h_per_block or num_heads % h_per_block == 0, (
num_heads <= h_per_block or num_heads % h_per_block == 0 "num_heads must be <=16 or divisible by 16"
), "num_heads must be <=16 or divisible by 16" )
head_blocks_per_seq = (num_heads + h_per_block - 1) // h_per_block head_blocks_per_seq = (num_heads + h_per_block - 1) // h_per_block
batch = 1 batch = 1
@@ -1594,9 +1593,7 @@ def dpsk_v4_fp8_partial_kernel(
sm_scale = sm_scale * log2e sm_scale = sm_scale * log2e
assert dim == 448 and tail_dim == 64 assert dim == 448 and tail_dim == 64
assert topk_1 % block_I == 0 assert topk_1 % block_I == 0
assert ( assert (topk_1 // block_I) % inner_iter_1 == 0, (
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}" 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 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: if is_dual:
assert inner_iter_2 > 0, "dual-cache call requires inner_iter_2 > 0" assert inner_iter_2 > 0, "dual-cache call requires inner_iter_2 > 0"
assert topk_2 % block_I == 0 assert topk_2 % block_I == 0
assert ( assert (topk_2 // block_I) % inner_iter_2 == 0, (
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}" 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 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 @T.prim_func
def main( def main(
Partial_O: T.Tensor( Partial_O: T.Tensor([batch, seq_len, n_groups, num_heads, DT], BF16), # type: ignore
[batch, seq_len, n_groups, num_heads, DT], BF16 Partial_LSE: T.Tensor([batch, seq_len, n_groups, num_heads], accum_dtype), # type: ignore
), # 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_1: T.Tensor([batch], INT32), # type: ignore
Topk_length_2: T.Tensor([batch], INT32), # type: ignore Topk_length_2: T.Tensor([batch], INT32), # type: ignore
Attn_sink: T.Tensor([num_heads], FP32), # type: ignore Attn_sink: T.Tensor([num_heads], FP32), # type: ignore
@@ -2369,12 +2360,8 @@ def dpsk_v4_combine_kernel(
@T.prim_func @T.prim_func
def main( def main(
Partial_O: T.Tensor( Partial_O: T.Tensor([batch, seq_len, n_groups, num_heads, DT], BF16), # type: ignore
[batch, seq_len, n_groups, num_heads, DT], BF16 Partial_LSE: T.Tensor([batch, seq_len, n_groups, num_heads], accum_dtype), # type: ignore
), # 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 Attn_sink: T.Tensor([num_heads], FP32), # type: ignore
Output: T.Tensor([batch, seq_len, num_heads, DT], BF16), # 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 LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype), # type: ignore
python/sglang/kernels/ops/attention/dsa/triton_kernel.py
+3 -3
View File
@@ -99,9 +99,9 @@ def act_quant(
- A tensor of scaling factors with dtype `torch.float32`. - A tensor of scaling factors with dtype `torch.float32`.
""" """
assert x.is_contiguous(), "Input tensor must be contiguous" assert x.is_contiguous(), "Input tensor must be contiguous"
assert ( assert x.size(-1) % block_size == 0, (
x.size(-1) % block_size == 0 f"Last dimension size must be divisible by block_size (block_size={block_size})"
), f"Last dimension size must be divisible by block_size (block_size={block_size})" )
# Flatten all dims except last # Flatten all dims except last
N = x.size(-1) N = x.size(-1)
python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py
+2 -6
View File
@@ -69,9 +69,7 @@ def _sparse_mla_fwd_kernel(
) # [H, D_V] ) # [H, D_V]
q_tail = tl.load( q_tail = tl.load(
q_rope_ptr + s_i * H * D_TAIL + h[:, None] * D_TAIL + dt[None, :] q_rope_ptr + s_i * H * D_TAIL + h[:, None] * D_TAIL + dt[None, :]
).to( ).to(q_nope_ptr.dtype.element_ty) # [H, D_TAIL]
q_nope_ptr.dtype.element_ty
) # [H, D_TAIL]
m_i = tl.full([H], -float("inf"), tl.float32) m_i = tl.full([H], -float("inf"), tl.float32)
l_i = tl.zeros([H], 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 ) # [BLOCK_N, D_V] -- reused as V
kv_tail = tl.load( kv_tail = tl.load(
kbase + (D_V + dt)[None, :], mask=valid[:, None], other=0.0 kbase + (D_V + dt)[None, :], mask=valid[:, None], other=0.0
).to( ).to(q_nope_ptr.dtype.element_ty) # [BLOCK_N, D_TAIL]
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_main, tl.trans(kv_main)).to(tl.float32)
qk += tl.dot(q_tail, tl.trans(kv_tail)).to(tl.float32) qk += tl.dot(q_tail, tl.trans(kv_tail)).to(tl.float32)
python/sglang/kernels/ops/attention/dsv4/index_buf_accessor.py
+3 -1
View File
@@ -209,7 +209,9 @@ def _set_k_and_s_torch(
== num_tokens_to_write_nope == num_tokens_to_write_nope
== num_tokens_to_write_rope == num_tokens_to_write_rope
== num_tokens_to_write_scale == 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 buf.dtype == torch.uint8
assert loc.dtype in [ assert loc.dtype in [
python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py
+9 -9
View File
@@ -110,9 +110,9 @@ def _init_compressed_attn_metadata_triton(
# no cache-write locations. Keep the write buffers unpadded and mask those # no cache-write locations. Keep the write buffers unpadded and mask those
# rows in the kernel. # rows in the kernel.
num_write_tokens = raw_out_loc.shape[0] num_write_tokens = raw_out_loc.shape[0]
assert ( assert num_write_tokens <= bs, (
num_write_tokens <= bs f"raw_out_loc has {num_write_tokens} rows, expected at most {bs} metadata rows"
), f"raw_out_loc has {num_write_tokens} rows, expected at most {bs} metadata rows" )
device = seq_lens.device device = seq_lens.device
c4_out_loc = torch.empty(num_write_tokens, dtype=torch.int64, device=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) c128_seq_lens_clamp1 = torch.empty(bs, dtype=torch.int32, device=device)
if compute_page_indices: if compute_page_indices:
assert ( assert page_table is not None, (
page_table is not None "page_table required when compute_page_indices=True"
), "page_table required when compute_page_indices=True" )
assert ( assert page_size >= 128 and page_size % 128 == 0, (
page_size >= 128 and page_size % 128 == 0 "page_size must be a multiple of 128 when compute_page_indices=True"
), "page_size must be a multiple of 128 when compute_page_indices=True" )
max_pages = page_table.shape[1] max_pages = page_table.shape[1]
c128_page_size = page_size // 128 c128_page_size = page_size // 128
c128_cur_max_seq_len = c128_page_size * max_pages c128_cur_max_seq_len = c128_page_size * max_pages

Some files were not shown because too many files have changed in this diff Show More