diff --git a/Magpie/modes/benchmark/gap_analysis.py b/Magpie/modes/benchmark/gap_analysis.py index 8527959..2ec4ff6 100644 --- a/Magpie/modes/benchmark/gap_analysis.py +++ b/Magpie/modes/benchmark/gap_analysis.py @@ -190,7 +190,14 @@ def to_csv( ] if find_kernel_sources: from Magpie.tools.amd_kernel_finder import KernelSourceInfo - headers.extend(KernelSourceInfo.csv_headers()) + source_headers = KernelSourceInfo.csv_headers() + source_column_start = len(headers) + 1 + headers.extend(source_headers) + source_column_end = len(headers) + f.write( + f"# Warning: columns {source_column_start}-{source_column_end} " + "from find_kernel_sources are experimental.\n" + ) writer.writerow(headers) diff --git a/Magpie/tools/amd_kernel_finder/finder.py b/Magpie/tools/amd_kernel_finder/finder.py index 259323f..2ba98f3 100644 --- a/Magpie/tools/amd_kernel_finder/finder.py +++ b/Magpie/tools/amd_kernel_finder/finder.py @@ -158,7 +158,7 @@ def search(self, kernel_name: str) -> KernelSourceInfo: source_match = self.searcher.search_source(parsed) # Search for test - test_match = self.searcher.search_test(parsed, source_match) + test_match = self.searcher.search_test(parsed, source_match, category=category) # Search for PyTorch eager baseline reference. We hand over the # already-computed test_match + category so the searcher does not diff --git a/Magpie/tools/amd_kernel_finder/parser.py b/Magpie/tools/amd_kernel_finder/parser.py index f6c3dec..1f5cdf8 100644 --- a/Magpie/tools/amd_kernel_finder/parser.py +++ b/Magpie/tools/amd_kernel_finder/parser.py @@ -13,6 +13,26 @@ from .models import KernelKind, KernelCategory, ParsedKernelName +SGLANG_KERNEL_TOKENS = ( + "_zn7sgl_hip", + "sgl_hip::", + "write_req_to_token_pool", + "create_flashinfer_kv_indices", + "future_token_ids", + "kn_get_mla_metadata", + "mla_metadata", + "clamp_position", + "compute_position", + "set_mla_kv_buffer", +) + + +def is_sglang_kernel_name(name: str) -> bool: + """Return True for SGLang-owned runtime helper kernels.""" + name_lc = name.lower() + return any(token in name_lc for token in SGLANG_KERNEL_TOKENS) + + class KernelNameParser: """Parse kernel names from profiler output.""" @@ -108,6 +128,11 @@ def _classify_kind(self, name: str) -> KernelKind: # Check for aiter kernels (before CK and Triton checks) if self.AITER_PATTERN.search(name): return KernelKind.AITER + + # SGLang owns HIP/C++ and Triton-emitted runtime kernels; keep the kind + # generic, like vLLM, and let source search route to $SGLANG_DIR. + if is_sglang_kernel_name(name): + return KernelKind.HIP_CPP # Check for CK tile (handles both mangled and readable names) if self.CK_PATTERN.search(name): @@ -321,12 +346,37 @@ def _parse_hip(self, name: str) -> ParsedKernelName: extra['namespace'] = match.group(1) function_name = match.group(2) else: + if 'act_and_mul_kernel' in name: + extra['namespace'] = 'sgl_hip' + function_name = 'act_and_mul_kernel' + elif 'write_req_to_token_pool_triton' in name: + function_name = 'write_req_to_token_pool_triton' + elif 'create_flashinfer_kv_indices_triton' in name: + function_name = 'create_flashinfer_kv_indices_triton' + elif 'kn_get_mla_metadata' in name: + function_name = 'kn_get_mla_metadata' + elif 'clamp_position_kernel' in name: + function_name = 'clamp_position_kernel' + elif 'compute_position_kernel' in name: + function_name = 'compute_position_kernel' + elif 'resolve_future_token_ids_kernel' in name: + function_name = 'resolve_future_token_ids_kernel' + elif 'set_mla_kv_buffer_kernel' in name: + function_name = 'set_mla_kv_buffer_kernel' + elif '_ZN7sgl_hip' in name: + extra['namespace'] = 'sgl_hip' + function_name = 'SGLang HIP kernel' # Try without namespace + else: + match = re.search(r'void (?:\(anonymous namespace\)::)?(\w+)', name) + if match: + function_name = match.group(1) + else: + function_name = "HIP kernel" + if function_name == "HIP kernel": match = re.search(r'void (\w+)<', name) if match: function_name = match.group(1) - else: - function_name = "HIP kernel" # Extract dtype if '__hip_bfloat16' in name: diff --git a/Magpie/tools/amd_kernel_finder/repo_config.py b/Magpie/tools/amd_kernel_finder/repo_config.py index 586cdb5..6acc70f 100644 --- a/Magpie/tools/amd_kernel_finder/repo_config.py +++ b/Magpie/tools/amd_kernel_finder/repo_config.py @@ -254,6 +254,26 @@ class RepoConfig: ], }, ), + "sglang": RepoConfig( + name="sglang", + var_name="$SGLANG_DIR", + github_base="https://github.com/sgl-project/sglang", + source_paths={ + "sglang": [ + "python/sglang/", + "sglang/", + "sgl-kernel/", + "csrc/", + ], + }, + test_paths={ + "sglang": [ + "test/", + "tests/", + "sgl-kernel/tests/", + ], + }, + ), } @@ -268,6 +288,7 @@ class RepoConfig: "$VLLM_DIR": ("$VLLM_DIR", ""), # vllm is its own repo "$PYTORCH_DIR": ("$PYTORCH_DIR", ""), # pytorch is its own repo "$AITER_DIR": ("$AITER_DIR", ""), # aiter is its own repo + "$SGLANG_DIR": ("$SGLANG_DIR", ""), # sglang is its own repo "$ROCM_SYSTEMS_DIR": ("$ROCM_SYSTEMS_DIR", ""), # rocm-systems super-repo (clr, hip, rocprofiler, etc) } @@ -280,6 +301,7 @@ class RepoConfig: "vllm": "https://github.com/vllm-project/vllm/blob/main/{path}", "pytorch": "https://github.com/pytorch/pytorch/blob/main/{path}", "aiter": "https://github.com/ROCm/aiter/blob/main/{path}", + "sglang": "https://github.com/sgl-project/sglang/blob/main/{path}", } @@ -313,6 +335,13 @@ def detect_repo_type(repo_path: str) -> Optional[str]: # Check for pytorch if (path / "aten").exists() and (path / "torch").exists(): return "pytorch" + + # Check for SGLang + if ( + (path / "python" / "sglang").exists() + or (path / "sglang").exists() + ) and ((path / "sgl-kernel").exists() or (path / "test").exists() or (path / "tests").exists()): + return "sglang" return None diff --git a/Magpie/tools/amd_kernel_finder/repo_manager.py b/Magpie/tools/amd_kernel_finder/repo_manager.py index 6c6906f..cead0bf 100644 --- a/Magpie/tools/amd_kernel_finder/repo_manager.py +++ b/Magpie/tools/amd_kernel_finder/repo_manager.py @@ -25,6 +25,7 @@ "pytorch": "https://github.com/pytorch/pytorch.git", "rocm-systems": "https://github.com/ROCm/rocm-systems.git", # ROCm super-repo (clr, hip, rocprofiler, etc) "aiter": "https://github.com/ROCm/aiter.git", + "sglang": "https://github.com/sgl-project/sglang.git", } KERNEL_REPO_MAP = { @@ -38,7 +39,7 @@ } # All repos to clone when force_all is True -ALL_REPOS = ["rocm-libraries", "triton", "vllm", "pytorch", "aiter", "rocm-systems"] +ALL_REPOS = ["rocm-libraries", "triton", "vllm", "pytorch", "aiter", "sglang", "rocm-systems"] @@ -153,11 +154,12 @@ def get_repos_for_kernels(self, kernel_names: List[str], force_all: bool = False if force_all: return ALL_REPOS.copy() - from .parser import KernelNameParser + from .parser import KernelNameParser, is_sglang_kernel_name parser = KernelNameParser() kinds = set() has_vllm = False + has_sglang = False for name in kernel_names: parsed = parser.parse(name) @@ -165,6 +167,8 @@ def get_repos_for_kernels(self, kernel_names: List[str], force_all: bool = False if "vllm::" in name or "vllm" in name.lower(): has_vllm = True + if is_sglang_kernel_name(name): + has_sglang = True repos = set() for kind in kinds: @@ -173,6 +177,8 @@ def get_repos_for_kernels(self, kernel_names: List[str], force_all: bool = False if has_vllm: repos.add("vllm") + if has_sglang: + repos.add("sglang") return list(repos) diff --git a/Magpie/tools/amd_kernel_finder/searcher.py b/Magpie/tools/amd_kernel_finder/searcher.py index f900f4e..7919ee0 100644 --- a/Magpie/tools/amd_kernel_finder/searcher.py +++ b/Magpie/tools/amd_kernel_finder/searcher.py @@ -24,6 +24,7 @@ BaselineRefMatch, TritonRefMatch, ) +from .parser import is_sglang_kernel_name from .repo_config import RepoConfig, SUBPROJECT_MAPPINGS, GITHUB_URL_TEMPLATES logger = logging.getLogger(__name__) @@ -801,13 +802,20 @@ def search_source(self, parsed: ParsedKernelName) -> Optional[SourceMatch]: return None - def search_test(self, parsed: ParsedKernelName, source: Optional[SourceMatch] = None) -> Optional[TestMatch]: + def search_test( + self, + parsed: ParsedKernelName, + source: Optional[SourceMatch] = None, + category: Optional[KernelCategory] = None, + ) -> Optional[TestMatch]: """ Search for test files and generate test command. Args: parsed: Parsed kernel name information source: Optional source match for context + category: Optional category fallback for kernels whose kind cannot + route to a repo-specific test searcher. Returns: TestMatch if found, None otherwise @@ -815,20 +823,27 @@ def search_test(self, parsed: ParsedKernelName, source: Optional[SourceMatch] = if parsed.kind == KernelKind.ANNOTATION: return None + match = None if parsed.kind == KernelKind.TRITON_JIT: - return self._search_triton_test(parsed, source) + match = self._search_triton_test(parsed, source) elif parsed.kind == KernelKind.TENSILE_GEMM: - return self._search_tensile_test(parsed) + match = self._search_tensile_test(parsed) elif parsed.kind == KernelKind.CK_TILE: - return self._search_ck_test(parsed) + match = self._search_ck_test(parsed) elif parsed.kind == KernelKind.ATEN_NATIVE: - return self._search_aten_test(parsed) + match = self._search_aten_test(parsed) elif parsed.kind == KernelKind.HIP_CPP: - return self._search_hip_test(parsed, source) + match = self._search_hip_test(parsed, source) elif parsed.kind == KernelKind.AITER: - return self._search_aiter_test(parsed, source) - - return None + match = self._search_aiter_test(parsed, source) + + if match is not None: + return match + + # Some profiler names lose the suffix/namespace that lets the parser + # identify their KernelKind. Still route them to a category-level test + # when we have a stable op category such as MoE GEMM or layernorm. + return self._search_category_test(category) # ------------------------------------------------------------------ # Baseline (PyTorch / Triton) reference search @@ -930,6 +945,46 @@ def search_baseline_ref( return None + def _search_category_test(self, category: Optional[KernelCategory]) -> Optional[TestMatch]: + """Route unknown-kind kernels to the category's canonical test file.""" + if category is None: + return None + + for display_path in CATEGORY_TO_TEST_FILES.get(category, []): + match = self._test_match_from_display_path(display_path) + if match is not None: + return match + + return None + + @staticmethod + def _test_match_from_display_path(display_path: str) -> Optional[TestMatch]: + """Create a TestMatch from a `$REPO_DIR/path` category mapping.""" + if not display_path: + return None + + repo_var = "" + test_file = display_path + if display_path.startswith("$"): + repo_var, _, test_file = display_path.partition("/") + if not test_file: + return None + + if repo_var == "$AITER_DIR": + test_cmd = f"cd {repo_var} && pytest {test_file} -v" + elif repo_var == "$VLLM_DIR": + test_cmd = f"cd {repo_var} && pytest {test_file} -q" + elif repo_var == "$PYTORCH_DIR": + test_cmd = f"pytest {display_path} -q" + else: + test_cmd = f"pytest {display_path} -q" + + return TestMatch( + test_file=test_file, + test_cmd=test_cmd, + repo_var=repo_var, + ) + def search_triton_ref( self, parsed: ParsedKernelName, @@ -1460,6 +1515,12 @@ def _search_hip_source(self, parsed: ParsedKernelName) -> Optional[SourceMatch]: repo_name="vllm", repo_var="$VLLM_DIR", ) + + # Check SGLang kernels by namespace/name. + if namespace == "sgl_hip" or self._is_sglang_kernel(original_name): + match = self._search_sglang_source(function_name, original_name) + if match is not None: + return match # Search in rocm-libraries rocm_libs = self._repo_var_map.get("$ROCM_LIBRARIES_DIR") @@ -1476,6 +1537,47 @@ def _search_hip_source(self, parsed: ParsedKernelName) -> Optional[SourceMatch]: ) return None + + @staticmethod + def _is_sglang_kernel(name: str) -> bool: + return is_sglang_kernel_name(name) + + def _search_sglang_source(self, function_name: str, original_name: str) -> Optional[SourceMatch]: + """Search SGLang source for runtime HIP/Triton helper kernels.""" + sglang_path = self._repo_var_map.get("$SGLANG_DIR") + if not sglang_path: + return None + + candidates = [function_name] + original_lc = original_name.lower() + for token in ( + "write_req_to_token_pool_triton", + "create_flashinfer_kv_indices_triton", + "resolve_future_token_ids_kernel", + "kn_get_mla_metadata", + "clamp_position_kernel", + "compute_position_kernel", + "set_mla_kv_buffer_kernel", + "act_and_mul_kernel", + ): + if token in original_lc and token not in candidates: + candidates.append(token) + + for candidate in candidates: + if not candidate or candidate == "HIP kernel": + continue + pattern = rf"(def|void|__global__|template).*{re.escape(candidate)}|{re.escape(candidate)}" + files = self._search_files(pattern, sglang_path, ["py", "cpp", "cu", "hip", "hpp"]) + if files: + rel_path = os.path.relpath(files[0], sglang_path) + return SourceMatch( + file_path=rel_path, + symbol=candidate, + repo_name="sglang", + repo_var="$SGLANG_DIR", + ) + + return None def _search_inductor_source(self, parsed: ParsedKernelName) -> Optional[SourceMatch]: """Search for torch.inductor generated kernel.""" @@ -1669,6 +1771,13 @@ def _search_hip_test(self, parsed: ParsedKernelName, """Search for HIP/CUDA kernel tests.""" namespace = parsed.namespace function_name = parsed.function_name + + if (source and source.repo_name == "sglang") or namespace == "sgl_hip" or self._is_sglang_kernel(parsed.original_name): + return TestMatch( + test_file="test/", + test_cmd="cd $SGLANG_DIR && pytest test/ -q", + repo_var="$SGLANG_DIR", + ) original_name = parsed.original_name # Known HIP kernel test mappings diff --git a/Magpie/tools/amd_kernel_finder/tests/test_refs.py b/Magpie/tools/amd_kernel_finder/tests/test_refs.py index d43fb0c..a3c332d 100644 --- a/Magpie/tools/amd_kernel_finder/tests/test_refs.py +++ b/Magpie/tools/amd_kernel_finder/tests/test_refs.py @@ -44,9 +44,11 @@ import pytest -from tools.amd_kernel_finder.models import KernelCategory -from tools.amd_kernel_finder.parser import KernelNameParser -from tools.amd_kernel_finder.searcher import ( +from Magpie.modes.benchmark.gap_analysis import GapAnalysisResult, KernelStat +from Magpie.tools.amd_kernel_finder.models import KernelCategory, KernelSourceInfo +from Magpie.tools.amd_kernel_finder.parser import KernelNameParser +from Magpie.tools.amd_kernel_finder.repo_manager import RepoManager +from Magpie.tools.amd_kernel_finder.searcher import ( KernelSourceSearcher, _tokenize_symbol, _extract_compute_dtypes, @@ -126,6 +128,139 @@ def test_quant_dtype_extraction(self): assert _extract_quant_dtypes("_gemm_a16_w16_kernel") == frozenset() +class TestRepoRouting: + """Unit tests for repo-owner routing without requiring cloned repos.""" + + def test_sglang_kernel_names_route_as_hip_cpp_without_new_kernel_kind(self): + parser = KernelNameParser() + + parsed = parser.parse( + "_ZN7sgl_hip10activation18act_and_mul_kernelI14__hip_bfloat16" + "TnPFT_RKS3_EXadL_Z4siluIS2_ES3_S5_EEEEvPS3_PS4_i" + ) + assert parsed.kind.value == "hip_cpp" + assert parsed.namespace == "sgl_hip" + assert parsed.function_name == "act_and_mul_kernel" + + parsed = parser.parse("write_req_to_token_pool_triton") + assert parsed.kind.value == "hip_cpp" + assert parsed.function_name == "write_req_to_token_pool_triton" + + parsed = parser.parse("flashinfer_attention.kd") + assert parsed.kind.value == "triton_jit" + + def test_sglang_kernel_names_trigger_sglang_repo_detection(self, tmp_path): + manager = RepoManager(base_dir=str(tmp_path / "repos")) + + repos = manager.get_repos_for_kernels([ + "write_req_to_token_pool_triton", + "create_flashinfer_kv_indices_triton", + ]) + + assert "sglang" in repos + + flashinfer_repos = manager.get_repos_for_kernels([ + "flashinfer_attention.kd", + ]) + + assert "sglang" not in flashinfer_repos + + def test_sglang_source_and_test_routing_uses_repo_owner_pattern(self, tmp_path): + sglang_repo = tmp_path / "sglang" + source_file = ( + sglang_repo + / "python" + / "sglang" + / "srt" + / "mem_cache" + / "common.py" + ) + source_file.parent.mkdir(parents=True, exist_ok=True) + source_file.write_text( + "def write_req_to_token_pool_triton(req):\n return req\n", + encoding="utf-8", + ) + test_file = sglang_repo / "test" / "test_placeholder.py" + test_file.parent.mkdir(parents=True, exist_ok=True) + test_file.write_text("def test_placeholder(): pass\n", encoding="utf-8") + + parser = KernelNameParser() + parsed = parser.parse("write_req_to_token_pool_triton") + searcher = KernelSourceSearcher([str(sglang_repo)], auto_install_ripgrep=False) + + source = searcher.search_source(parsed) + test = searcher.search_test(parsed, source) + + assert source is not None + assert source.repo_name == "sglang" + assert source.repo_var == "$SGLANG_DIR" + assert source.file_path == "python/sglang/srt/mem_cache/common.py" + assert test is not None + assert test.display_path == "$SGLANG_DIR/test/" + assert test.test_cmd == "cd $SGLANG_DIR && pytest test/ -q" + + def test_sglang_source_search_returns_none_when_no_source_found(self, tmp_path): + sglang_repo = tmp_path / "sglang" + (sglang_repo / "python" / "sglang").mkdir(parents=True) + (sglang_repo / "test").mkdir() + + parser = KernelNameParser() + parsed = parser.parse("write_req_to_token_pool_triton") + searcher = KernelSourceSearcher([str(sglang_repo)], auto_install_ripgrep=False) + + source = searcher.search_source(parsed) + test = searcher.search_test(parsed, source) + + assert source is None + assert test is not None + assert test.display_path == "$SGLANG_DIR/test/" + + def test_unknown_moe_gemm_uses_category_test_fallback(self): + parser = KernelNameParser() + parsed = parser.parse("moe_gemm1_0") + searcher = KernelSourceSearcher([], auto_install_ripgrep=False) + + test = searcher.search_test(parsed, category=KernelCategory.MOE_GEMM) + + assert test is not None + assert test.display_path == "$AITER_DIR/op_tests/test_moe.py" + assert test.test_cmd == "cd $AITER_DIR && pytest op_tests/test_moe.py -v" + + +class TestGapAnalysisCsv: + """CSV behavior for find_kernel_sources metadata columns.""" + + def test_find_kernel_sources_warning_uses_dynamic_column_range(self, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + + result = GapAnalysisResult( + merged_kernels=[ + KernelStat( + name="flashinfer_attention.kd", + total_duration_us=10.0, + calls=1, + durations_us=[10.0], + ) + ], + total_duration_us=10.0, + ) + + output_path = result.to_csv( + tmp_path / "gap_analysis.csv", + find_kernel_sources=True, + auto_clone_repos=False, + ) + + lines = output_path.read_text(encoding="utf-8").splitlines() + source_column_start = 6 + 1 + source_column_end = 6 + len(KernelSourceInfo.csv_headers()) + + assert ( + f"# Warning: columns {source_column_start}-{source_column_end} " + "from find_kernel_sources are experimental." + ) in lines + + # ---------------------------------------------------------------------------- # 2. End-to-end golden corpus # ----------------------------------------------------------------------------