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

Co-authored-by: Alison Shao <a.shao@wustl.edu>

.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py

@@ -541,7 +541,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
             "fused_flashmla_metadata",
         ),
         rationale_hint=(
-            "NSA replay metadata copies are already fused into one-kernel" " families."
+            "NSA replay metadata copies are already fused into one-kernel families."
         ),
         min_share=0.02,
         likely_share=0.2,
@@ -787,7 +787,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
             ("softmax", "sampling"),
         ),
         rationale_hint=(
-            "Decode-time sampling already has fused temperature and softmax" " kernels."
+            "Decode-time sampling already has fused temperature and softmax kernels."
         ),
         min_share=0.05,
         likely_share=0.5,
@@ -1218,8 +1218,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
     FusionPatternSpec(
         pattern="vLLM fused residual add + RMSNorm",
         candidate_path=(
-            "vllm/_custom_ops.py"
+            "vllm/_custom_ops.py<br>vllm/compilation/passes/fusion/rms_quant_fusion.py"
-            "<br>vllm/compilation/passes/fusion/rms_quant_fusion.py"
         ),
         active_keywords=(
             "fused_add_rms_norm",
@@ -1236,8 +1235,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
     FusionPatternSpec(
         pattern="vLLM fused activation-and-mul",
         candidate_path=(
-            "vllm/_custom_ops.py"
+            "vllm/_custom_ops.py<br>vllm/compilation/passes/fusion/act_quant_fusion.py"
-            "<br>vllm/compilation/passes/fusion/act_quant_fusion.py"
         ),
         active_keywords=(
             "silu_and_mul",

.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_proof_generator.py

@@ -256,7 +256,9 @@ def _module_assign_names(text: str) -> set:
         targets = (
             node.targets
             if isinstance(node, ast.Assign)
-            else [node.target] if isinstance(node, ast.AnnAssign) else []
+            else [node.target]
+            if isinstance(node, ast.AnnAssign)
+            else []
         )
         names |= {t.id for t in targets if isinstance(t, ast.Name)}
     return names

.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_reproduction_utils.py

@@ -173,9 +173,9 @@ def _find_unique_def(
         if isinstance(node, definition) and node.name == name
     ]
     assert matches, f"{name} not found in {where}"
-    assert (
+    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]
 
 
@@ -287,9 +287,9 @@ def _lowered_call_text(text: str, node: ast.Call) -> str:
     """
     receiver = node.args[0]
     receiver_src = _node_slice(text, receiver)
-    assert (
+    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(
         text,
         node.func.end_lineno,
@@ -722,9 +722,9 @@ class Repro:
                     )
                 existing = [alias_text(a.name, a.asname) for a in node.names]
                 added = alias_text(name, asname)
-                assert (
+                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
                 lines[node.lineno - 1 : node.end_lineno] = [rebuilt]
                 _write_source(path, "".join(lines))
@@ -825,9 +825,9 @@ class Repro:
                 for node in tree.body
                 if isinstance(node, (ast.Import, ast.ImportFrom))
             ]
-            assert (
+            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
             lines[insert_at:insert_at] = [
                 nl,
@@ -864,9 +864,9 @@ class Repro:
                     replaced = lines[node.lineno - 1].replace(
                         f"from {spelled} import", f"from {new_module} import", 1
                     )
-                    assert (
+                    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
                     changed = True
             assert changed, f"nested import of {name} from {old_module} not in {rel}"
@@ -974,9 +974,9 @@ class Repro:
         ``self: Target`` annotation is dropped (redundant inside the class). The body is moved
         verbatim; the formatter normalises the surrounding blank lines.
         """
-        assert (
+        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:
             src_path = root / src
@@ -1255,15 +1255,17 @@ class Repro:
                 targets = (
                     node.targets
                     if isinstance(node, ast.Assign)
-                    else [node.target] if isinstance(node, ast.AnnAssign) else []
+                    else [node.target]
+                    if isinstance(node, ast.AnnAssign)
+                    else []
                 )
                 names = {t.id for t in targets if isinstance(t, ast.Name)}
                 hit = names & dropped
                 if not hit:
                     continue
-                assert len(names) == len(
+                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
                 for dropped_name in hit:
                     removed_assigns[dropped_name] = value_src
@@ -1289,15 +1291,17 @@ class Repro:
                 else:
                     assign_spans.append((node.lineno, node.end_lineno))
                 found_assigns |= hit
-            assert (
+            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] = {}
             for node in tree.body:
                 targets = (
                     node.targets
                     if isinstance(node, ast.Assign)
-                    else [node.target] if isinstance(node, ast.AnnAssign) else []
+                    else [node.target]
+                    if isinstance(node, ast.AnnAssign)
+                    else []
                 )
                 names = [t.id for t in targets if isinstance(t, ast.Name)]
                 if not names or set(names) & dropped:
@@ -1383,9 +1387,9 @@ class Repro:
             src_text = _read_source(src_path)
             assert src_text.count(body) == 1, f"block not found uniquely in {src}"
             at = src_text.find(body)
-            assert (
+            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))
 
             dst_path = root / dst

.claude/skills/mechanical-refactor-verify/scripts/tests/proof_generator/test_infer_moves.py

@@ -370,12 +370,7 @@ def test_infer_recipe_module_level_def_shadowed_by_method_name(repo: Path) -> No
                 "        return foo(x=self.x)\n"
             ),
             "util.py": (
-                "def keep():\n"
+                "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(
         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",
         },
     )
@@ -448,9 +441,7 @@ def test_infer_recipe_move_leaving_a_forwarding_delegate(repo: Path) -> None:
         repo,
         **{
             "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": (
                 "class C:\n"

.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_cli/cli_testlib.py

@@ -2,13 +2,9 @@ import subprocess
 from pathlib import Path
 
 _PASSING_PROOF = (
-    "import sys\n"
+    '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:

.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_add_imports.py

@@ -110,13 +110,7 @@ def test_add_typechecking_import_inserts_in_block(tmp_path: Path) -> None:
 def test_add_typechecking_import_creates_missing_block(tmp_path: Path) -> None:
     """With no TYPE_CHECKING block, one is created after the trailing module import."""
     (tmp_path / "m.py").write_text(
-        "from typing import TYPE_CHECKING\n"
+        "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")
     _apply(r, tmp_path)
@@ -215,12 +209,7 @@ def test_add_typechecking_import_raises_without_imports(tmp_path: Path) -> None:
 def test_add_typechecking_import_drops_a_lone_pass_placeholder(tmp_path: Path) -> None:
     """Populating a `pass`-only TYPE_CHECKING block replaces the placeholder."""
     (tmp_path / "m.py").write_text(
-        "from typing import TYPE_CHECKING\n"
+        "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")
     _apply(r, tmp_path)

.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_extract_symbols_to_new_module.py

@@ -111,11 +111,7 @@ def test_extract_symbols_to_new_module_drops_relocated_assigns(tmp_path: Path) -
         "    return _FLAG\n"
     )
     header = (
-        "from __future__ import annotations\n"
+        "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(
         "src.py",

.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_assign.py

@@ -19,13 +19,7 @@ def test_move_assign_relocates_a_module_constant(tmp_path: Path) -> None:
     _apply(r, tmp_path)
     assert "LIMIT" not in (tmp_path / "src.py").read_text().split("def stay")[0]
     assert (tmp_path / "dst.py").read_text() == (
-        "import sys\n"
+        "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)
     assert "LIMIT" not in (tmp_path / "src.py").read_text().split("def stay")[0]
     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

@@ -260,13 +260,7 @@ def test_move_symbol_dedent_leaves_string_literal_interior_lines(
     )
     _apply(r, tmp_path)
     assert (tmp_path / "dst.py").read_text() == (
-        "import os\n"
+        "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

@@ -101,8 +101,7 @@ def format_summary_line(filename: str, result: Dict[str, Any]) -> str:
     if result.get("ok"):
         return f"{filename}: ok"
     return (
-        f"{filename}: failed status={result.get('status')} "
+        f"{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 = (
         max(timestamps) - min(timestamps)
         if len(timestamps) >= 2
-        else 0.0 if len(timestamps) == 1 else None
+        else 0.0
+        if len(timestamps) == 1
+        else None
     )
 
     lines = [

.pre-commit-config.yaml

@@ -48,10 +48,7 @@ repos:
           python/sglang/srt/grpc/.*_pb2\.pyi$|
           python/sglang/srt/grpc/.*_pb2_grpc\.pyi$|
           )$
-  - repo: https://github.com/psf/black
+      - 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$'
   - repo: https://github.com/codespell-project/codespell
     rev: v2.4.1

3rdparty/amd/tuning/benchmark_moe_rocm.py

@@ -187,8 +187,10 @@ def run_grid(bs, model, method, tp_size, dtype: str):
 
     configs = union_of_list_of_dicts(prune_configs_1, prune_configs_2)
 
-    print(f"{bs=} || {len(full_configs)=} | {len(prune_configs_1)=} | \
+    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_time_us = 1e20

benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py

@@ -243,9 +243,7 @@ def run(task, fi, tri, device, dtype, args):
                     )  # noqa: E731
                 else:
                     inp = make_decode_inputs(B, H, HV, K, V, pool, device, dtype)
-                    corr = lambda kern: call_decode(
+                    corr = lambda kern: call_decode(kern, inp, inp["ssm"].clone())  # noqa: E731
-                        kern, inp, inp["ssm"].clone()
-                    )  # noqa: E731
                     ssm_t = inp["ssm"].clone()
                     timed = lambda kern: call_decode(kern, inp, ssm_t)  # noqa: E731
 

benchmark/gsm8k/bench_sglang.py

@@ -53,9 +53,9 @@ def main(args):
     if args.enable_thinking:
         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(
             args.tokenizer_path, trust_remote_code=True
         )

benchmark/hicache/bench_serving.py

@@ -75,9 +75,9 @@ async def async_request_openai_completions(
     pbar: Optional[tqdm] = None,
 ) -> RequestFuncOutput:
     api_url = request_func_input.api_url
-    assert api_url.endswith(
+    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:
         payload = {

benchmark/hicache/nextqa.py

@@ -120,7 +120,7 @@ class NExTQALoader(VideoLoader):
         video = Video(video_path, num_frames)
         prompt = entry["question"] + "?"
         if self.task == "MC":  # add choices
-            prompt += f' a0: {entry["a0"]}, a1: {entry["a1"]}, a2: {entry["a2"]}, a3: {entry["a3"]}'
+            prompt += f" a0: {entry['a0']}, a1: {entry['a1']}, a2: {entry['a2']}, a3: {entry['a3']}"
         return VideoPrompt(video_path, num_frames, prompt)
 
     def __iter__(self):

benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py

@@ -149,9 +149,9 @@ def _check_correctness():
             cos = torch.nn.functional.cosine_similarity(
                 (q.float() * scale).flatten(), ref_deq.flatten(), dim=0
             ).item()
-            assert (
+            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)")
 
 

benchmark/kernels/deepep/deepep_utils.py

@@ -191,9 +191,9 @@ def bench_kineto(
     kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names
     assert all([isinstance(name, str) for name in kernel_names])
     for name in kernel_names:
-        assert (
+        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
     if trace_path is not None:

benchmark/kernels/deepep/tuning_deepep.py

@@ -155,7 +155,7 @@ def test_main(
                 for with_topk in (False, True):
                     if local_rank == 0:
                         print(
-                            f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
+                            f"[testing] Running with {'FP8' if isinstance(current_x, tuple) else 'BF16'}, {'with' if with_topk else 'without'} top-k (async={async_mode}, previous={previous_mode}) ...",
                             flush=True,
                             end="",
                         )
@@ -198,9 +198,9 @@ def test_main(
 
                     # Checks
                     recv_gbl_rank_prefix_sum = handle[-4]
-                    assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
+                    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 (
                         gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
                         == recv_num_tokens_per_expert_list
@@ -325,11 +325,14 @@ def test_main(
                 tune_args = {"x": current_x, "handle": handle, "config": config}
                 t = bench(lambda: buffer.dispatch(**tune_args))[0]
                 if t < best_time:
-                    best_time, best_results = t, (
+                    best_time, best_results = (
+                        t,
+                        (
                             num_sms,
                             nvl_chunk_size,
                             rdma_chunk_size,
                             config_kwargs,
+                        ),
                     )
                 if local_rank == 0:
                     print(
@@ -338,7 +341,7 @@ def test_main(
                     )
         if local_rank == 0:
             print(
-                f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
+                f"[tuning] Best dispatch ({'FP8' if isinstance(current_x, tuple) else 'BF16'}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
                 flush=True,
             )
             print("", flush=True)
@@ -399,11 +402,14 @@ def test_main(
                     flush=True,
                 )
                 if t < best_time:
-                    best_time, best_results = t, (
+                    best_time, best_results = (
+                        t,
+                        (
                             num_sms,
                             nvl_chunk_size,
                             rdma_chunk_size,
                             config_kwargs,
+                        ),
                     )
 
     if local_rank == 0:

benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py

@@ -59,7 +59,6 @@ def tl_gemm(
             bx,
             by,
         ):
-
             A_shared = T.alloc_shared(A_shared_shape, in_dtype)
             B_shared = T.alloc_shared(B_shared_shape, in_dtype)
             C_shared = T.alloc_shared(C_shared_shape, out_dtype)

benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py

@@ -243,8 +243,7 @@ def main():
             else:
                 speedup = f"{legacy_us / us:.2f}x"
             print(
-                f"{s_q:>5} {kv_len:>7} {name:>10} {us:>9.1f} "
+                f"{s_q:>5} {kv_len:>7} {name:>10} {us:>9.1f} {tbps:>9.3f} {speedup:>8}"
-                f"{tbps:>9.3f} {speedup:>8}"
             )
         print()
 

benchmark/kernels/elementwise/benchmark_concat_mla.py

@@ -143,9 +143,7 @@ output_exp = execute_and_get_output(fn_cuda, data)
 if not torch.all(output_ref == output_exp):
     abs_delta = torch.abs(output_ref - output_exp)
     raise AssertionError(
-        f"{output_ref=} {output_exp=} "
+        f"{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

@@ -535,7 +535,6 @@ class BestConfigTrace:
 
 
 class BenchmarkWorker:
-
     def __init__(self, seed: int, server_args: ServerArgs) -> None:
         torch.set_default_device("cuda")
         torch.cuda.manual_seed_all(0)
@@ -729,8 +728,7 @@ class BenchmarkWorker:
                 down_use_tma_map[block_m] = time_cost_all[2] > time_cost_all[3]
 
             print(
-                f"Round 1 done. Down TMA decisions per BLOCK_SIZE_M: "
+                f"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 ===

benchmark/lean_kernel_sweep.py

@@ -137,7 +137,9 @@ def main():
                 # b32 x 128K on the 8-KV-head config exceeds the microbench's single
                 # contiguous KV tensor (faults the GPU); real serving uses a paged pool.
                 if H_KV == 8 and B == 32 and S == 131072:
-                    print(f"{B:>5} {S//1024:>5}K {'skipped (contiguous-KV limit)':>30}")
+                    print(
+                        f"{B:>5} {S // 1024:>5}K {'skipped (contiguous-KV limit)':>30}"
+                    )
                     rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,skip,")
                     continue
                 try:

benchmark/ocr/generate_report.py

@@ -164,7 +164,7 @@ def _render_sample(sample: dict, failures_only: bool) -> str:
 
     parts = [f'<div class="sample">']
     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"<span>📄 {html.escape(pdf)}  &nbsp;·&nbsp; page {page}</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>'
                     )
             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"):
                 before = ti.get("before", "")
                 after = ti.get("after", "")

benchmark/prefill_only/util.py

@@ -745,7 +745,6 @@ async def run_generic_benchmark(
     async with aiohttp.ClientSession(
         timeout=aiohttp.ClientTimeout(total=300)
     ) as session:
-
         # Send START_PROFILE if profiling is enabled
         if config.profile:
             await send_profile_request("START_PROFILE", http_url, session=session)

benchmark/scheduler/bench_token_storage.py

@@ -254,7 +254,7 @@ def print_breakdown(title: str, results: dict[str, dict[str, float]]) -> None:
 
 
 def microbench_torch_tensor_paths(
-    sizes: tuple[int, ...] = (1_000, 10_000, 100_000)
+    sizes: tuple[int, ...] = (1_000, 10_000, 100_000),
 ) -> None:
     """Compare three CPU-buffer -> pinned cuda tensor paths.
 
@@ -293,9 +293,11 @@ def microbench_torch_tensor_paths(
         ),
         (
             "(C) from_numpy(frombuf(array('q'))).pin() -> cuda",
-            lambda x: torch.from_numpy(np.frombuffer(x, dtype=np.int64))
+            lambda x: (
+                torch.from_numpy(np.frombuffer(x, dtype=np.int64))
                 .pin_memory()
-            .to("cuda", non_blocking=True),
+                .to("cuda", non_blocking=True)
+            ),
         ),
     ]:
         cells = []

examples/frontend_language/usage/llava_video/srt_example_llava_v.py

@@ -124,7 +124,6 @@ def batch(video_dir, save_dir, cur_chunk, num_chunks, num_frames=16, batch_size=
 
 
 if __name__ == "__main__":
-
     url = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4"
 
     cache_dir = os.path.expanduser("~/.cache")

examples/profiler/nsys_profile_tools/gputrc2graph.py

@@ -182,9 +182,9 @@ class GPUTrace2Graph:
 
     def is_valid_file(self, base_file):
         """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):
         """figure out if new file should be generated from base_file"""

examples/runtime/engine/offline_batch_inference_qwen_1m.py

@@ -17,8 +17,7 @@ def load_prompt() -> str:
     # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/1m.txt
 
     with urlopen(
-        "https://qianwen-res.oss-cn-beijing.aliyuncs.com"
+        "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/64k.txt",
-        "/Qwen2.5-1M/test-data/64k.txt",
         timeout=5,
     ) as response:
         prompt = response.read().decode("utf-8")
@@ -41,9 +40,7 @@ def process_requests(llm: sgl.Engine, prompts: list[str]) -> None:
     for output in outputs:
         prompt_token_ids = output["meta_info"]["prompt_tokens"]
         generated_text = output["text"]
-        print(
+        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.

experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py

@@ -194,9 +194,9 @@ def _direct_warm(worker_url: str, model_id: str, prefix: str) -> None:
         },
         timeout=60.0,
     )
-    assert (
+    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:
@@ -212,9 +212,9 @@ def _route_through(router_url: str, model_id: str, prompt: str) -> str:
     after = _success_counts_by_worker(router_url)
     deltas = {w: after.get(w, 0) - before.get(w, 0) for w in set(after) | set(before)}
     winners = [w for w, d in deltas.items() if d > 0]
-    assert (
+    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]
 
 
@@ -309,15 +309,15 @@ def test_routers_route_by_prefix_content(
                         landed = _route_through(
                             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(
                             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:
                     _dump_logs(logs)
                     raise

experimental/sgl-router/tests/e2e/chat_completions/test_validation.py

@@ -50,9 +50,9 @@ def test_chat_non_streaming_returns_assistant_message(
                 body = resp.json()
                 choice = body["choices"][0]
                 assert choice["message"]["role"] == "assistant"
-                assert choice["message"][
+                assert choice["message"]["content"], (
-                    "content"
+                    f"empty assistant content: {choice!r}"
-                ], f"empty assistant content: {choice!r}"
+                )
                 assert choice.get("finish_reason"), choice
     finally:
         gpu_allocator.release(gpu)
@@ -91,8 +91,8 @@ def test_chat_streaming_emits_sse_chunks_with_done(
                         if line.startswith("data:"):
                             chunks.append(line.strip())
                 assert len(chunks) >= 2, f"expected >=2 SSE chunks, got: {chunks}"
-                assert any(
+                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:
         gpu_allocator.release(gpu)

experimental/sgl-router/tests/e2e/k8s_integration/test_discovery.py

@@ -66,7 +66,8 @@ def test_router_discovers_multiple_workers(router_url):
     # Scale down to 1 — router should still route after reconverging
     _scale_fake_worker(1)
     _poll_until(
-        lambda: httpx.post(
+        lambda: (
+            httpx.post(
                 f"{router_url}/v1/chat/completions",
                 json={
                     "model": "tiny",
@@ -74,7 +75,8 @@ def test_router_discovers_multiple_workers(router_url):
                 },
                 timeout=10.0,
             ).status_code
-        == 200,
+            == 200
+        ),
         "router routes after scale-down to 1",
         timeout=60,
         interval=3,

experimental/sgl-router/tests/e2e/test_chat_smoke.py

@@ -16,9 +16,9 @@ def test_models(router: str) -> None:
     assert resp.status_code == 200, resp.text
     data = resp.json()
     ids = [m["id"] for m in data.get("data", [])]
-    assert any(
+    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:
@@ -59,6 +59,6 @@ def test_chat_streaming(router: str) -> None:
                 chunks.append(line)
 
     assert len(chunks) >= 2, f"Expected >=2 SSE chunks, got {len(chunks)}: {chunks}"
-    assert any(
+    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

@@ -20,9 +20,9 @@ def test_tokenize_round_trip(router: str) -> None:
     )
     assert tok_resp.status_code == 200, tok_resp.text
     tokens = tok_resp.json()["tokens"]
-    assert (
+    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
     detok_resp = httpx.post(
@@ -32,6 +32,6 @@ def test_tokenize_round_trip(router: str) -> None:
     )
     assert detok_resp.status_code == 200, detok_resp.text
     recovered = detok_resp.json()["text"]
-    assert (
+    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

@@ -85,7 +85,7 @@ def load_tokenizer_with_fallback(primary, fallback, slug):
                 raise
             continue
     raise RuntimeError(
-        f"No accessible tokenizer for slug={slug} " f"(tried: {primary}, {fallback})"
+        f"No accessible tokenizer for slug={slug} (tried: {primary}, {fallback})"
     )
 
 

python/sglang/benchmark/dspark_sps_profiler.py

@@ -687,8 +687,7 @@ def run_one_round(
     rank_rows = fetch_rank_rows(base_url=context.base_url)
     if len(rank_rows) != len(watermarks):
         raise RuntimeError(
-            f"DP rank count changed mid-profile: {len(watermarks)} -> "
+            f"DP rank count changed mid-profile: {len(watermarks)} -> {len(rank_rows)}."
-            f"{len(rank_rows)}."
         )
     new_rank_rows = [
         [row for row in rows if row.forward_ct > watermark]

python/sglang/benchmark/offline_throughput.py

@@ -136,13 +136,13 @@ class BenchArgs:
             "--gsp-system-prompt-len",
             type=int,
             default=BenchArgs.gsp_system_prompt_len,
-            help="System prompt length, used" "only for generate-shared-prefix",
+            help="System prompt length, usedonly for generate-shared-prefix",
         )
         parser.add_argument(
             "--gsp-question-len",
             type=int,
             default=BenchArgs.gsp_question_len,
-            help="Question length, used" "only for generate-shared-prefix",
+            help="Question length, usedonly for generate-shared-prefix",
         )
         parser.add_argument(
             "--gsp-output-len",
@@ -259,9 +259,9 @@ def throughput_test_once(
     ]
 
     if profile:
-        assert (
+        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)
         known_files = None
         backend.start_profile(

python/sglang/benchmark/one_batch_server.py

@@ -1241,9 +1241,9 @@ def run_benchmark_internal(
             skip_max_running_requests_threshold = float("inf")
             skip_token_capacity_threshold = float("inf")
         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
 
         print(f"{max_running_requests_per_dp=}")
@@ -1288,9 +1288,9 @@ def run_benchmark_internal(
             "--lora-request-distribution=distinct/skewed requires more than "
             "one adapter via --lora-name."
         )
-    assert (
+    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 (
         bench_args.fixed_prompt_file or bench_args.dataset_name in REPLAY_TEXT_DATASETS

python/sglang/benchmark/serving.py

@@ -261,9 +261,9 @@ async def async_request_openai_completions(
     pbar: Optional[tqdm] = None,
 ) -> RequestFuncOutput:
     api_url = request_func_input.api_url
-    assert api_url.endswith(
+    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
 
@@ -392,9 +392,9 @@ async def async_request_openai_chat_completions(
                            latency, TTFT, ITL, and success status.
     """
     api_url = request_func_input.api_url
-    assert api_url.endswith(
+    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
     if getattr(args, "print_requests", False):
@@ -1296,9 +1296,9 @@ def _normalize_round_messages(turn: Any) -> Optional[List[Dict[str, str]]]:
 
 
 def wrap_multi_turn_request_func(request_func: Callable, backend: str) -> Callable:
-    assert (
+    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(
         request_func_input: RequestFuncInput,
@@ -1534,9 +1534,9 @@ async def benchmark(
                 lora_name = lora_names[lora_idx]
                 lora_idx = (lora_idx + 1) % len(lora_names)
             else:
-                assert (
+                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)
         else:
@@ -2000,9 +2000,9 @@ def run_benchmark(args_: argparse.Namespace):
         extra_request_body["bootstrap_room"] = 0
 
     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
     if args.port is None:
@@ -2079,18 +2079,18 @@ def run_benchmark(args_: argparse.Namespace):
 
     if args.dataset_name in ["image", "mmmu"]:
         args.apply_chat_template = True
-        assert (
+        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"]:
-        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")
 
@@ -2364,13 +2364,13 @@ def cli_main():
         "--image-format",
         type=str,
         default="jpeg",
-        help=("Format of images for image dataset. " "Supports jpeg and png."),
+        help=("Format of images for image dataset. Supports jpeg and png."),
     )
     parser.add_argument(
         "--image-content",
         type=str,
         default="random",
-        help=("Content for images for image dataset. " "Supports random and blank."),
+        help=("Content for images for image dataset. Supports random and blank."),
     )
     parser.add_argument(
         "--request-rate",

python/sglang/cli/killall.py

@@ -315,8 +315,7 @@ def _print_diagnostics(unkillable_pids):
             print(f"  {line}")
     else:
         print(
-            "\n[killall] Diagnostic — no sglang/python/gpu processes "
+            "\n[killall] Diagnostic — no sglang/python/gpu processes in this container"
-            "in this container"
         )
 
 

python/sglang/cli/serve.py

@@ -195,7 +195,7 @@ def serve(args, extra_argv):
         else:
             registered = registry.get(backend_name)
             logger.info(
-                "Dispatch override enabled: --model-type=%s " "(skip auto detection)",
+                "Dispatch override enabled: --model-type=%s (skip auto detection)",
                 backend_name,
             )
 

python/sglang/cli/utils.py

@@ -126,8 +126,7 @@ def get_model_path(extra_argv):
             )
         else:
             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
 

python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py

@@ -114,10 +114,8 @@ def benchmark(batch_size, seq_len, provider, block_size, num_kv_splits):
 
     q_size = qn.numel() * qn.element_size() + qr.numel() * qr.element_size()
 
-    gbps = (
+    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
         / (ms * 1e-3)
     )

python/sglang/kernels/aot/benchmark/bench_fp4_gemm.py

@@ -368,9 +368,9 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
             res_fi,
             backend="cudnn",
         )
-        assert torch.allclose(
+        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(
             a_fp4,
             b_fp4_T,
@@ -381,9 +381,9 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
             res_fi,
             backend="trtllm",
         )
-        assert torch.allclose(
+        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:
         with open(csv_file, "a", newline="") as f:

python/sglang/kernels/aot/benchmark/bench_int8_gemm.py

@@ -127,8 +127,8 @@ def benchmark(batch_size, provider, N, K):
             lambda: vllm_scaled_mm(a, b, scale_a, scale_b, torch.float16, bias),
             quantiles=quantiles,
         )
-    gbps = (
+    gbps = lambda ms: (
-        lambda ms: (
+        (
             (2 * M * N * K - M * N) * a.element_size()
             + (3 * M * N) * scale_a.element_size()
         )

python/sglang/kernels/aot/python/sgl_kernel/attention.py

@@ -38,9 +38,9 @@ def cutlass_mla_decode(
 ) -> torch.Tensor:
     assert q_nope.ndim == 3, f"q_nope must be a 3D tensor, but got {q_nope.ndim}"
     assert q_pe.ndim == 3, f"q_pe must be a 3D tensor, but got {q_pe.ndim}"
-    assert (
+    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_2, H_2, D_q_pe = q_pe.shape
@@ -77,12 +77,12 @@ def cutlass_mla_decode(
         torch.bfloat16,
     ), f"q_nope.dtype needs to be fp16 or bf16 but got {q_nope.dtype}."
     assert q_nope.dtype == q_pe.dtype == kv_c_and_k_pe_cache.dtype
-    assert (
+    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))
 

python/sglang/kernels/aot/python/sgl_kernel/elementwise.py

@@ -247,12 +247,12 @@ def gemma_fused_add_rmsnorm(
 
 def _check_shape(input: torch.Tensor, output: torch.Tensor) -> None:
     assert input.ndim == output.ndim, f"{input.ndim} != {output.ndim}"
-    assert (
+    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:

python/sglang/kernels/aot/python/sgl_kernel/flash_mla.py

@@ -159,9 +159,9 @@ def flash_mla_with_kvcache(
     assert extra_topk_length is None
     if indices is not None:
         assert causal == False, "causal must be `false` if sparse attention is enabled."
-    assert (descale_q is None) == (
+    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:
         out, softmax_lse = torch.ops.sgl_kernel.fwd_kvcache_mla_fp8.default(
@@ -257,9 +257,9 @@ def _flash_mla_with_kvcache_sched_meta(
         assert sched_meta.config.causal == causal, helper_msg
         assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, helper_msg
         assert sched_meta.config.topk == topk, helper_msg
-        assert (
+        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
 
     if topk is not None:

python/sglang/kernels/aot/python/sgl_kernel/metal.py

@@ -76,11 +76,11 @@ def rope_pool_fused(
 
     if q_shape != (q_shape[0], num_qo_heads, head_dim):
         raise ValueError(
-            "q shape must be [num_tokens, num_qo_heads, head_dim], " f"got {q.shape}"
+            f"q shape must be [num_tokens, num_qo_heads, head_dim], got {q.shape}"
         )
     if k_shape != (q_shape[0], num_kv_heads, head_dim):
         raise ValueError(
-            "k shape must be [num_tokens, num_kv_heads, head_dim], " f"got {k.shape}"
+            f"k shape must be [num_tokens, num_kv_heads, head_dim], got {k.shape}"
         )
     if v_shape != k_shape:
         raise ValueError(f"v shape must match k shape, got {v.shape} vs {k.shape}")

python/sglang/kernels/aot/python/sgl_kernel/musa.py

@@ -86,9 +86,9 @@ def musa_fused_gemv(
     out_shape = x.shape[:-1] + (
         qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2,
     )
-    assert not (
+    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 gamma is None:
@@ -113,9 +113,9 @@ def musa_fused_gemv(
         return output
     # w4a16 gemv
     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
         out_shape = x.shape[:-1] + (
             qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2,

python/sglang/kernels/aot/python/sgl_kernel/scalar_type.py

@@ -70,9 +70,9 @@ class ScalarType:
     """
 
     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
         if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN:
@@ -80,9 +80,9 @@ class ScalarType:
 
         max_exponent = (1 << self.exponent) - 2
         if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN or self.nan_repr == NanRepr.NONE:
-            assert (
+            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
 
         # adjust the exponent to match that of a double
@@ -109,25 +109,25 @@ class ScalarType:
         if self.is_floating_point():
             return self._floating_point_max()
         else:
-            assert (
+            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
 
     def _raw_min(self) -> Union[int, float]:
         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
 
             max_raw = self._floating_point_max_int()
             min_raw = max_raw | sign_bit_double
             return struct.unpack("!d", struct.pack("!Q", min_raw))[0]
         else:
-            assert (
+            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():
                 return -(1 << (self.size_bits - 1))

python/sglang/kernels/aot/python/sgl_kernel/test_utils.py

@@ -94,9 +94,9 @@ def _compute_imbalanced_split(
 
 
 def assert_all_close_or_tiny_diff(a: torch.Tensor, b: torch.Tensor):
-    assert (a.shape == b.shape) and (
+    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()
 
     if a.dtype == torch.float8_e4m3fn:

python/sglang/kernels/aot/python/sgl_kernel/testing/rotary_embedding.py

@@ -112,9 +112,9 @@ class RotaryEmbedding(torch.nn.Module):
         fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """A PyTorch-native implementation of forward()."""
-        assert (
+        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:
             positions = positions + offsets
@@ -182,9 +182,9 @@ class SglKernelRotaryEmbedding(RotaryEmbedding):
         offsets: Optional[torch.Tensor] = None,
         fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        assert (
+        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:
             self.cos_sin_cache = self.cos_sin_cache.to(query.dtype)
         torch.ops.sgl_kernel.rotary_embedding(

python/sglang/kernels/aot/python/sgl_kernel/top_k.py

@@ -33,9 +33,9 @@ def fast_topk_v2(
     Returns:
         The topk indices tensor of shape (B, topk)
     """
-    assert (
+    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
     topk_indices = score.new_empty((score.size(0), topk), dtype=torch.int32)
     torch.ops.sgl_kernel.fast_topk(score, topk_indices, lengths, row_starts)
@@ -68,9 +68,9 @@ def fast_topk_transform_fused(
     Returns:
         The topk indices tensor of shape (B, topk)
     """
-    assert (
+    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
     src_page_table = page_table_size_1
     dst_page_table = score.new_empty((score.shape[0], topk), dtype=torch.int32)
@@ -138,9 +138,9 @@ def fast_topk_transform_ragged_fused(
     Returns:
         The topk indices tensor of shape (B, topk)
     """
-    assert (
+    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
     topk_indices_ragged = score.new_empty((score.shape[0], topk), dtype=torch.int32)
     torch.ops.sgl_kernel.fast_topk_transform_ragged_fused(

python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py

@@ -116,15 +116,15 @@ def test_tree_speculative_sampling_target_only(
         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__":

python/sglang/kernels/aot/tests/test_custom_allreduce.py

@@ -92,9 +92,9 @@ def multi_process_parallel(
 
     for i in range(world_size):
         procs[i].join()
-        assert (
+        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):

python/sglang/kernels/aot/tests/test_flash_attn_sparse.py

@@ -251,12 +251,14 @@ def test_sparse_attention(
 
     ref_out, ref_lse = ref_attn(q, k, v)
 
-    torch.testing.assert_close(
+    (
-        out, ref_out, atol=2e-2, rtol=1e-2
+        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

python/sglang/kernels/aot/tests/test_flashmla.py

@@ -198,9 +198,7 @@ def reference_torch_prefill(
 
     kvs = torch.index_select(
         kvs, 0, indices.masked_fill(invalid_indices_mask, 0).flatten()
-    ).view(
+    ).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.masked_fill_(invalid_indices_mask.unsqueeze(1), float("-inf"))
     attn_score *= sm_scale * math.log2(math.e)

python/sglang/kernels/aot/tests/test_fused_qk_norm_rope.py

@@ -76,9 +76,9 @@ def torch_ref_rms_norm_rope(
     v_size = num_heads_v * head_dim
 
     # Verify dimensions match
-    assert (
+    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
     q = qkv[:, :q_size]

python/sglang/kernels/aot/tests/test_moe_topk_sigmoid.py

@@ -44,13 +44,13 @@ def test_topk_sigmoid(num_tokens, num_experts, topk):
     topk_weights_ref, topk_indices_ref = torch.topk(sigmoid_output, topk, dim=-1)
 
     # Verify the top-k weights and indices match the torch native ones
-    assert torch.allclose(
+    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(
@@ -87,13 +87,13 @@ def test_topk_sigmoid_dtype_regression(num_tokens, num_experts, topk, dtype):
         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(
@@ -136,13 +136,13 @@ def test_topk_sigmoid_renormalize(num_tokens, num_experts, topk):
     )
     topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True)
 
-    assert torch.allclose(
+    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(
@@ -180,13 +180,13 @@ def test_topk_sigmoid_renormalize_correction_bias(num_tokens, num_experts, topk)
     topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True)
 
     # Verify the top-k weights and indices match the torch native ones
-    assert torch.allclose(
+    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__":

python/sglang/kernels/aot/tests/test_moe_topk_softmax.py

@@ -41,13 +41,13 @@ def test_topkfast_softmax(num_tokens, num_experts, topk):
     topk_weights_ref, topk_indices_ref = torch.topk(softmax_output, topk, dim=-1)
 
     # Verify the top-k weights and indices match the torch native ones
-    assert torch.allclose(
+    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(
@@ -79,13 +79,13 @@ def test_topk_softmax(num_tokens, num_experts, topk):
     topk_weights_ref, topk_indices_ref = torch.topk(softmax_output, topk, dim=-1)
 
     # Verify the top-k weights and indices match the torch native ones
-    assert torch.allclose(
+    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(
@@ -122,13 +122,13 @@ def test_topk_softmax_dtype_regression(num_tokens, num_experts, topk, dtype):
         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(
@@ -171,13 +171,13 @@ def test_topk_softmax_renormalize(num_tokens, num_experts, topk):
     )
     topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True)
 
-    assert torch.allclose(
+    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__":

python/sglang/kernels/jit/__main__.py

@@ -72,9 +72,9 @@ def generate_clangd():
         arch = make_jit_cuda_arch(int(major), int(minor))
     else:
         arch = get_jit_cuda_arch()
-        assert (
+        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 = [
         "-xcuda",

python/sglang/kernels/jit/benchmark/marker.py

@@ -253,9 +253,9 @@ class Benchmark(Generic[F]):
                 f"parametrize name {name!r} is not a parameter of "
                 f"{self._fn.__name__}; available: {list(self._fn_params)}"
             )
-            assert (
+            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._configs.insert(0, (names, vals))
 
@@ -305,8 +305,7 @@ class Benchmark(Generic[F]):
             if p.default is inspect.Parameter.empty and p.kind in kinds
         } - (set(flat_names) | {self._line_arg})
         assert not missing, (
-            f"parameters not parametrized for {self._fn.__name__}: "
+            f"parameters not parametrized for {self._fn.__name__}: {sorted(missing)}"
-            f"{sorted(missing)}"
         )
 
         results, bandwidths, should_log_bw = self._collect_results()
@@ -360,13 +359,13 @@ def parametrize(names: str, vals: List[Any], ci_vals: Optional[List[Any]] = None
             return [(v,) for v in vs]
         out: List[Tuple[Any, ...]] = []
         for v in vs:
-            assert isinstance(
+            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)
-            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)
         return out
 

python/sglang/kernels/jit/utils/compile/loader.py

@@ -121,7 +121,7 @@ def load_jit(
             # Also the benign case where a concurrent GC unlinked the leaf
             # between the lookup and the load.
             logger.warning(
-                "Cached JIT module %s failed to load; rebuilding. " "Got error: %s",
+                "Cached JIT module %s failed to load; rebuilding. Got error: %s",
                 spec.module_name,
                 e,
             )

python/sglang/kernels/kda_kernels/causal_conv3d_cat_pad_jit.py

@@ -25,7 +25,7 @@ def _jit_causal_conv3d_cat_pad_module(dtype: torch.dtype) -> Module:
         cuda_wrappers=[
             (
                 "causal_conv3d_cat_pad",
-                "causal_conv3d_cat_pad::" f"CausalConv3dCatPadKernel<{args}>::run",
+                f"causal_conv3d_cat_pad::CausalConv3dCatPadKernel<{args}>::run",
             )
         ],
     )

python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py

@@ -353,9 +353,9 @@ class _Qwen3xNvfp4Sm120Kernel:
             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,

python/sglang/kernels/kda_kernels/residual_gate_add_jit.py

@@ -34,11 +34,11 @@ def _jit_residual_gate_add_module(dtype: torch.dtype) -> Module:
         cuda_wrappers=[
             (
                 "residual_gate_add",
-                "residual_gate_add::" f"ResidualGateAddKernel<{args}>::run",
+                f"residual_gate_add::ResidualGateAddKernel<{args}>::run",
             ),
             (
                 "residual_gate_add_transposed",
-                "residual_gate_add::" f"ResidualGateAddKernel<{args}>::run_transposed",
+                f"residual_gate_add::ResidualGateAddKernel<{args}>::run_transposed",
             ),
         ],
     )

python/sglang/kernels/ops/activation/activation.py

@@ -157,9 +157,9 @@ def run_unary_activation(
     Unlike :func:`run_activation`, there is no gate/up split — ``input`` and
     ``out`` share the same shape.
     """
-    assert (
+    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:
         out = torch.empty_like(input)
     _run_unary_activation_inplace(op_name, input, out)

python/sglang/kernels/ops/activation/softcap.py

@@ -101,9 +101,9 @@ def softcap_inplace_logits(full_logits, final_logit_softcapping):
         row_stride = ncols
     else:
         assert full_logits.ndim == 2, "non-contiguous softcap requires 2D tensor"
-        assert (
+        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
         row_stride = full_logits.stride(0)
 

python/sglang/kernels/ops/attention/cutedsl_fp8_paged_mqa_logits.py

@@ -221,12 +221,12 @@ class FP8MQALogitsKernel:
         self.block_kv = block_kv
         self.phys_block_kv = phys_block_kv
         self.num_blocks_per_mma = block_kv // phys_block_kv
-        assert (
+        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.early_tmem_copy = early_tmem_copy
         self.smem_subpartition_opt = smem_subpartition_opt

python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py

@@ -3080,9 +3080,9 @@ def gated_delta_rule_mtp_wide_vec(
     assert K_val == 128 and V_val == 128
     assert initial_state_source.dtype == torch.bfloat16
     assert tile_v in (32, 64, 128), f"tile_v must be 32/64/128, got {tile_v}"
-    assert (
+    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:
         assert replayssm_rawv is not None and replayssm_rawk is not None
@@ -3194,13 +3194,13 @@ def gated_delta_rule_mtp_wide_vec(
     )
 
     # Validate recovery_steps for fused recovery+decode mode.
-    assert (
+    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:
-        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, (
             "recovery_steps > 0 requires state writeback "
             "(disable_state_update=False); the boundary writeback at i_t=K-1 "
@@ -3220,12 +3220,12 @@ def gated_delta_rule_mtp_wide_vec(
     # accepted_steps[i] is the per-request phase boundary.
     per_request_accepted_steps = accepted_steps is not None
     if per_request_accepted_steps:
-        assert accepted_steps.shape == (
+        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
 
     # FLA-style per-token pool scatter (vLLM API compat). When the public
@@ -3236,15 +3236,15 @@ def gated_delta_rule_mtp_wide_vec(
     # this entry point hit the same fail-fast errors.
     per_token_pool_scatter = ssm_state_indices is not None
     if per_token_pool_scatter:
-        assert (
+        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, (
             f"ssm_state_indices requires T >= 2 (got T={T_val}); "
             f"for T=1 use output_state_indices"
@@ -3253,9 +3253,9 @@ def gated_delta_rule_mtp_wide_vec(
             f"ssm_state_indices must have shape [B={B_val}, T={T_val}], "
             f"got {tuple(ssm_state_indices.shape)}"
         )
-        assert (
+        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
 
     phase_b_unroll = _select_wide_vec_phase_b_unroll(
@@ -3517,9 +3517,9 @@ def gated_delta_rule_t1_wide_vec(
     assert K_val == 128 and V_val == 128
     assert initial_state_source.dtype == torch.bfloat16
     assert tile_v in (32, 64, 128), f"tile_v must be 32/64/128, got {tile_v}"
-    assert (
+    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:
         scale = 1.0 / math.sqrt(K_val)
@@ -3827,9 +3827,9 @@ def gated_delta_rule_mtp(
             f"intermediate_states_buffer dim 0 ({buffer_size}) must equal "
             f"batch size B={B}; the buffer is batch-scoped, not pool-scoped"
         )
-        assert (
+        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
         intermediate_states = intermediate_states_buffer.reshape(
             B * cache_steps * HV, V, K
@@ -3860,28 +3860,28 @@ def gated_delta_rule_mtp(
     # results/2026-06-03/FLA_SCATTER_MODE_PLAN.md.
     per_token_pool_scatter = ssm_state_indices is not None
     if per_token_pool_scatter:
-        assert (
+        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), (
             f"ssm_state_indices must have shape [B={B}, T={T}], "
             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
     # lower per-CTA parallelism. ``_select_wide_vec_tile_v`` picks tile_v
@@ -3960,12 +3960,12 @@ def gated_delta_rule_mtp(
     # Per-request K opt-in (see gated_delta_rule_mtp_wide_vec for full rationale).
     per_request_accepted_steps = accepted_steps is not None
     if per_request_accepted_steps:
-        assert accepted_steps.shape == (
+        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
 
     # Contiguous pool -> sentinel keys + slot dim marked dynamic (pool-size

python/sglang/kernels/ops/attention/cutedsl_kda.py

@@ -1427,12 +1427,12 @@ def cutedsl_fused_sigmoid_gating_kda_update(
     N = initial_state_indices.shape[0]
 
     assert K == TILE_K, f"Current CuTe DSL KDA kernel requires K={TILE_K}, got {K}"
-    assert (
+    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, (
         "Small-batch KDA kernel requires num_v_tiles_small divisible by "
         f"{NUM_BLOCKS_PER_STATE_SMALL}, got V={V}"

python/sglang/kernels/ops/attention/decode_attention.py

@@ -1483,7 +1483,6 @@ def _lean_attention_decode_kernel(
     # Use a regular while loop instead of tl.static_range with a dynamic bound to avoid
     # Triton compiler crashes in the Coalesce pass (max_output_tile_cnt is runtime-computed).
     while iter < cta_end_tile_gid:
-
         tile_row_idx = iter // tiles_per_khead
         tile_idx = tile_row_idx * batch_size
         tile_iter = tile_row_idx * tiles_per_khead

python/sglang/kernels/ops/attention/deepseek_v4_rope.py

@@ -354,9 +354,9 @@ def apply_rotary_emb_triton(
     grid = (batch_size, n_heads if is_3d else 1, num_blocks_dim)
 
     if positions is not None:
-        assert positions.shape == (
+        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](
             x,
@@ -374,9 +374,9 @@ def apply_rotary_emb_triton(
             BLOCK_SIZE=BLOCK_SIZE,
         )
     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](
             x,
@@ -621,9 +621,9 @@ def fused_norm_rope_inplace_triton(
     if weight is not None:
         assert weight.shape == (head_dim,)
     if positions is None:
-        assert (
+        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:
         assert positions.shape == (M,) and positions.dim() == 1
 

python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py

@@ -181,9 +181,9 @@ def dequantize_k_cache_paged(
         output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache
     """
     dim_quant = quant_k_cache.shape[-1]
-    assert (
+    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))
 
     # 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

@@ -308,9 +308,9 @@ def _set_k_and_s_triton(
     assert scale_dim == 1
     if _is_hip:
         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:
         assert page_size == 64
 

python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py

@@ -155,9 +155,9 @@ def act_quant(
             - A tensor of scaling factors with dtype `torch.float32`.
     """
     assert x.is_contiguous(), "Input tensor must be contiguous"
-    assert (
+    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)
     if _is_fp8_fnuz:
         y = torch.empty_like(x, dtype=torch.float8_e4m3fnuz)
@@ -272,16 +272,16 @@ def sparse_attention_fwd_kernel_v1(
     num_stages=2,
     threads=256,
 ):
-    assert dim == tilelang.math.next_power_of_2(
+    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 (
+    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:
         sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504  # log2(e)
     else:
@@ -361,7 +361,6 @@ def sparse_attention_fwd_kernel_v1(
             T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared)
 
             for i_i in T.Pipelined(NI, num_stages=num_stages):
-
                 for bi_i in T.Parallel(BI):
                     mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] >= 0
 
@@ -446,15 +445,15 @@ def sparse_attention_fwd_kernel_v2(
     sm_scale: Optional[float] = None,
     block_I: int = 64,
 ):
-    assert dim == tilelang.math.next_power_of_2(
+    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:
         sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504  # log2(e)
     else:
@@ -1078,9 +1077,9 @@ def sparse_mla_fwd_decode_partial_fp8(
     threads=256,
 ):
     assert d_v == 512, f"only support d_v=512"
-    assert (
+    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
     # to better utilize FP8 dynamic range, then apply the inverse scale after GEMM.
@@ -1104,9 +1103,9 @@ def sparse_mla_fwd_decode_partial_fp8(
     h_per_block = 16
     # Match bf16 partial behavior: keep fixed 16-head tiles and use
     # sliced T.copy on H0:H1 for tail handling.
-    assert (
+    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
 
     batch = 1
@@ -1594,9 +1593,7 @@ def dpsk_v4_fp8_partial_kernel(
         sm_scale = sm_scale * log2e
     assert dim == 448 and tail_dim == 64
     assert topk_1 % block_I == 0
-    assert (
+    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}"
     )
     assert block_size_kv_1 > 0 and (block_size_kv_1 & (block_size_kv_1 - 1)) == 0
@@ -1605,9 +1602,7 @@ def dpsk_v4_fp8_partial_kernel(
     if is_dual:
         assert inner_iter_2 > 0, "dual-cache call requires inner_iter_2 > 0"
         assert topk_2 % block_I == 0
-        assert (
+        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}"
         )
         assert block_size_kv_2 > 0 and (block_size_kv_2 & (block_size_kv_2 - 1)) == 0
@@ -2256,12 +2251,8 @@ def dpsk_v4_combine_kernel(
 
         @T.prim_func
         def main(
-            Partial_O: T.Tensor(
+            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_2: T.Tensor([batch], INT32),  # type: ignore
             Attn_sink: T.Tensor([num_heads], FP32),  # type: ignore
@@ -2369,12 +2360,8 @@ def dpsk_v4_combine_kernel(
 
     @T.prim_func
     def main(
-        Partial_O: T.Tensor(
+        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
         Output: T.Tensor([batch, seq_len, num_heads, DT], BF16),  # type: ignore
         LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype),  # type: ignore

python/sglang/kernels/ops/attention/dsa/triton_kernel.py

@@ -99,9 +99,9 @@ def act_quant(
             - A tensor of scaling factors with dtype `torch.float32`.
     """
     assert x.is_contiguous(), "Input tensor must be contiguous"
-    assert (
+    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
     N = x.size(-1)

python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py

@@ -69,9 +69,7 @@ def _sparse_mla_fwd_kernel(
     )  # [H, D_V]
     q_tail = tl.load(
         q_rope_ptr + s_i * H * D_TAIL + h[:, None] * D_TAIL + dt[None, :]
-    ).to(
+    ).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)
     l_i = tl.zeros([H], tl.float32)
@@ -89,9 +87,7 @@ def _sparse_mla_fwd_kernel(
         )  # [BLOCK_N, D_V] -- reused as V
         kv_tail = tl.load(
             kbase + (D_V + dt)[None, :], mask=valid[:, None], other=0.0
-        ).to(
+        ).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_tail, tl.trans(kv_tail)).to(tl.float32)

python/sglang/kernels/ops/attention/dsv4/index_buf_accessor.py

@@ -209,7 +209,9 @@ def _set_k_and_s_torch(
         == num_tokens_to_write_nope
         == num_tokens_to_write_rope
         == num_tokens_to_write_scale
-    ), f"{num_tokens_to_write=} {num_tokens_to_write_nope=} {num_tokens_to_write_rope=} {num_tokens_to_write_scale=}"
+    ), (
+        f"{num_tokens_to_write=} {num_tokens_to_write_nope=} {num_tokens_to_write_rope=} {num_tokens_to_write_scale=}"
+    )
 
     assert buf.dtype == torch.uint8
     assert loc.dtype in [

python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py

@@ -110,9 +110,9 @@ def _init_compressed_attn_metadata_triton(
     # no cache-write locations. Keep the write buffers unpadded and mask those
     # rows in the kernel.
     num_write_tokens = raw_out_loc.shape[0]
-    assert (
+    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
 
     c4_out_loc = torch.empty(num_write_tokens, dtype=torch.int64, device=device)
@@ -126,12 +126,12 @@ def _init_compressed_attn_metadata_triton(
     c128_seq_lens_clamp1 = torch.empty(bs, dtype=torch.int32, device=device)
 
     if compute_page_indices:
-        assert (
+        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]
         c128_page_size = page_size // 128
         c128_cur_max_seq_len = c128_page_size * max_pages

python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py

@@ -1090,9 +1090,9 @@ class FlashAttentionForwardSm120(FlashAttentionForwardBase):
     ):
         assert blocksparse_tensors is None, "Block sparsity is not supported on SM120"
         assert (mBias is not None) == self.has_bias
-        assert (
+        assert mPageTable is None or self.paged_kv, (
-            mPageTable is None or self.paged_kv
+            "SM120 paged KV requires the dedicated DMA-warp specialization"
-        ), "SM120 paged KV requires the dedicated DMA-warp specialization"
+        )
         self._check_type(
             *(
                 t.element_type if t is not None else None
@@ -1251,7 +1251,9 @@ class FlashAttentionForwardSm120(FlashAttentionForwardBase):
         TileScheduler = (
             Sm120UniformBatchScheduler
             if is_varlen and self.direct_uniform_batch
-            else SingleTileVarlenScheduler if is_varlen else SingleTileScheduler
+            else SingleTileVarlenScheduler
+            if is_varlen
+            else SingleTileScheduler
         )
         tile_sched_args = TileSchedulerArguments(
             num_block=cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m),

python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py

@@ -75,9 +75,9 @@ class Sm120UniformBatchScheduler:
         loc=None,
         ip=None,
     ) -> Params:
-        assert (
+        assert scheduling_mode == SchedulingMode.STATIC, (
-            scheduling_mode == SchedulingMode.STATIC
+            f"SM120 uniform-batch scheduler only supports STATIC, got {scheduling_mode!r}"
-        ), f"SM120 uniform-batch scheduler only supports STATIC, got {scheduling_mode!r}"
+        )
         return Sm120UniformBatchScheduler.Params.create(args, loc=loc, ip=ip)
 
     @staticmethod

python/sglang/kernels/ops/attention/fla/chunk.py

@@ -85,7 +85,6 @@ def chunk_gated_delta_rule_fwd(
 
 
 class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
-
     @staticmethod
     @input_guard
     @autocast_custom_fwd
@@ -207,12 +206,12 @@ def chunk_gated_delta_rule(
         )
     """
     assert q.dtype == k.dtype == v.dtype
-    assert (
+    assert q.dtype != torch.float32, (
-        q.dtype != torch.float32
+        "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
-    ), "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
+    )
-    assert (
+    assert len(beta.shape) == 3, (
-        len(beta.shape) == 3
+        "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise."
-    ), "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise."
+    )
 
     if head_first:
         raise DeprecationWarning(

python/sglang/kernels/ops/attention/fla/chunk_delta_h.py

@@ -82,9 +82,10 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
     i_v, i_nh = tl.program_id(0), tl.program_id(1)
     i_n, i_h = i_nh // H, i_nh % H
     if IS_VARLEN:
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+        bos, eos = (
-            cu_seqlens + i_n + 1
+            tl.load(cu_seqlens + i_n).to(tl.int32),
-        ).to(tl.int32)
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
         NT = tl.cdiv(T, BT)
         boh = tl.load(chunk_offsets + i_n).to(tl.int32)
@@ -326,9 +327,9 @@ def chunk_gated_delta_rule_fwd_h(
     chunk_indices: Optional[torch.LongTensor] = None,
     use_exp2: bool = False,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    assert not (
+    assert not (use_exp2 and g is not None), (
-        use_exp2 and g is not None
+        "use_exp2 covers only the per-channel gk path; scalar g stays natural-exp"
-    ), "use_exp2 covers only the per-channel gk path; scalar g stays natural-exp"
+    )
     B, T, Hg, K, V = *k.shape, u.shape[-1]
     H = u.shape[-2]
     BT = CHUNK_SIZE

python/sglang/kernels/ops/attention/fla/chunk_fwd.py

@@ -71,12 +71,14 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel(
     i_b, i_h = i_bh // H, i_bh % H
 
     if IS_VARLEN:
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
+        i_n, i_t = (
-            chunk_indices + i_t * 2 + 1
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
-        ).to(tl.int32)
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+        )
-            cu_seqlens + i_n + 1
+        bos, eos = (
-        ).to(tl.int32)
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
     else:
         bos, eos = i_b * T, i_b * T + T

python/sglang/kernels/ops/attention/fla/chunk_intra.py

@@ -88,12 +88,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
     i_b, i_h = i_bh // H, i_bh % H
 
     if IS_VARLEN:
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
+        i_n, i_t = (
-            chunk_indices + i_t * 2 + 1
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
-        ).to(tl.int32)
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+        )
-            cu_seqlens + i_n + 1
+        bos, eos = (
-        ).to(tl.int32)
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
     else:
         bos, eos = i_b * T, i_b * T + T
@@ -819,12 +821,14 @@ def chunk_kda_fwd_kernel_intra_sub_chunk(
     i_b, i_h = i_bh // H, i_bh % H
 
     if IS_VARLEN:
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
+        i_n, i_t = (
-            chunk_indices + i_t * 2 + 1
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
-        ).to(tl.int32)
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+        )
-            cu_seqlens + i_n + 1
+        bos, eos = (
-        ).to(tl.int32)
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
     else:
         bos, eos = i_b * T, i_b * T + T

python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py

@@ -62,9 +62,10 @@ def chunk_kda_fwd_kernel_intra_token_parallel(
                     left = mid + 1
         i_n = left
 
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+        bos, eos = (
-            cu_seqlens + i_n + 1
+            tl.load(cu_seqlens + i_n).to(tl.int32),
-        ).to(tl.int32)
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
         i_t = i_tg - bos
     else:

python/sglang/kernels/ops/attention/fla/chunk_o.py

@@ -53,12 +53,14 @@ def chunk_fwd_kernel_o(
 
     if IS_VARLEN:
         i_tg = i_t
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
+        i_n, i_t = (
-            chunk_indices + i_t * 2 + 1
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
-        ).to(tl.int32)
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+        )
-            cu_seqlens + i_n + 1
+        bos, eos = (
-        ).to(tl.int32)
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
         NT = tl.cdiv(T, BT)
     else:

python/sglang/kernels/ops/attention/fla/cumsum.py

@@ -37,12 +37,14 @@ def chunk_local_cumsum_scalar_kernel(
     i_t, i_bh = tl.program_id(0), tl.program_id(1)
     i_b, i_h = i_bh // H, i_bh % H
     if IS_VARLEN:
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
+        i_n, i_t = (
-            chunk_indices + i_t * 2 + 1
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
-        ).to(tl.int32)
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+        )
-            cu_seqlens + i_n + 1
+        bos, eos = (
-        ).to(tl.int32)
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
     else:
         bos, eos = i_b * T, i_b * T + T
@@ -97,12 +99,14 @@ def chunk_local_cumsum_vector_kernel(
     i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
     i_b, i_h = i_bh // H, i_bh % H
     if IS_VARLEN:
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
+        i_n, i_t = (
-            chunk_indices + i_t * 2 + 1
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
-        ).to(tl.int32)
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+        )
-            cu_seqlens + i_n + 1
+        bos, eos = (
-        ).to(tl.int32)
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
     else:
         bos, eos = i_b * T, i_b * T + T
@@ -169,9 +173,9 @@ def chunk_local_cumsum_scalar(
         B, H, T = g.shape
     else:
         B, T, H = g.shape
-    assert chunk_size == 2 ** (
+    assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
-        chunk_size.bit_length() - 1
+        "chunk_size must be a power of 2"
-    ), "chunk_size must be a power of 2"
+    )
     BT = chunk_size
     if chunk_indices is None and cu_seqlens is not None:
         chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
@@ -216,9 +220,9 @@ def chunk_local_cumsum_vector(
     if chunk_indices is None and cu_seqlens is not None:
         chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
     NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
-    assert chunk_size == 2 ** (
+    assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
-        chunk_size.bit_length() - 1
+        "chunk_size must be a power of 2"
-    ), "chunk_size must be a power of 2"
+    )
 
     g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
 
@@ -260,9 +264,9 @@ def chunk_local_cumsum(
     **kwargs,
 ) -> torch.Tensor:
     if cu_seqlens is not None:
-        assert (
+        assert g.shape[0] == 1, (
-            g.shape[0] == 1
+            "Only batch size 1 is supported when cu_seqlens are provided"
-        ), "Only batch size 1 is supported when cu_seqlens are provided"
+        )
     if len(g.shape) == 3:
         return chunk_local_cumsum_scalar(
             g=g,

python/sglang/kernels/ops/attention/fla/fused_norm_gate.py

@@ -390,9 +390,9 @@ class FusedRMSNormGated(nn.Module):
         residual_in_fp32: bool = False,
     ) -> torch.Tensor:
         if _use_cpu:
-            assert (
+            assert self.activation == "silu", (
-                self.activation == "silu"
+                "CPU rmsnorm_gated currently only supports activation silu"
-            ), "CPU rmsnorm_gated currently only supports activation silu"
+            )
             return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu(
                 x, self.weight, g, self.eps
             )

python/sglang/kernels/ops/attention/fla/fused_recurrent.py

@@ -43,9 +43,10 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
     i_n, i_hv = i_nh // HV, i_nh % HV
     i_h = i_hv // (HV // H)
     if IS_VARLEN:
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(
+        bos, eos = (
-            cu_seqlens + i_n + 1
+            tl.load(cu_seqlens + i_n).to(tl.int64),
-        ).to(tl.int64)
+            tl.load(cu_seqlens + i_n + 1).to(tl.int64),
+        )
         all = T
         T = eos - bos
     else:
@@ -708,7 +709,6 @@ def fused_recurrent_kda_packed_decode(
 
 
 class FusedRecurrentFunction(torch.autograd.Function):
-
     @staticmethod
     @input_guard
     def forward(
@@ -907,9 +907,10 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
     i_n, i_hv = i_nh // HV, i_nh % HV
     i_h = i_hv // (HV // H)
     if IS_VARLEN:
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(
+        bos, eos = (
-            cu_seqlens + i_n + 1
+            tl.load(cu_seqlens + i_n).to(tl.int64),
-        ).to(tl.int64)
+            tl.load(cu_seqlens + i_n + 1).to(tl.int64),
+        )
         all = T
         T = eos - bos
     else:
@@ -1144,7 +1145,6 @@ def fused_recurrent_gated_delta_rule_update_fwd(
 
 
 class FusedRecurrentUpdateFunction(torch.autograd.Function):
-
     @staticmethod
     @input_guard
     def forward(

python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py

@@ -678,9 +678,9 @@ def _launch_gdn_spec(
     num_slots, HV, V, K = checkpoint_state.shape
     H = k.shape[1]
     B = query_start_loc.shape[0] - 1
-    assert (
+    assert max_cache_len & (max_cache_len - 1) == 0, (
-        max_cache_len & (max_cache_len - 1) == 0
+        "circular cache requires power-of-two max_cache_len"
-    ), "circular cache requires power-of-two max_cache_len"
+    )
     assert d_cache.shape[2] == max_cache_len
 
     BK = triton.next_power_of_2(K)

python/sglang/kernels/ops/attention/fla/kda.py

@@ -1046,18 +1046,18 @@ def kda_gate_chunk_cumsum(
         Cumulative-summed gated tensor of shape [B, T, H, K].
     """
     if cu_seqlens is not None:
-        assert (
+        assert g.shape[0] == 1, (
-            g.shape[0] == 1
+            "Only batch size 1 is supported when cu_seqlens are provided"
-        ), "Only batch size 1 is supported when cu_seqlens are provided"
+        )
     assert len(g.shape) == 4
     B, T, H, S = g.shape
     BT = chunk_size
     if chunk_indices is None and cu_seqlens is not None:
         chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
     NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
-    assert chunk_size == 2 ** (
+    assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
-        chunk_size.bit_length() - 1
+        "chunk_size must be a power of 2"
-    ), "chunk_size must be a power of 2"
+    )
 
     g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
 

python/sglang/kernels/ops/attention/fla/l2norm.py

@@ -120,7 +120,6 @@ def l2norm_fwd(
 
 
 class L2NormFunction(torch.autograd.Function):
-
     @staticmethod
     @input_guard
     def forward(ctx, x, eps=1e-6, output_dtype=None):
@@ -137,7 +136,6 @@ l2_norm = l2norm
 
 
 class L2Norm(nn.Module):
-
     def __init__(self, eps: float = 1e-6, output_dtype: Optional[torch.dtype] = None):
         super().__init__()
         self.eps = eps

python/sglang/kernels/ops/attention/fla/layernorm_gated.py

@@ -345,7 +345,6 @@ def rms_norm_gated(
 
 
 class LayerNormFn(torch.autograd.Function):
-
     @staticmethod
     def forward(
         ctx,
@@ -389,7 +388,6 @@ def layernorm_fn(
 
 
 class LayerNorm(torch.nn.Module):
-
     def __init__(
         self,
         hidden_size,
@@ -431,7 +429,6 @@ class LayerNorm(torch.nn.Module):
 
 
 class RMSNorm(torch.nn.Module):
-
     def __init__(
         self,
         hidden_size,
@@ -465,7 +462,9 @@ class RMSNorm(torch.nn.Module):
                 self.norm_before_gate
                 and self.group_size is None
                 and self.activation == "swish"
-            ), "CPU rmsnorm_gated currently only supports norm before gate without group size or activation other than swish"
+            ), (
+                "CPU rmsnorm_gated currently only supports norm before gate without group size or activation other than swish"
+            )
             return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu(
                 x, self.weight, z, self.eps
             )

python/sglang/kernels/ops/attention/fla/utils.py

@@ -326,9 +326,9 @@ if torch_release >= (2, 4):
         return device_torch_lib.device(index)
 
 else:
-    assert (
+    assert device == "cuda", (
-        device == "cuda"
+        "Only cuda device is supported for PyTorch version < 2.4.0."
-    ), "Only cuda device is supported for PyTorch version < 2.4.0."
+    )
     autocast_custom_fwd = device_torch_lib.amp.custom_fwd
     autocast_custom_bwd = device_torch_lib.amp.custom_bwd
 

python/sglang/kernels/ops/attention/fla/wy_fast.py

@@ -43,12 +43,14 @@ def recompute_w_u_fwd_kernel(
     i_t, i_bh = tl.program_id(0), tl.program_id(1)
     i_b, i_h = i_bh // H, i_bh % H
     if IS_VARLEN:
-        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
+        i_n, i_t = (
-            chunk_indices + i_t * 2 + 1
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
-        ).to(tl.int32)
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
-        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+        )
-            cu_seqlens + i_n + 1
+        bos, eos = (
-        ).to(tl.int32)
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
         T = eos - bos
     else:
         bos, eos = i_b * T, i_b * T + T

python/sglang/kernels/ops/attention/flash_mla_sm120.py

@@ -77,9 +77,7 @@ def _gather_and_dequant(k_cache, indices, page_size):
     raw_pages = k_cache.as_strided(
         (num_pages, page_bytes),
         (page_bytes, 1),
-    ).view(
+    ).view(torch.uint8)  # (num_pages, page_bytes) uint8
-        torch.uint8
-    )  # (num_pages, page_bytes) uint8
     # Note: float8_e4m3fn and uint8 are both 1 byte, view is safe
 
     # Compute byte offsets within each page