1use serde::{Deserialize, Serialize};
6use serde_json;
7use std::env;
8use std::fs;
9use std::path::{Path, PathBuf};
10use std::process::Command;
11use std::time::Instant;
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct BenchmarkResult {
15 pub name: String,
16 pub success: bool,
17 pub execution_time_ms: u128,
18 pub metrics: serde_json::Value,
19 pub errors: Vec<String>,
20}
21
22fn helper_dir_candidates() -> Vec<PathBuf> {
23 let mut dirs = Vec::new();
24 if let Some(home) = dirs::home_dir() {
25 dirs.push(home.join(".rusty-stack").join("tmp"));
26 }
27 dirs.push(PathBuf::from("/tmp").join("rusty-stack-tmp"));
28 dirs
29}
30
31fn ensure_helper_script() -> Result<PathBuf, String> {
32 let body = PY_HELPER.trim_start();
33 let mut last_error = String::new();
34
35 for dir in helper_dir_candidates() {
36 if let Err(err) = fs::create_dir_all(&dir) {
37 last_error = format!("failed to create helper dir {}: {}", dir.display(), err);
38 continue;
39 }
40 let path = dir.join("rusty_bench.py");
41 match fs::write(&path, body) {
42 Ok(_) => return Ok(path),
43 Err(err) => {
44 last_error = format!("failed to write helper script {}: {}", path.display(), err);
45 continue;
46 }
47 }
48 }
49
50 Err(last_error)
51}
52
53fn resolve_benchmark_python() -> String {
54 for key in [
55 "MLSTACK_BENCHMARK_PYTHON",
56 "MLSTACK_PYTHON_BIN",
57 "UV_PYTHON",
58 ] {
59 if let Ok(value) = env::var(key) {
60 let value = value.trim();
61 if !value.is_empty() {
62 return value.to_string();
63 }
64 }
65 }
66
67 {
68 let candidate = "/usr/local/bin/python3";
69 let path = Path::new(candidate);
70 if path.exists() {
71 return candidate.to_string();
72 }
73 }
74
75 for candidate in ["python3", "python"] {
76 if Command::new(candidate).arg("--version").output().is_ok() {
77 return candidate.to_string();
78 }
79 }
80
81 "python3".to_string()
82}
83
84fn extract_helper_payload(stdout: &str) -> Option<serde_json::Value> {
85 let marker = "---BENCHMARK_RESULTS_START---";
86 let search = if let Some(pos) = stdout.rfind(marker) {
87 &stdout[pos + marker.len()..]
88 } else {
89 stdout
90 };
91
92 if let Some(start) = search.find('{') {
93 let mut depth = 0usize;
94 let mut in_string = false;
95 let mut escaped = false;
96 for (idx, ch) in search[start..].char_indices() {
97 if in_string {
98 if escaped {
99 escaped = false;
100 continue;
101 }
102 match ch {
103 '\\' => escaped = true,
104 '"' => in_string = false,
105 _ => {}
106 }
107 continue;
108 }
109
110 match ch {
111 '"' => in_string = true,
112 '{' => depth += 1,
113 '}' => {
114 if depth == 0 {
115 return None;
116 }
117 depth -= 1;
118 if depth == 0 {
119 let end = start + idx + 1;
120 return serde_json::from_str::<serde_json::Value>(&search[start..end]).ok();
121 }
122 }
123 _ => {}
124 }
125 }
126 }
127
128 None
129}
130
131fn run_python_benchmark(name: &str) -> BenchmarkResult {
132 let start = Instant::now();
133 let python_bin = resolve_benchmark_python();
134 let helper = match ensure_helper_script() {
135 Ok(path) => path,
136 Err(err) => {
137 return BenchmarkResult {
138 name: name.to_string(),
139 success: false,
140 execution_time_ms: start.elapsed().as_millis(),
141 metrics: serde_json::json!({}),
142 errors: vec![format!("Unable to create benchmark helper script: {}", err)],
143 };
144 }
145 };
146
147 let mut command = Command::new(&python_bin);
148 command.arg(helper).arg(name).arg("--json");
149 if env::var_os("PYTHONHASHSEED").is_none() {
150 command.env("PYTHONHASHSEED", "0");
151 }
152 let output = command.output();
153
154 let parse_payload_result = |parsed: serde_json::Value| BenchmarkResult {
155 name: parsed
156 .get("name")
157 .and_then(|v| v.as_str())
158 .unwrap_or(name)
159 .to_string(),
160 success: parsed
161 .get("success")
162 .and_then(|v| v.as_bool())
163 .unwrap_or(false),
164 execution_time_ms: parsed
165 .get("execution_time_ms")
166 .and_then(|v| v.as_u64())
167 .unwrap_or(start.elapsed().as_millis() as u64) as u128,
168 metrics: parsed.get("metrics").cloned().unwrap_or_default(),
169 errors: parsed
170 .get("errors")
171 .and_then(|v| v.as_array())
172 .map(|arr| {
173 arr.iter()
174 .filter_map(|e| e.as_str().map(|s| s.to_string()))
175 .collect()
176 })
177 .unwrap_or_default(),
178 };
179
180 match output {
181 Ok(out) => {
182 let stdout_str = String::from_utf8_lossy(&out.stdout);
183 if let Some(parsed) = extract_helper_payload(&stdout_str) {
184 let mut result = parse_payload_result(parsed);
185 if !result.success {
186 let stderr = String::from_utf8_lossy(&out.stderr);
187 let stderr_trim = stderr.trim();
188 if !stderr_trim.is_empty() {
189 let stderr_tail = stderr_trim
190 .lines()
191 .rev()
192 .take(24)
193 .collect::<Vec<_>>()
194 .into_iter()
195 .rev()
196 .collect::<Vec<_>>()
197 .join(" | ");
198 if !stderr_tail.is_empty() {
199 result
200 .errors
201 .push(format!("Helper stderr tail: {}", stderr_tail));
202 }
203 }
204 }
205 return result;
206 }
207
208 if out.status.success() {
209 let stderr = String::from_utf8_lossy(&out.stderr).trim().to_string();
210 let parse_error = if stderr.is_empty() {
211 format!("Failed to parse helper JSON output using {}", python_bin)
212 } else {
213 format!(
214 "Failed to parse helper JSON output using {} (stderr: {})",
215 python_bin, stderr
216 )
217 };
218 return BenchmarkResult {
219 name: name.to_string(),
220 success: false,
221 execution_time_ms: start.elapsed().as_millis(),
222 metrics: serde_json::json!({}),
223 errors: vec![parse_error],
224 };
225 }
226
227 BenchmarkResult {
228 name: name.to_string(),
229 success: false,
230 execution_time_ms: start.elapsed().as_millis(),
231 metrics: serde_json::json!({}),
232 errors: {
233 let stderr = String::from_utf8_lossy(&out.stderr).trim().to_string();
234 let stdout = String::from_utf8_lossy(&out.stdout).trim().to_string();
235 let detail = if !stderr.is_empty() {
236 stderr
237 } else if !stdout.is_empty() {
238 stdout
239 } else {
240 "no output".to_string()
241 };
242 vec![format!(
243 "Helper failed with status {} using {}: {}",
244 out.status, python_bin, detail
245 )]
246 },
247 }
248 }
249 Err(err) => BenchmarkResult {
250 name: name.to_string(),
251 success: false,
252 execution_time_ms: start.elapsed().as_millis(),
253 metrics: serde_json::json!({}),
254 errors: vec![format!("Failed to invoke {}: {}", python_bin, err)],
255 },
256 }
257}
258
259pub fn run_gpu_capability_benchmark() -> BenchmarkResult {
260 run_python_benchmark("gpu-info")
261}
262pub fn run_memory_bandwidth_benchmark() -> BenchmarkResult {
263 run_python_benchmark("memory-bandwidth")
264}
265pub fn run_tensor_core_benchmark() -> BenchmarkResult {
266 run_python_benchmark("tensor-core")
267}
268pub fn run_gemm_benchmark() -> BenchmarkResult {
269 run_python_benchmark("gemm")
270}
271pub fn run_pytorch_benchmark() -> BenchmarkResult {
272 run_python_benchmark("pytorch")
273}
274pub fn run_llama_cpp_benchmark() -> BenchmarkResult {
275 run_python_benchmark("llama-cpp")
276}
277pub fn run_flash_attention_benchmark() -> BenchmarkResult {
278 run_python_benchmark("flash-attention")
279}
280pub fn run_vllm_benchmark() -> BenchmarkResult {
281 run_python_benchmark("vllm")
282}
283pub fn run_deepspeed_benchmark() -> BenchmarkResult {
284 run_python_benchmark("deepspeed")
285}
286pub fn run_megatron_benchmark() -> BenchmarkResult {
287 run_python_benchmark("megatron")
288}
289pub fn run_onnx_benchmark() -> BenchmarkResult {
290 run_python_benchmark("onnx")
291}
292
293const PY_HELPER: &str = r#"
297import argparse
298import json
299import os
300import shutil
301import subprocess
302import sys
303import time
304from glob import glob
305
306_GPU_RUNTIME_CACHE = None
307_DEFAULT_TINY_SAFETENSORS_MODEL = "HuggingFaceTB/SmolLM2-135M-Instruct"
308
309
310def _set_amd_gpu_id_table_env():
311 if os.environ.get("AMDGPU_ASIC_ID_TABLE_PATH"):
312 return
313 candidate = "/usr/share/libdrm/amdgpu.ids"
314 if os.path.isfile(candidate):
315 os.environ["AMDGPU_ASIC_ID_TABLE_PATH"] = candidate
316 os.environ.setdefault("AMDGPU_ASIC_ID_TABLE_PATHS", "/usr/share/libdrm")
317
318
319def _env_or_default(name, default):
320 value = os.environ.get(name)
321 if value is None:
322 return default
323 value = value.strip()
324 return value if value else default
325
326
327def _resolve_gguf_model_path(raw):
328 raw = (raw or "").strip()
329 if not raw:
330 return ""
331 expanded = os.path.expanduser(raw)
332 if os.path.isfile(expanded):
333 return expanded
334 if os.path.isdir(expanded):
335 for entry in sorted(os.listdir(expanded)):
336 if entry.lower().endswith(".gguf"):
337 return os.path.join(expanded, entry)
338 return ""
339
340
341def _find_llama_cpp_binary():
342 candidates = [
343 os.path.expanduser("~/.mlstack/components/llama-cpp/bin/llama-bench"),
344 shutil.which("llama-bench"),
345 ]
346 for candidate in candidates:
347 if candidate and os.path.isfile(candidate) and os.access(candidate, os.X_OK):
348 return candidate
349 return ""
350
351
352def _find_gguf_model():
353 candidates = []
354 candidates.extend(sorted(glob(os.path.expanduser("~/.mlstack/models/*.gguf"))))
355 candidates.extend(sorted(glob(os.path.expanduser("~/.cache/huggingface/hub/*/snapshots/*/*.gguf"))))
356 candidates.extend(sorted(glob(os.path.expanduser("~/.cache/**/*.gguf"), recursive=True)))
357 for candidate in candidates:
358 if os.path.isfile(candidate):
359 return candidate
360 return ""
361
362
363def _detect_llama_cpp_gpus():
364 gpu_names = []
365 try:
366 proc = subprocess.run(
367 ["rocm-smi", "--showproductname"],
368 check=False,
369 capture_output=True,
370 text=True,
371 )
372 output = (proc.stdout or "") + "\n" + (proc.stderr or "")
373 for line in output.splitlines():
374 line = line.strip()
375 if not line:
376 continue
377 lowered = line.lower()
378 if "gfx" in lowered or "amd" in lowered or "radeon" in lowered or "instinct" in lowered:
379 gpu_names.append(line)
380 except Exception:
381 pass
382 if gpu_names:
383 return gpu_names
384
385 try:
386 proc = subprocess.run(
387 "rocminfo | grep gfx",
388 shell=True,
389 check=False,
390 capture_output=True,
391 text=True,
392 )
393 output = (proc.stdout or "") + "\n" + (proc.stderr or "")
394 for line in output.splitlines():
395 line = line.strip()
396 if "gfx" in line.lower():
397 gpu_names.append(line)
398 except Exception:
399 pass
400 return gpu_names
401
402
403def _llama_cpp():
404 start = time.perf_counter()
405 bench = _find_llama_cpp_binary()
406 if not bench:
407 return False, {
408 "name": "llama-cpp",
409 "success": False,
410 "execution_time_ms": 0,
411 "metrics": {},
412 "errors": ["llama-cpp not installed"],
413 }, ["llama-cpp not installed"]
414
415 model = _find_gguf_model()
416 if not model:
417 elapsed = int((time.perf_counter() - start) * 1000)
418 return False, {
419 "name": "llama-cpp",
420 "success": False,
421 "execution_time_ms": elapsed,
422 "metrics": {},
423 "errors": ["no GGUF model found"],
424 }, ["no GGUF model found"]
425
426 gpus = _detect_llama_cpp_gpus()
427 if not gpus:
428 elapsed = int((time.perf_counter() - start) * 1000)
429 return False, {
430 "name": "llama-cpp",
431 "success": False,
432 "execution_time_ms": elapsed,
433 "metrics": {},
434 "errors": ["no ROCm GPUs detected"],
435 }, ["no ROCm GPUs detected"]
436
437 metrics = {}
438 errors = []
439 for gpu_idx, _ in enumerate(gpus):
440 env = os.environ.copy()
441 env["ROCM_VISIBLE_DEVICES"] = str(gpu_idx)
442 try:
443 proc = subprocess.run(
444 [bench, "-m", model, "-p", "512", "-p", "2048", "-p", "8192", "-n", "128", "-o", "json", "-r", "3"],
445 check=False,
446 capture_output=True,
447 text=True,
448 env=env,
449 )
450 raw = (proc.stdout or "").strip()
451 parsed = json.loads(raw) if raw else []
452 if isinstance(parsed, dict):
453 parsed = [parsed]
454 for entry in parsed:
455 if not isinstance(entry, dict):
456 continue
457 n_prompt = int(entry.get("n_prompt", 0) or 0)
458 n_gen = int(entry.get("n_gen", 0) or 0)
459 avg_ts = entry.get("avg_ts")
460 if avg_ts is None:
461 avg_ts = entry.get("avg_tps")
462 stddev_ts = entry.get("stddev_ts", 0.0)
463 prefix = "prefill" if n_prompt > 0 and n_gen == 0 else "decode" if n_prompt == 0 and n_gen > 0 else "other"
464 if prefix == "other" or avg_ts is None:
465 continue
466 context = n_prompt if prefix == "prefill" else n_gen
467 metrics[f"{prefix}_{context}_tps_gpu{gpu_idx}"] = float(avg_ts)
468 metrics[f"{prefix}_{context}_stddev_tps_gpu{gpu_idx}"] = float(stddev_ts or 0.0)
469 except Exception as exc:
470 errors.append(str(exc))
471
472 elapsed = int((time.perf_counter() - start) * 1000)
473 success = bool(metrics) and not errors
474 result = {
475 "name": "llama-cpp",
476 "success": success,
477 "execution_time_ms": elapsed,
478 "metrics": metrics,
479 "errors": errors,
480 }
481 print("---BENCHMARK_RESULTS_START---")
482 print(json.dumps(result, indent=2, sort_keys=True))
483 return success, result, errors
484
485
486def _resolve_vllm_model_candidates():
487 safetensors_model = _env_or_default(
488 "MLSTACK_BENCH_VLLM_SAFETENSORS_MODEL",
489 _DEFAULT_TINY_SAFETENSORS_MODEL,
490 )
491 gguf_hint = os.environ.get("MLSTACK_BENCH_VLLM_GGUF_MODEL_PATH", "")
492 gguf_tokenizer = _env_or_default("MLSTACK_BENCH_VLLM_GGUF_TOKENIZER", safetensors_model)
493 gguf_path = _resolve_gguf_model_path(gguf_hint)
494
495 candidates = []
496 if gguf_path:
497 candidates.append({
498 "format": "gguf",
499 "model": gguf_path,
500 "tokenizer": gguf_tokenizer,
501 })
502 candidates.append({
503 "format": "safetensors",
504 "model": safetensors_model,
505 })
506 return candidates
507
508
509def _probe_gpu_runtime():
510 global _GPU_RUNTIME_CACHE
511 if _GPU_RUNTIME_CACHE is not None:
512 return _GPU_RUNTIME_CACHE
513
514 _set_amd_gpu_id_table_env()
515 torch = _load_torch()
516 if isinstance(torch, tuple):
517 _, errors = torch
518 reason = "; ".join(errors) if errors else "unable to import torch"
519 _GPU_RUNTIME_CACHE = (False, reason, {"available": False})
520 return _GPU_RUNTIME_CACHE
521
522 if not torch.cuda.is_available():
523 _GPU_RUNTIME_CACHE = (False, "torch.cuda is not available", {"available": False})
524 return _GPU_RUNTIME_CACHE
525
526 arch = "unknown"
527 try:
528 props = torch.cuda.get_device_properties(torch.device("cuda:0"))
529 arch = getattr(props, "gcnArchName", "unknown")
530 except Exception:
531 pass
532
533 probe_code = (
534 "import torch;"
535 "x=torch.ones(1, device='cuda', dtype=torch.float32);"
536 "y=x+1;"
537 "torch.cuda.synchronize();"
538 "print(float(y.item()))"
539 )
540 try:
541 probe = subprocess.run(
542 [sys.executable, "-c", probe_code],
543 check=False,
544 capture_output=True,
545 text=True,
546 timeout=20,
547 )
548 if probe.returncode == 0:
549 _GPU_RUNTIME_CACHE = (True, "", {"available": True, "gcn_arch": arch})
550 return _GPU_RUNTIME_CACHE
551
552 detail = (probe.stderr or probe.stdout or "").strip()
553 if not detail:
554 detail = f"exit code {probe.returncode}"
555 if probe.returncode < 0:
556 detail = f"signal {-probe.returncode}: {detail}"
557 reason = f"GPU runtime probe failed on {arch}: {detail}"
558 _GPU_RUNTIME_CACHE = (False, reason, {"available": True, "gcn_arch": arch})
559 return _GPU_RUNTIME_CACHE
560 except subprocess.TimeoutExpired:
561 reason = f"GPU runtime probe timed out on {arch}"
562 _GPU_RUNTIME_CACHE = (False, reason, {"available": True, "gcn_arch": arch})
563 return _GPU_RUNTIME_CACHE
564 except Exception as exc:
565 reason = f"GPU runtime probe error on {arch}: {exc}"
566 _GPU_RUNTIME_CACHE = (False, reason, {"available": True, "gcn_arch": arch})
567 return _GPU_RUNTIME_CACHE
568
569
570def _degraded_metrics(name, reason, extra=None):
571 payload = {
572 "mode": "degraded",
573 "component": name,
574 "gpu_runtime_stable": False,
575 "reason": reason,
576 }
577 if isinstance(extra, dict):
578 payload.update(extra)
579 return payload
580
581
582def _load_torch():
583 try:
584 import torch # noqa: F401
585 return torch
586 except Exception as exc: # pragma: no cover
587 return None, [f"Unable to import torch: {exc}"]
588
589
590def _parse_visible_gpu_indices():
591 for key in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
592 raw = os.environ.get(key, "").strip()
593 if not raw:
594 continue
595
596 values = []
597 valid = True
598 for token in raw.split(","):
599 token = token.strip()
600 if not token:
601 continue
602 if not token.isdigit():
603 valid = False
604 break
605 values.append(int(token))
606
607 if valid and values:
608 return values
609 return []
610
611
612def _coerce_smi_value(entry, key):
613 raw = entry.get(key)
614 if isinstance(raw, dict):
615 value = raw.get("value")
616 if value is None:
617 return None
618 return value
619 return raw
620
621
622def _discrete_hint_from_name(name):
623 if not name:
624 return None
625
626 lowered = str(name).lower()
627 if any(token in lowered for token in ("radeon rx", " rx ", "radeon pro", "instinct", "firepro")):
628 return True
629 if "ryzen" in lowered:
630 return False
631 if "radeon graphics" in lowered and "rx" not in lowered and "pro" not in lowered:
632 return False
633 if "integrated" in lowered or "igpu" in lowered or "apu" in lowered:
634 return False
635 return None
636
637
638def _is_integrated_name(name):
639 hint = _discrete_hint_from_name(name)
640 if hint is not None:
641 return not hint
642 lowered = str(name or "").lower()
643 for token in (
644 "raphael",
645 "phoenix",
646 "rembrandt",
647 "cezanne",
648 "barcelo",
649 "yellow carp",
650 "green sardine",
651 "pink sardine",
652 "hawk point",
653 ):
654 if token in lowered:
655 return True
656 return False
657
658
659def _gpu_info():
660 torch = _load_torch()
661 if isinstance(torch, tuple):
662 _, errors = torch
663 return False, {}, errors
664 if not torch.cuda.is_available():
665 return False, {}, ["torch.cuda is not available"]
666
667 visible_indices = _parse_visible_gpu_indices()
668 smi_available = shutil.which("rocm-smi") is not None
669 devices = []
670 for idx in range(torch.cuda.device_count()):
671 global_idx = visible_indices[idx] if idx < len(visible_indices) else idx
672 device = torch.device(f"cuda:{idx}")
673 props = torch.cuda.get_device_properties(device)
674 info = {
675 "index": global_idx,
676 "local_index": idx,
677 "gpu_model": getattr(props, "name", "Unknown"),
678 "gcn_arch": getattr(props, "gcnArchName", "unknown"),
679 "vram_gb": round(getattr(props, "total_memory", 0) / 1e9, 2),
680 "compute_units": getattr(props, "multi_processor_count", 0),
681 "max_clock_mhz": getattr(props, "max_frequency", 0) / 1e6,
682 "tensor_cores": bool(getattr(props, "gcnArchName", "").startswith("gfx11")),
683 }
684
685 smi_discrete_hint = None
686 if smi_available:
687 try:
688 smi = subprocess.run(
689 [
690 "rocm-smi",
691 "--showproductname",
692 "--showtemp",
693 "--showpower",
694 "--showuse",
695 "--showclocks",
696 "--json",
697 "-d",
698 str(global_idx),
699 ],
700 check=False,
701 capture_output=True,
702 text=True,
703 )
704 if smi.returncode == 0 and smi.stdout.strip().startswith("{"):
705 data = json.loads(smi.stdout)
706 first_key = next(iter(data.keys()))
707 entry = data.get(first_key, {})
708
709 smi_name = (
710 _coerce_smi_value(entry, "Card Series")
711 or _coerce_smi_value(entry, "Card series")
712 or _coerce_smi_value(entry, "Card Model")
713 or _coerce_smi_value(entry, "Card model")
714 )
715 if smi_name:
716 info["gpu_model"] = str(smi_name).strip()
717 smi_discrete_hint = _discrete_hint_from_name(smi_name)
718
719 temp = entry.get("Temperature (Sensor die)", {}).get("value")
720 power = entry.get("Average Graphics Package Power", {}).get("value")
721 if temp is not None:
722 info["temperature_c"] = float(temp)
723 if power is not None:
724 info["power_watts"] = float(power)
725
726 # capture additional metrics requested
727 util = entry.get("GPU use (%)", {}).get("value")
728 if util is not None:
729 info["utilization_percent"] = float(util)
730
731 mem_util = entry.get("GPU memory use (%)", {}).get("value")
732 if mem_util is not None:
733 info["memory_percent"] = float(mem_util)
734
735 sclk = entry.get("GFX Clock (MHz)", {}).get("value")
736 if sclk is not None:
737 info["sclk_mhz"] = int(float(sclk))
738
739 mclk = entry.get("Memory Clock (MHz)", {}).get("value")
740 if mclk is not None:
741 info["mclk_mhz"] = int(float(mclk))
742 except Exception:
743 pass
744
745 if smi_discrete_hint is False:
746 continue
747 if smi_discrete_hint is None and _is_integrated_name(info.get("gpu_model")):
748 continue
749 devices.append(info)
750
751 probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
752 metrics = {"gpus": devices}
753 metrics.update({"gpu_runtime_stable": probe_ok})
754 if isinstance(probe_meta, dict):
755 metrics.update({f"probe_{k}": v for k, v in probe_meta.items()})
756 if probe_ok:
757 return True, metrics, []
758 return True, metrics, [probe_reason]
759
760
761def _time_fn(fn, warmup=1, repeat=5):
762 import torch
763
764 torch.cuda.synchronize()
765 for _ in range(warmup):
766 fn()
767 times = []
768 for _ in range(repeat):
769 torch.cuda.synchronize()
770 start = time.perf_counter()
771 fn()
772 torch.cuda.synchronize()
773 times.append(time.perf_counter() - start)
774 return sum(times) / len(times)
775
776
777def _memory_bandwidth():
778 torch = _load_torch()
779 if isinstance(torch, tuple):
780 _, errors = torch
781 return False, {}, errors
782 if not torch.cuda.is_available():
783 return False, {}, ["torch.cuda not available"]
784 probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
785 if not probe_ok:
786 extra = {
787 "hbm_peak_gb_s": 0.0,
788 "system_peak_gb_s": 0.0,
789 "hbm_ratio": 0.0,
790 "hbm_samples_gbps": [],
791 "system_samples_gbps": [],
792 }
793 if isinstance(probe_meta, dict):
794 extra.update(probe_meta)
795 return True, _degraded_metrics("memory-bandwidth", probe_reason, extra), [probe_reason]
796
797 device = torch.device("cuda:0")
798 sizes_mb = [64, 128, 256, 512]
799 hbm_samples = []
800 system_samples = []
801
802 for size_mb in sizes_mb:
803 numel = size_mb * 1024 * 1024 // 4
804 a = torch.empty(numel, device=device, dtype=torch.float32)
805 def op():
806 a.add_(1.0)
807 t = _time_fn(op, warmup=1, repeat=4)
808 bytes_moved = a.numel() * a.element_size() * 2 # read + write
809 hbm_samples.append(bytes_moved / t / 1e9)
810
811 host = torch.empty_like(a.cpu())
812 def h2d():
813 _ = host.to(device, non_blocking=True)
814 t_h2d = _time_fn(h2d, warmup=1, repeat=3)
815 system_samples.append(bytes_moved / t_h2d / 1e9 if t_h2d > 0 else 0)
816
817 hbm_peak = max(hbm_samples) if hbm_samples else 0
818 sys_peak = max(system_samples) if system_samples else 0
819 hbm_ratio = hbm_peak / sys_peak if sys_peak > 0 else 0
820
821 return True, {
822 "hbm_peak_gb_s": round(hbm_peak, 2),
823 "system_peak_gb_s": round(sys_peak, 2),
824 "hbm_ratio": round(hbm_ratio, 2),
825 "hbm_samples_gbps": [round(x, 2) for x in hbm_samples],
826 "system_samples_gbps": [round(x, 2) for x in system_samples],
827 }, []
828
829
830def _tensor_core():
831 torch = _load_torch()
832 if isinstance(torch, tuple):
833 _, errors = torch
834 return False, {}, errors
835 if not torch.cuda.is_available():
836 return False, {}, ["torch.cuda not available"]
837 probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
838 if not probe_ok:
839 extra = {
840 "fp16_tflops": 0.0,
841 "bf16_tflops": 0.0,
842 "tf32_tflops": 0.0,
843 "fp32_tflops": 0.0,
844 "fp16_samples_tflops": [],
845 }
846 if isinstance(probe_meta, dict):
847 extra.update(probe_meta)
848 return True, _degraded_metrics("tensor-core", probe_reason, extra), [probe_reason]
849
850 device = torch.device("cuda:0")
851 sizes = [512, 1024, 2048]
852
853 results = {}
854
855 # FP16
856 fp16_samples = []
857 for n in sizes:
858 a = torch.randn((n, n), device=device, dtype=torch.float16)
859 b = torch.randn((n, n), device=device, dtype=torch.float16)
860 t = _time_fn(lambda: torch.matmul(a, b), warmup=2, repeat=3)
861 flops = 2 * n * n * n
862 fp16_samples.append(flops / t / 1e12)
863 results["fp16_tflops"] = round(max(fp16_samples), 2)
864 results["fp16_samples"] = [round(x, 2) for x in fp16_samples]
865
866 # BF16
867 bf16_samples = []
868 try:
869 for n in sizes:
870 a = torch.randn((n, n), device=device, dtype=torch.bfloat16)
871 b = torch.randn((n, n), device=device, dtype=torch.bfloat16)
872 t = _time_fn(lambda: torch.matmul(a, b), warmup=2, repeat=3)
873 flops = 2 * n * n * n
874 bf16_samples.append(flops / t / 1e12)
875 results["bf16_tflops"] = round(max(bf16_samples), 2)
876 except Exception:
877 results["bf16_tflops"] = 0.0
878
879 # FP32
880 fp32_samples = []
881 for n in sizes:
882 a = torch.randn((n, n), device=device, dtype=torch.float32)
883 b = torch.randn((n, n), device=device, dtype=torch.float32)
884 t = _time_fn(lambda: torch.matmul(a, b), warmup=2, repeat=3)
885 flops = 2 * n * n * n
886 fp32_samples.append(flops / t / 1e12)
887 results["fp32_tflops"] = round(max(fp32_samples), 2)
888
889 # TF32 is NVIDIA specific, on AMD we'll report 0 or skip
890 results["tf32_tflops"] = 0.0
891
892 return True, {
893 "fp16_tflops": results["fp16_tflops"],
894 "bf16_tflops": results["bf16_tflops"],
895 "tf32_tflops": results["tf32_tflops"],
896 "fp32_tflops": results["fp32_tflops"],
897 "fp16_samples_tflops": results["fp16_samples"],
898 }, []
899
900
901def _gemm():
902 torch = _load_torch()
903 if isinstance(torch, tuple):
904 _, errors = torch
905 return False, {}, errors
906 if not torch.cuda.is_available():
907 return False, {}, ["torch.cuda not available"]
908 probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
909 if not probe_ok:
910 extra = {
911 "fp16_peak_gflops": 0.0,
912 "fp16_samples_gflops": [],
913 }
914 if isinstance(probe_meta, dict):
915 extra.update(probe_meta)
916 return True, _degraded_metrics("gemm", probe_reason, extra), [probe_reason]
917
918 device = torch.device("cuda:0")
919 shapes = [(1024, 1024, 1024), (1536, 1536, 1536), (2048, 2048, 2048)]
920 samples = []
921 for m, k, n in shapes:
922 a = torch.randn((m, k), device=device, dtype=torch.float16)
923 b = torch.randn((k, n), device=device, dtype=torch.float16)
924 def op():
925 return torch.matmul(a, b)
926 t = _time_fn(op, warmup=1, repeat=3)
927 flops = 2 * m * n * k
928 samples.append(flops / t / 1e9)
929 peak = max(samples) if samples else 0
930 return True, {
931 "fp16_peak_gflops": round(peak, 2),
932 "fp16_samples_gflops": [round(x, 2) for x in samples],
933 }, []
934
935
936def _pytorch():
937 torch = _load_torch()
938 if isinstance(torch, tuple):
939 _, errors = torch
940 return False, {}, errors
941 if not torch.cuda.is_available():
942 return False, {}, ["torch.cuda not available"]
943 probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
944 if not probe_ok:
945 extra = {
946 "gemm_gflops": 0.0,
947 "convolution_gflops": 0.0,
948 "gemm_samples_gflops": [],
949 "conv_samples_gflops": [],
950 "autograd_overhead_percent": 0.0,
951 }
952 if isinstance(probe_meta, dict):
953 extra.update(probe_meta)
954 return True, _degraded_metrics("pytorch", probe_reason, extra), [probe_reason]
955
956 device = torch.device("cuda:0")
957 m = k = n = 1024
958 a = torch.randn((m, k), device=device, dtype=torch.float32, requires_grad=True)
959 b = torch.randn((k, n), device=device, dtype=torch.float32, requires_grad=True)
960 gemm_samples = []
961 for _ in range(3):
962 gemm_t = _time_fn(lambda: torch.matmul(a, b), warmup=1, repeat=1)
963 gemm_samples.append(2 * m * n * k / gemm_t / 1e9)
964 gemm_peak = max(gemm_samples) if gemm_samples else 0
965
966 conv = torch.nn.Conv2d(64, 64, kernel_size=3, padding=1).to(device)
967 x = torch.randn((32, 64, 64, 64), device=device)
968 conv_samples = []
969 for _ in range(3):
970 conv_t = _time_fn(lambda: conv(x), warmup=1, repeat=1)
971 conv_ops = 2 * 32 * 64 * 64 * 64 * 64 * 3 * 3
972 conv_samples.append(conv_ops / conv_t / 1e9)
973 conv_peak = max(conv_samples) if conv_samples else 0
974
975 y = a @ b
976 y.sum().backward()
977 torch.cuda.synchronize()
978 start = time.perf_counter()
979 y = a @ b
980 loss = y.sum()
981 loss.backward()
982 torch.cuda.synchronize()
983 full = time.perf_counter() - start
984
985 start = time.perf_counter()
986 _ = a @ b
987 torch.cuda.synchronize()
988 fwd = time.perf_counter() - start
989 overhead = (full - fwd) / full * 100 if full > 0 else 0
990
991 return True, {
992 "gemm_gflops": round(gemm_peak, 2),
993 "convolution_gflops": round(conv_peak, 2),
994 "gemm_samples_gflops": [round(x, 2) for x in gemm_samples],
995 "conv_samples_gflops": [round(x, 2) for x in conv_samples],
996 "autograd_overhead_percent": round(overhead, 2),
997 }, []
998
999
1000def _flash_attention():
1001 torch = _load_torch()
1002 if isinstance(torch, tuple):
1003 _, errors = torch
1004 return False, {}, errors
1005 if not torch.cuda.is_available():
1006 return False, {}, ["torch.cuda not available"]
1007 probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
1008 if not probe_ok:
1009 extra = {
1010 "standard_attention_speed": 0.0,
1011 "flash_attention_speed": 0.0,
1012 "standard_samples_tok_s": [],
1013 "flash_samples_tok_s": [],
1014 "speedup": 0.0,
1015 "memory_savings_gb": 0.0,
1016 }
1017 if isinstance(probe_meta, dict):
1018 extra.update(probe_meta)
1019 return True, _degraded_metrics("flash-attention", probe_reason, extra), [probe_reason]
1020
1021 try:
1022 from torch.backends.cuda import sdp_kernel
1023 except Exception as exc: # pragma: no cover
1024 return False, {}, [f"flash attention control unavailable: {exc}"]
1025
1026 device = torch.device("cuda:0")
1027 seqs = [128, 256, 512, 1024]
1028 bsz, heads, dim = 8, 8, 64
1029 flash_samples = []
1030 std_samples = []
1031 mem_std = 0
1032 mem_flash = 0
1033
1034 for seqlen in seqs:
1035 q = torch.randn((bsz, heads, seqlen, dim), device=device, dtype=torch.float16)
1036 k = torch.randn_like(q)
1037 v = torch.randn_like(q)
1038
1039 def run_flash():
1040 with sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
1041 return torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=False)
1042
1043 def run_standard():
1044 with sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False):
1045 return torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=False)
1046
1047 torch.cuda.reset_peak_memory_stats()
1048 _ = run_standard()
1049 mem_std = max(mem_std, torch.cuda.max_memory_allocated() / 1e9)
1050
1051 torch.cuda.reset_peak_memory_stats()
1052 _ = run_flash()
1053 mem_flash = max(mem_flash, torch.cuda.max_memory_allocated() / 1e9)
1054
1055 flash_t = _time_fn(run_flash, warmup=1, repeat=2)
1056 std_t = _time_fn(run_standard, warmup=1, repeat=2)
1057
1058 tokens = bsz * seqlen
1059 flash_samples.append(tokens / flash_t)
1060 std_samples.append(tokens / std_t)
1061
1062 flash_peak = max(flash_samples) if flash_samples else 0
1063 std_peak = max(std_samples) if std_samples else 0
1064
1065 return True, {
1066 "standard_attention_speed": round(std_peak, 2),
1067 "flash_attention_speed": round(flash_peak, 2),
1068 "standard_samples_tok_s": [round(x, 2) for x in std_samples],
1069 "flash_samples_tok_s": [round(x, 2) for x in flash_samples],
1070 "speedup": round(flash_peak / std_peak, 3) if std_peak > 0 else 0,
1071 "memory_savings_gb": round(max(0, mem_std - mem_flash), 4),
1072 }, []
1073
1074
1075def _vllm():
1076 # Some shell setups can propagate an empty/whitespace-only target device.
1077 # vLLM rejects empty device strings, so normalize aggressively before import.
1078 target_device = _env_or_default("VLLM_TARGET_DEVICE", "rocm")
1079 if target_device not in {"rocm", "cuda", "cpu"}:
1080 target_device = "rocm"
1081 os.environ["VLLM_TARGET_DEVICE"] = target_device
1082
1083 def _ensure_writable_triton_cache_env():
1084 import tempfile
1085
1086 def _writable_dir(path):
1087 if not path:
1088 return False
1089 try:
1090 os.makedirs(path, exist_ok=True)
1091 probe = os.path.join(path, f".mlstack_probe_{os.getpid()}")
1092 with open(probe, "w", encoding="utf-8") as fp:
1093 fp.write("ok")
1094 os.remove(probe)
1095 return True
1096 except Exception:
1097 return False
1098
1099 triton_cache = (os.environ.get("TRITON_CACHE_DIR") or "").strip()
1100 if triton_cache and _writable_dir(triton_cache):
1101 triton_home = (os.environ.get("TRITON_HOME") or os.path.dirname(triton_cache)).strip()
1102 os.environ.setdefault("TRITON_HOME", triton_home or os.path.expanduser("~/.cache/mlstack/triton"))
1103 os.environ.setdefault("TRITON_DUMP_DIR", os.path.join(os.environ["TRITON_HOME"], "dump"))
1104 os.environ.setdefault("TRITON_OVERRIDE_DIR", os.path.join(os.environ["TRITON_HOME"], "override"))
1105 _writable_dir(os.environ.get("TRITON_DUMP_DIR", ""))
1106 _writable_dir(os.environ.get("TRITON_OVERRIDE_DIR", ""))
1107 return
1108
1109 home = os.path.expanduser("~")
1110 candidates = [
1111 os.environ.get("MLSTACK_TRITON_HOME", "").strip(),
1112 os.environ.get("TRITON_HOME", "").strip(),
1113 os.path.join(home, ".cache", "mlstack", "triton"),
1114 os.path.join(tempfile.gettempdir(), f"mlstack-triton-{os.getuid()}"),
1115 ]
1116 for root in candidates:
1117 if not root:
1118 continue
1119 cache_dir = os.path.join(root, "cache")
1120 dump_dir = os.path.join(root, "dump")
1121 override_dir = os.path.join(root, "override")
1122 if _writable_dir(cache_dir) and _writable_dir(dump_dir) and _writable_dir(override_dir):
1123 os.environ["MLSTACK_TRITON_HOME"] = root
1124 os.environ["TRITON_HOME"] = root
1125 os.environ["TRITON_CACHE_DIR"] = cache_dir
1126 os.environ["TRITON_DUMP_DIR"] = dump_dir
1127 os.environ["TRITON_OVERRIDE_DIR"] = override_dir
1128 return
1129
1130 _ensure_writable_triton_cache_env()
1131
1132 def _normalize_visible_devices():
1133 visible = os.environ.get("HIP_VISIBLE_DEVICES") or os.environ.get("CUDA_VISIBLE_DEVICES") or ""
1134 candidates = []
1135 seen = set()
1136 for token in visible.split(","):
1137 item = token.strip()
1138 if not item.isdigit() or item in seen:
1139 continue
1140 seen.add(item)
1141 candidates.append(item)
1142 if not candidates:
1143 candidates = ["0"]
1144
1145 normalized = ",".join(candidates)
1146 primary = candidates[0]
1147 os.environ["HIP_VISIBLE_DEVICES"] = normalized
1148 os.environ["CUDA_VISIBLE_DEVICES"] = normalized
1149 if not str(os.environ.get("PYTORCH_ROCM_DEVICE", "")).strip():
1150 os.environ["PYTORCH_ROCM_DEVICE"] = primary
1151 return normalized, primary
1152
1153 # Preserve the full discrete-GPU set. Do not collapse to a single GPU.
1154 visible_devices, primary_visible = _normalize_visible_devices()
1155 if target_device == "rocm":
1156 os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
1157
1158 candidates = _resolve_vllm_model_candidates()
1159 candidate_names = [f"{c.get('format')}:{c.get('model')}" for c in candidates]
1160 torch = _load_torch()
1161 if isinstance(torch, tuple):
1162 _, errors = torch
1163 reason = "; ".join(errors) if errors else "unable to import torch"
1164 return True, _degraded_metrics("vllm", reason, {
1165 "model": "vLLM (Unavailable)",
1166 "model_format": "unavailable",
1167 "throughput_tokens_per_sec": 0.0,
1168 "latency_ms": 0.0,
1169 "throughput_samples": [],
1170 "target_device": target_device,
1171 "candidate_models": candidate_names,
1172 }), [reason]
1173 runtime_probe_hint = ""
1174 if target_device in {"rocm", "cuda"}:
1175 try:
1176 if not torch.cuda.is_available():
1177 runtime_probe_hint = "ROCm torch reports no available HIP GPUs in benchmark environment"
1178 else:
1179 torch.cuda.get_device_properties(0)
1180 except Exception as exc:
1181 runtime_probe_hint = f"ROCm torch GPU initialization failed before vLLM load: {exc}"
1182 def _patch_amdsmi_for_vllm_arch_probe():
1183 try:
1184 import sys
1185 import types
1186
1187 if os.environ.get("MLSTACK_VLLM_DISABLE_AMDSMI_SHIM", "").strip() in {"1", "true", "TRUE"}:
1188 return
1189
1190 arch_hint = (
1191 os.environ.get("GPU_ARCH")
1192 or os.environ.get("PYTORCH_ROCM_ARCH")
1193 or os.environ.get("HSA_OVERRIDE_GFX_VERSION")
1194 or ""
1195 ).strip()
1196 if not arch_hint:
1197 return
1198 if not arch_hint.startswith("gfx") and "." in arch_hint:
1199 parts = [p for p in arch_hint.split(".") if p]
1200 if len(parts) >= 2:
1201 arch_hint = f"gfx{parts[0]}{parts[1]}"
1202
1203 shim = types.ModuleType("amdsmi")
1204
1205 class AmdSmiException(Exception):
1206 pass
1207
1208 def amdsmi_init():
1209 return None
1210
1211 def amdsmi_shut_down():
1212 return None
1213
1214 def amdsmi_get_processor_handles():
1215 return [0]
1216
1217 def amdsmi_get_gpu_asic_info(_handle):
1218 return {"target_graphics_version": arch_hint}
1219
1220 def amdsmi_topo_get_link_type(*_args, **_kwargs):
1221 return (0, 0)
1222
1223 shim.AmdSmiException = AmdSmiException
1224 shim.amdsmi_init = amdsmi_init
1225 shim.amdsmi_shut_down = amdsmi_shut_down
1226 shim.amdsmi_get_processor_handles = amdsmi_get_processor_handles
1227 shim.amdsmi_get_gpu_asic_info = amdsmi_get_gpu_asic_info
1228 shim.amdsmi_topo_get_link_type = amdsmi_topo_get_link_type
1229
1230 sys.modules["amdsmi"] = shim
1231 except Exception:
1232 # Best-effort compatibility shim.
1233 pass
1234
1235 _patch_amdsmi_for_vllm_arch_probe()
1236 try:
1237 import vllm
1238 from vllm import LLM, SamplingParams
1239 import logging
1240 import importlib
1241 # Suppress vLLM and related logging to avoid breaking JSON parsing
1242 logging.getLogger("vllm").setLevel(logging.ERROR)
1243 except Exception as exc:
1244 err = f"vLLM not available: {exc}"
1245 return True, _degraded_metrics("vllm", err, {
1246 "model": "vLLM (Unavailable)",
1247 "model_format": "unavailable",
1248 "throughput_tokens_per_sec": 0.0,
1249 "latency_ms": 0.0,
1250 "throughput_samples": [],
1251 "target_device": target_device,
1252 "candidate_models": candidate_names,
1253 }), [err]
1254 abi_errors = []
1255 for ext_name in ("vllm._C", "vllm._rocm_C"):
1256 try:
1257 importlib.import_module(ext_name)
1258 except Exception as exc:
1259 abi_errors.append(f"{ext_name} import failed: {exc}")
1260 if len(abi_errors) == 2:
1261 reason = "vLLM ROCm native extensions failed to load; " + "; ".join(abi_errors)
1262 return True, _degraded_metrics("vllm", reason, {
1263 "model": "vLLM (Unavailable)",
1264 "model_format": "unavailable",
1265 "throughput_tokens_per_sec": 0.0,
1266 "latency_ms": 0.0,
1267 "throughput_samples": [],
1268 "target_device": target_device,
1269 "candidate_models": candidate_names,
1270 }), [reason]
1271 platform_device_type = ""
1272 platform_class = "unknown"
1273 try:
1274 from vllm.platforms import current_platform
1275 platform_device_type = str(getattr(current_platform, "device_type", "") or "").strip()
1276 platform_class = current_platform.__class__.__name__
1277 except Exception:
1278 pass
1279 if not platform_device_type:
1280 platform_reason = (
1281 "vLLM platform detection returned an empty device_type (UnspecifiedPlatform). "
1282 "This commonly means amdsmi is missing, so ROCm platform detection failed. "
1283 "Install amdsmi in the benchmark/runtime Python environment and rerun."
1284 )
1285 return True, _degraded_metrics("vllm", platform_reason, {
1286 "model": "vLLM (Detected)",
1287 "model_format": "unavailable",
1288 "throughput_tokens_per_sec": 0.0,
1289 "latency_ms": 0.0,
1290 "throughput_samples": [],
1291 "target_device": target_device,
1292 "candidate_models": candidate_names,
1293 "vllm_platform_class": platform_class,
1294 "vllm_platform_device_type": platform_device_type,
1295 }), [platform_reason]
1296 probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
1297 if not probe_ok:
1298 runtime_probe_hint = probe_reason if not runtime_probe_hint else f"{runtime_probe_hint}; {probe_reason}"
1299
1300 prompts = [
1301 "The future of AI is",
1302 "ROCm performance on AMD is",
1303 "vLLM is a fast",
1304 ]
1305 max_tokens = 24
1306 attempt_errors = []
1307
1308 def _run_vllm_attempt_subprocess(llm_kwargs, prompts, max_tokens, env_overrides=None):
1309 payload = {
1310 "llm_kwargs": llm_kwargs,
1311 "prompts": prompts,
1312 "max_tokens": int(max_tokens),
1313 }
1314 runner = r"""
1315import json
1316import logging
1317import os
1318import sys
1319import time
1320import traceback
1321
1322def _emit(obj):
1323 print(json.dumps(obj))
1324
1325try:
1326 payload = json.loads(sys.stdin.read())
1327 llm_kwargs = payload.get("llm_kwargs", {})
1328 prompts = payload.get("prompts", [])
1329 max_tokens = int(payload.get("max_tokens", 24))
1330 try:
1331 import types
1332 if os.environ.get("MLSTACK_VLLM_DISABLE_AMDSMI_SHIM", "").strip() not in {"1", "true", "TRUE"}:
1333 arch_hint = (
1334 os.environ.get("GPU_ARCH")
1335 or os.environ.get("PYTORCH_ROCM_ARCH")
1336 or os.environ.get("HSA_OVERRIDE_GFX_VERSION")
1337 or ""
1338 ).strip()
1339 if arch_hint and not arch_hint.startswith("gfx") and "." in arch_hint:
1340 parts = [p for p in arch_hint.split(".") if p]
1341 if len(parts) >= 2:
1342 arch_hint = f"gfx{parts[0]}{parts[1]}"
1343 if arch_hint:
1344 shim = types.ModuleType("amdsmi")
1345 class AmdSmiException(Exception):
1346 pass
1347 def amdsmi_init():
1348 return None
1349 def amdsmi_shut_down():
1350 return None
1351 def amdsmi_get_processor_handles():
1352 return [0]
1353 def amdsmi_get_gpu_asic_info(_handle):
1354 return {"target_graphics_version": arch_hint}
1355 def amdsmi_topo_get_link_type(*_args, **_kwargs):
1356 return (0, 0)
1357 shim.AmdSmiException = AmdSmiException
1358 shim.amdsmi_init = amdsmi_init
1359 shim.amdsmi_shut_down = amdsmi_shut_down
1360 shim.amdsmi_get_processor_handles = amdsmi_get_processor_handles
1361 shim.amdsmi_get_gpu_asic_info = amdsmi_get_gpu_asic_info
1362 shim.amdsmi_topo_get_link_type = amdsmi_topo_get_link_type
1363 sys.modules["amdsmi"] = shim
1364 except Exception:
1365 pass
1366 import vllm
1367 from vllm import LLM, SamplingParams
1368 logging.getLogger("vllm").setLevel(logging.ERROR)
1369 sampling_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
1370 llm = LLM(**llm_kwargs)
1371 start = time.perf_counter()
1372 outputs = llm.generate(prompts, sampling_params)
1373 elapsed = time.perf_counter() - start
1374
1375 throughput_samples = []
1376 for prompt in prompts:
1377 sample_start = time.perf_counter()
1378 sample_out = llm.generate([prompt], sampling_params)
1379 sample_elapsed = time.perf_counter() - sample_start
1380 tokens = len(sample_out[0].outputs[0].token_ids)
1381 throughput_samples.append(tokens / sample_elapsed if sample_elapsed > 0 else 0.0)
1382
1383 total_tokens = sum(len(output.outputs[0].token_ids) for output in outputs)
1384 throughput = total_tokens / elapsed if elapsed > 0 else 0.0
1385 _emit({
1386 "ok": True,
1387 "throughput_tokens_per_sec": round(throughput, 2),
1388 "latency_ms": round((elapsed / max(len(prompts), 1)) * 1000, 2),
1389 "throughput_samples": [round(x, 2) for x in throughput_samples],
1390 })
1391except Exception as exc:
1392 tb = traceback.format_exc().strip()
1393 tail_lines = [line.strip() for line in tb.splitlines()[-12:] if line.strip()]
1394 msg = str(exc).strip() or repr(exc)
1395 if tail_lines:
1396 msg = f"{msg} :: traceback_tail: {' | '.join(tail_lines)}"
1397 _emit({"ok": False, "error": msg})
1398"""
1399 env = os.environ.copy()
1400 if isinstance(env_overrides, dict):
1401 for key, value in env_overrides.items():
1402 if value is None:
1403 env.pop(key, None)
1404 else:
1405 env[key] = str(value)
1406 try:
1407 proc = subprocess.run(
1408 [sys.executable, "-c", runner],
1409 input=json.dumps(payload),
1410 capture_output=True,
1411 text=True,
1412 timeout=240,
1413 env=env,
1414 )
1415 except Exception as exc:
1416 return {"ok": False, "error": f"subprocess launch failed: {exc}"}
1417
1418 merged = (proc.stdout or "").splitlines()
1419 if proc.stderr:
1420 merged.extend(proc.stderr.splitlines())
1421 for line in reversed(merged):
1422 text = line.strip()
1423 if not text.startswith("{") or not text.endswith("}"):
1424 continue
1425 try:
1426 parsed = json.loads(text)
1427 except Exception:
1428 continue
1429 if isinstance(parsed, dict) and "ok" in parsed:
1430 if not parsed.get("ok", False):
1431 log_tail = []
1432 for raw in merged:
1433 msg = raw.strip()
1434 if not msg:
1435 continue
1436 if msg.startswith("{") and msg.endswith("}"):
1437 continue
1438 lower = msg.lower()
1439 if any(
1440 marker in lower
1441 for marker in (
1442 "error",
1443 "exception",
1444 "traceback",
1445 "runtimeerror",
1446 "valueerror",
1447 "fatal",
1448 "enginecore",
1449 "failed",
1450 )
1451 ):
1452 log_tail.append(msg)
1453 if log_tail:
1454 tail = " | ".join(log_tail[-12:])[:3000]
1455 base = str(parsed.get("error") or "unknown vLLM subprocess failure")
1456 if tail and tail not in base:
1457 parsed["error"] = f"{base}; engine_log_tail: {tail}"
1458 return parsed
1459
1460 detail = " ".join([line.strip() for line in merged if line.strip()])[:4000]
1461 if not detail:
1462 detail = f"exit code {proc.returncode}"
1463 return {"ok": False, "error": detail}
1464
1465 for candidate in candidates:
1466 model_name = candidate.get("model")
1467 model_format = candidate.get("format", "unknown")
1468 if not model_name:
1469 continue
1470
1471 llm_kwargs = {
1472 "model": model_name,
1473 "trust_remote_code": True,
1474 "enforce_eager": True,
1475 "gpu_memory_utilization": 0.65,
1476 "disable_log_stats": True,
1477 "tensor_parallel_size": 1,
1478 "dtype": "float16",
1479 }
1480 if model_format == "gguf":
1481 tokenizer_name = candidate.get("tokenizer")
1482 if tokenizer_name:
1483 llm_kwargs["tokenizer"] = tokenizer_name
1484
1485 attempt = _run_vllm_attempt_subprocess(llm_kwargs, prompts, max_tokens)
1486 if attempt.get("ok"):
1487 return True, {
1488 "model": model_name,
1489 "model_format": model_format,
1490 "target_device": target_device,
1491 "visible_devices": os.environ.get("HIP_VISIBLE_DEVICES", primary_visible),
1492 "throughput_tokens_per_sec": float(attempt.get("throughput_tokens_per_sec", 0.0)),
1493 "latency_ms": float(attempt.get("latency_ms", 0.0)),
1494 "throughput_samples": [
1495 float(x) for x in (attempt.get("throughput_samples") or [])
1496 ],
1497 "candidate_models": candidate_names,
1498 }, []
1499 err_msg = str(attempt.get("error") or "unknown vLLM execution error")
1500 engine_core_failure = (
1501 "Engine core initialization failed" in err_msg
1502 or "Failed core proc" in err_msg
1503 or "No HIP GPUs are available" in err_msg
1504 or "torch.cuda is not available" in err_msg
1505 )
1506 # Defensive fallback for intermittent vLLM device/runtime failures.
1507 # Preserve full visible GPU list for multi-GPU hosts.
1508 if "Device string must not be empty" in err_msg or engine_core_failure:
1509 retry_profiles = [
1510 (
1511 "spawn-v1-default",
1512 {
1513 "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
1514 "VLLM_ENABLE_V1_MULTIPROCESSING": "1",
1515 "VLLM_USE_V1": "1",
1516 },
1517 ),
1518 (
1519 "fork-v1-default",
1520 {
1521 "VLLM_WORKER_MULTIPROC_METHOD": "fork",
1522 "VLLM_ENABLE_V1_MULTIPROCESSING": "1",
1523 "VLLM_USE_V1": "1",
1524 },
1525 ),
1526 (
1527 "spawn-v0-singleproc",
1528 {
1529 "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
1530 "VLLM_ENABLE_V1_MULTIPROCESSING": "0",
1531 "VLLM_USE_V1": "0",
1532 },
1533 ),
1534 (
1535 "spawn-v0-no-aiter",
1536 {
1537 "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
1538 "VLLM_ENABLE_V1_MULTIPROCESSING": "0",
1539 "VLLM_USE_V1": "0",
1540 "VLLM_ROCM_USE_AITER": "0",
1541 },
1542 ),
1543 ]
1544 retry_failures = []
1545 for profile_name, profile_overrides in retry_profiles:
1546 retry_overrides = {
1547 "HIP_VISIBLE_DEVICES": visible_devices,
1548 "CUDA_VISIBLE_DEVICES": visible_devices,
1549 "PYTORCH_ROCM_DEVICE": primary_visible,
1550 "VLLM_TARGET_DEVICE": target_device,
1551 }
1552 retry_overrides.update(profile_overrides)
1553 retry = _run_vllm_attempt_subprocess(
1554 llm_kwargs,
1555 prompts,
1556 max_tokens,
1557 env_overrides=retry_overrides,
1558 )
1559 if retry.get("ok"):
1560 return True, {
1561 "model": model_name,
1562 "model_format": model_format,
1563 "target_device": target_device,
1564 "throughput_tokens_per_sec": float(retry.get("throughput_tokens_per_sec", 0.0)),
1565 "latency_ms": float(retry.get("latency_ms", 0.0)),
1566 "throughput_samples": [float(x) for x in (retry.get("throughput_samples") or [])],
1567 "candidate_models": candidate_names,
1568 "visible_devices_after_retry": visible_devices,
1569 "fallback_mode": profile_name,
1570 }, []
1571 retry_failures.append(f"{profile_name}: {retry.get('error', 'unknown retry failure')}")
1572 if retry_failures:
1573 err_msg = f"{err_msg}; retry_profiles_failed: {' || '.join(retry_failures)}"
1574 if "Entry Not Found" in err_msg or "not found on the Hugging Face Hub" in err_msg:
1575 err_msg = (
1576 f"Model {model_name} not found. Download/copy it first or override "
1577 "MLSTACK_BENCH_VLLM_SAFETENSORS_MODEL / MLSTACK_BENCH_VLLM_GGUF_MODEL_PATH"
1578 )
1579 elif "authentication" in err_msg.lower():
1580 err_msg = "Hugging Face authentication required. Run 'huggingface-cli login' first."
1581 attempt_errors.append(f"{model_format} model {model_name} failed: {err_msg}")
1582
1583 failure_reason = "No tiny vLLM benchmark model could be loaded"
1584 if runtime_probe_hint:
1585 failure_reason = f"{failure_reason}; runtime_probe={runtime_probe_hint}"
1586 if attempt_errors:
1587 failure_reason = f"{failure_reason}; {attempt_errors[0]}"
1588
1589 return True, _degraded_metrics("vllm", failure_reason, {
1590 "model": "vLLM (Detected)",
1591 "model_format": "unavailable",
1592 "target_device": target_device,
1593 "visible_devices": os.environ.get("HIP_VISIBLE_DEVICES", primary_visible),
1594 "throughput_tokens_per_sec": 0.0,
1595 "latency_ms": 0.0,
1596 "throughput_samples": [],
1597 "candidate_models": candidate_names,
1598 }), attempt_errors or [failure_reason]
1599
1600
1601def _deepspeed():
1602 _set_amd_gpu_id_table_env()
1603 try:
1604 import deepspeed
1605 import torch
1606 import os
1607 except Exception as exc:
1608 err = f"DeepSpeed or Torch not available: {exc}"
1609 return True, _degraded_metrics("deepspeed", err, {
1610 "throughput_samples_per_sec": 0.0,
1611 "avg_latency_ms": 0.0,
1612 "stage": 1,
1613 "accelerator": "rocm",
1614 "samples": [],
1615 }), [err]
1616 probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
1617 if not probe_ok:
1618 extra = {
1619 "throughput_samples_per_sec": 0.0,
1620 "avg_latency_ms": 0.0,
1621 "stage": 1,
1622 "accelerator": "rocm",
1623 "samples": [],
1624 }
1625 if isinstance(probe_meta, dict):
1626 extra.update(probe_meta)
1627 return True, _degraded_metrics("deepspeed", probe_reason, extra), [probe_reason]
1628
1629 # Ensure ROCm environment
1630 os.environ["DS_ACCELERATOR"] = "rocm"
1631
1632 try:
1633 from deepspeed.accelerator import get_accelerator
1634
1635 accelerator = get_accelerator()
1636 device_name = accelerator.device_name()
1637 device = torch.device(device_name if ":" in device_name else f"{device_name}:0")
1638
1639 class SimpleModel(torch.nn.Module):
1640 def __init__(self):
1641 super().__init__()
1642 self.net = torch.nn.Sequential(
1643 torch.nn.Linear(1024, 1024),
1644 torch.nn.ReLU(),
1645 torch.nn.Linear(1024, 1024),
1646 )
1647
1648 def forward(self, x):
1649 return self.net(x)
1650
1651 model = SimpleModel().to(device).half()
1652 batch_size = 4
1653 input_data = torch.randn(batch_size, 1024, device=device, dtype=torch.float16)
1654
1655 ds_config = {
1656 "train_batch_size": batch_size,
1657 "train_micro_batch_size_per_gpu": batch_size,
1658 "gradient_accumulation_steps": 1,
1659 "zero_optimization": {"stage": 0},
1660 "fp16": {"enabled": True},
1661 # DeepSpeed internally computes `step % steps_per_print`; keep this non-zero.
1662 "steps_per_print": 1,
1663 "wall_clock_breakdown": False,
1664 }
1665 optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
1666
1667 try:
1668 model_engine, _, _, _ = deepspeed.initialize(
1669 config=ds_config,
1670 model=model,
1671 model_parameters=model.parameters(),
1672 optimizer=optimizer,
1673 )
1674 mode = "train"
1675 except Exception as init_exc:
1676 init_err = str(init_exc)
1677 if "fused_adam" in init_err.lower() or "cuda_runtime_api.h" in init_err:
1678 model_engine = deepspeed.init_inference(
1679 model,
1680 mp_size=1,
1681 dtype=torch.float16,
1682 replace_with_kernel_inject=False,
1683 )
1684 mode = "inference"
1685 else:
1686 raise
1687
1688 times = []
1689 if mode == "train":
1690 for _ in range(2):
1691 out = model_engine(input_data)
1692 model_engine.backward(out.sum())
1693 model_engine.step()
1694
1695 for _ in range(5):
1696 accelerator.synchronize()
1697 start = time.perf_counter()
1698 out = model_engine(input_data)
1699 model_engine.backward(out.sum())
1700 model_engine.step()
1701 accelerator.synchronize()
1702 times.append(time.perf_counter() - start)
1703 else:
1704 model_engine.eval()
1705 for _ in range(2):
1706 with torch.no_grad():
1707 _ = model_engine(input_data)
1708 accelerator.synchronize()
1709
1710 for _ in range(5):
1711 accelerator.synchronize()
1712 start = time.perf_counter()
1713 with torch.no_grad():
1714 _ = model_engine(input_data)
1715 accelerator.synchronize()
1716 times.append(time.perf_counter() - start)
1717
1718 avg_time = sum(times) / len(times)
1719 throughput = batch_size / avg_time
1720
1721 return True, {
1722 "throughput_samples_per_sec": round(throughput, 2),
1723 "avg_latency_ms": round(avg_time * 1000, 2),
1724 "stage": 0 if mode == "train" else 0,
1725 "accelerator": "rocm",
1726 "mode": mode,
1727 "samples": [round(batch_size / t, 2) for t in times],
1728 }, []
1729 except Exception as exc:
1730 err = f"DeepSpeed benchmark failed: {exc}"
1731 return True, _degraded_metrics("deepspeed", err, {
1732 "throughput_samples_per_sec": 0.0,
1733 "avg_latency_ms": 0.0,
1734 "stage": 1,
1735 "accelerator": "rocm",
1736 "samples": [],
1737 }), [err]
1738
1739
1740def _megatron():
1741 try:
1742 import torch
1743 import megatron
1744 except Exception as exc:
1745 err = f"Megatron not available: {exc}"
1746 return True, _degraded_metrics("megatron", err, {
1747 "megatron_backend": "unavailable",
1748 "megatron_throughput_samples_per_sec": 0.0,
1749 "megatron_avg_latency_ms": 0.0,
1750 "megatron_samples": [],
1751 }), [err]
1752
1753 probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
1754 if not probe_ok:
1755 extra = {
1756 "megatron_backend": "megatron-core",
1757 "megatron_throughput_samples_per_sec": 0.0,
1758 "megatron_avg_latency_ms": 0.0,
1759 "megatron_samples": [],
1760 }
1761 if isinstance(probe_meta, dict):
1762 extra.update(probe_meta)
1763 return True, _degraded_metrics("megatron", probe_reason, extra), [probe_reason]
1764
1765 try:
1766 device = torch.device("cuda:0")
1767 model = torch.nn.Sequential(
1768 torch.nn.Linear(4096, 4096),
1769 torch.nn.GELU(),
1770 torch.nn.Linear(4096, 4096),
1771 ).to(device).half()
1772
1773 input_data = torch.randn(8, 4096, device=device, dtype=torch.float16)
1774 samples = []
1775
1776 for _ in range(2):
1777 out = model(input_data)
1778 _ = out.sum()
1779 torch.cuda.synchronize()
1780
1781 for _ in range(5):
1782 torch.cuda.synchronize()
1783 start = time.perf_counter()
1784 out = model(input_data)
1785 _ = out.sum()
1786 torch.cuda.synchronize()
1787 elapsed = time.perf_counter() - start
1788 samples.append(8.0 / elapsed if elapsed > 0 else 0.0)
1789
1790 avg_tput = sum(samples) / len(samples) if samples else 0.0
1791 avg_latency_ms = (1000.0 / avg_tput * 8.0) if avg_tput > 0 else 0.0
1792 return True, {
1793 "megatron_backend": getattr(megatron, "__name__", "megatron"),
1794 "megatron_throughput_samples_per_sec": round(avg_tput, 2),
1795 "megatron_avg_latency_ms": round(avg_latency_ms, 2),
1796 "megatron_samples": [round(x, 2) for x in samples],
1797 }, []
1798 except Exception as exc:
1799 err = f"Megatron benchmark failed: {exc}"
1800 return True, _degraded_metrics("megatron", err, {
1801 "megatron_backend": "megatron-core",
1802 "megatron_throughput_samples_per_sec": 0.0,
1803 "megatron_avg_latency_ms": 0.0,
1804 "megatron_samples": [],
1805 }), [err]
1806
1807
1808def _onnx():
1809 import os
1810 try:
1811 import onnxruntime as ort
1812 except Exception as exc:
1813 err = f"onnxruntime not available: {exc}"
1814 return True, _degraded_metrics("onnx", err, {
1815 "ort_version": "unavailable",
1816 "provider": "none",
1817 "providers_available": [],
1818 "model_load_ms": 0.0,
1819 "session_create_ms": 0.0,
1820 "inference_latency_p50_ms": 0.0,
1821 "inference_latency_p95_ms": 0.0,
1822 "inference_latency_p99_ms": 0.0,
1823 "throughput_inf_per_sec": 0.0,
1824 "input_shape": [],
1825 "output_shape": [],
1826 "graph_opt_level": "unavailable",
1827 "inference_samples": [],
1828 }), [err]
1829
1830 ort_version = ort.__version__
1831 all_providers = ort.get_available_providers()
1832 providers_available = list(all_providers)
1833
1834 if "MIGraphXExecutionProvider" in all_providers:
1835 provider = "MIGraphXExecutionProvider"
1836 elif "DmlExecutionProvider" in all_providers:
1837 provider = "DmlExecutionProvider"
1838 elif "ROCMExecutionProvider" in all_providers:
1839 provider = "ROCMExecutionProvider"
1840 elif "CUDAExecutionProvider" in all_providers:
1841 provider = "CUDAExecutionProvider"
1842 else:
1843 provider = "CPUExecutionProvider"
1844
1845 try:
1846 import numpy as np
1847 except Exception:
1848 err = "numpy not available for ONNX benchmark"
1849 return True, _degraded_metrics("onnx", err, {
1850 "ort_version": ort_version,
1851 "provider": provider,
1852 "providers_available": providers_available,
1853 "model_load_ms": 0.0,
1854 "session_create_ms": 0.0,
1855 "inference_latency_p50_ms": 0.0,
1856 "inference_latency_p95_ms": 0.0,
1857 "inference_latency_p99_ms": 0.0,
1858 "throughput_inf_per_sec": 0.0,
1859 "input_shape": [],
1860 "output_shape": [],
1861 "graph_opt_level": "ORT_ENABLE_ALL",
1862 "inference_samples": [],
1863 }), [err]
1864
1865 try:
1866 import onnx
1867 from onnx import helper, TensorProto, numpy_helper
1868 except Exception:
1869 onnx = None
1870
1871 batch_size = 8
1872 hidden = 512
1873 input_shape = [batch_size, hidden]
1874
1875 def _build_test_model():
1876 X = helper.make_tensor_value_info("input", TensorProto.FLOAT, input_shape)
1877 Y = helper.make_tensor_value_info("output", TensorProto.FLOAT, input_shape)
1878 W_init = numpy_helper.from_array(
1879 np.random.randn(hidden, hidden).astype(np.float32) * 0.01, name="W"
1880 )
1881 B_init = numpy_helper.from_array(
1882 np.zeros(hidden, dtype=np.float32), name="B"
1883 )
1884 matmul = helper.make_node("MatMul", ["input", "W"], ["matmul_out"])
1885 add = helper.make_node("Add", ["matmul_out", "B"], ["add_out"])
1886 relu = helper.make_node("Relu", ["add_out"], ["output"])
1887 graph = helper.make_graph(
1888 [matmul, add, relu], "bench_graph", [X], [Y], initializer=[W_init, B_init]
1889 )
1890 model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
1891 model.ir_version = 8
1892 onnx.checker.validate = False
1893 return onnx._serialize(model)
1894
1895 def _build_quantized_model():
1896 X = helper.make_tensor_value_info("input", TensorProto.FLOAT, input_shape)
1897 Y = helper.make_tensor_value_info("output", TensorProto.FLOAT, input_shape)
1898 W_init = numpy_helper.from_array(
1899 np.random.randn(hidden, hidden).astype(np.float32) * 0.01, name="W"
1900 )
1901 B_init = numpy_helper.from_array(
1902 np.zeros(hidden, dtype=np.float32), name="B"
1903 )
1904 dq = helper.make_node(
1905 "DynamicQuantizeLinear", ["input"], ["input_quant", "input_scale", "input_zp"]
1906 )
1907 matmul_int = helper.make_node(
1908 "MatMulInteger", ["input_quant", "W"], ["matmul_int_out"]
1909 )
1910 cast_out = helper.make_node(
1911 "Cast", ["matmul_int_out"], ["matmul_fp_out"], to=TensorProto.FLOAT
1912 )
1913 add = helper.make_node("Add", ["matmul_fp_out", "B"], ["output"])
1914 graph = helper.make_graph(
1915 [dq, matmul_int, cast_out, add], "quant_bench_graph",
1916 [X], [Y], initializer=[W_init, B_init]
1917 )
1918 model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
1919 model.ir_version = 8
1920 onnx.checker.validate = False
1921 return onnx._serialize(model)
1922
1923 import tempfile
1924 sess_opts = ort.SessionOptions()
1925 sess_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
1926 sess_opts.log_severity_level = 3
1927
1928 results = {}
1929 errors = []
1930
1931 # --- FP32 model benchmark ---
1932 try:
1933 if onnx is not None:
1934 model_bytes = _build_test_model()
1935 else:
1936 model_bytes = None
1937
1938 if model_bytes is None:
1939 raise RuntimeError("onnx package required to build test model")
1940
1941 with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
1942 f.write(model_bytes)
1943 fp32_path = f.name
1944
1945 t0 = time.perf_counter()
1946 session = ort.InferenceSession(fp32_path, sess_opts, providers=[provider])
1947 session_create_ms = (time.perf_counter() - t0) * 1000.0
1948
1949 input_meta = session.get_inputs()[0]
1950 output_meta = session.get_outputs()[0]
1951 input_name = input_meta.name
1952 actual_input_shape = input_meta.shape
1953 output_shape = output_meta.shape
1954
1955 if isinstance(actual_input_shape[0], int):
1956 bench_batch = actual_input_shape[0]
1957 else:
1958 bench_batch = batch_size
1959 actual_input_shape_resolved = [bench_batch] + [
1960 d if isinstance(d, int) else hidden for d in actual_input_shape[1:]
1961 ]
1962
1963 feed = {input_name: np.random.randn(*actual_input_shape_resolved).astype(np.float32)}
1964
1965 for _ in range(3):
1966 session.run(None, feed)
1967
1968 num_runs = 20
1969 latencies = []
1970 for _ in range(num_runs):
1971 t_start = time.perf_counter()
1972 session.run(None, feed)
1973 latencies.append((time.perf_counter() - t_start) * 1000.0)
1974
1975 latencies.sort()
1976 p50 = latencies[len(latencies) // 2]
1977 p95 = latencies[int(len(latencies) * 0.95)]
1978 p99 = latencies[int(len(latencies) * 0.99)]
1979 avg_latency = sum(latencies) / len(latencies)
1980 throughput = bench_batch / (avg_latency / 1000.0) if avg_latency > 0 else 0.0
1981
1982 mem_rss_mb = 0.0
1983 try:
1984 import resource
1985 mem_rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0
1986 except Exception:
1987 pass
1988
1989 results.update({
1990 "fp32_session_create_ms": round(session_create_ms, 2),
1991 "fp32_inference_latency_p50_ms": round(p50, 3),
1992 "fp32_inference_latency_p95_ms": round(p95, 3),
1993 "fp32_inference_latency_p99_ms": round(p99, 3),
1994 "fp32_avg_latency_ms": round(avg_latency, 3),
1995 "fp32_throughput_inf_per_sec": round(throughput, 2),
1996 "fp32_inference_samples": [round(l, 3) for l in latencies],
1997 "input_shape": actual_input_shape_resolved,
1998 "output_shape": list(output_shape) if output_shape else [],
1999 "input_dtype": str(input_meta.type),
2000 "output_dtype": str(output_meta.type),
2001 "peak_rss_mb": round(mem_rss_mb, 1),
2002 "num_warmup": 3,
2003 "num_timed_runs": num_runs,
2004 })
2005
2006 os.unlink(fp32_path)
2007 del session
2008 except Exception as exc:
2009 results["fp32_error"] = str(exc)
2010 errors.append(f"FP32 benchmark failed: {exc}")
2011
2012 # --- Quantized model benchmark ---
2013 quantized_ok = False
2014 try:
2015 if onnx is not None:
2016 q_model_bytes = _build_quantized_model()
2017 if q_model_bytes is not None:
2018 with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
2019 f.write(q_model_bytes)
2020 quant_path = f.name
2021
2022 t0 = time.perf_counter()
2023 q_session = ort.InferenceSession(
2024 quant_path, sess_opts, providers=[provider]
2025 )
2026 q_create_ms = (time.perf_counter() - t0) * 1000.0
2027
2028 q_input = q_session.get_inputs()[0]
2029 q_feed = {
2030 q_input.name: np.random.randn(
2031 *([batch_size] + [
2032 d if isinstance(d, int) else hidden
2033 for d in q_input.shape[1:]
2034 ])
2035 ).astype(np.float32)
2036 }
2037
2038 for _ in range(3):
2039 q_session.run(None, q_feed)
2040
2041 q_latencies = []
2042 for _ in range(20):
2043 t_start = time.perf_counter()
2044 q_session.run(None, q_feed)
2045 q_latencies.append((time.perf_counter() - t_start) * 1000.0)
2046
2047 q_latencies.sort()
2048 q_p50 = q_latencies[len(q_latencies) // 2]
2049 q_avg = sum(q_latencies) / len(q_latencies)
2050 q_throughput = batch_size / (q_avg / 1000.0) if q_avg > 0 else 0.0
2051
2052 results.update({
2053 "quantized_session_create_ms": round(q_create_ms, 2),
2054 "quantized_inference_latency_p50_ms": round(q_p50, 3),
2055 "quantized_avg_latency_ms": round(q_avg, 3),
2056 "quantized_throughput_inf_per_sec": round(q_throughput, 2),
2057 "quantized_inference_samples": [round(l, 3) for l in q_latencies],
2058 "quantized_ops": ["DynamicQuantizeLinear", "MatMulInteger"],
2059 })
2060 quantized_ok = True
2061 os.unlink(quant_path)
2062 del q_session
2063 except Exception as exc:
2064 results["quantized_error"] = str(exc)
2065 errors.append(f"Quantized benchmark failed: {exc}")
2066
2067 p50_val = results.get("fp32_inference_latency_p50_ms", 0.0)
2068 p95_val = results.get("fp32_inference_latency_p95_ms", 0.0)
2069 p99_val = results.get("fp32_inference_latency_p99_ms", 0.0)
2070 tput_val = results.get("fp32_throughput_inf_per_sec", 0.0)
2071
2072 metrics = {
2073 "ort_version": ort_version,
2074 "provider": provider,
2075 "providers_available": providers_available,
2076 "graph_opt_level": "ORT_ENABLE_ALL",
2077 "model_load_ms": results.get("fp32_session_create_ms", 0.0),
2078 "session_create_ms": results.get("fp32_session_create_ms", 0.0),
2079 "inference_latency_p50_ms": p50_val,
2080 "inference_latency_p95_ms": p95_val,
2081 "inference_latency_p99_ms": p99_val,
2082 "throughput_inf_per_sec": tput_val,
2083 "input_shape": results.get("input_shape", input_shape),
2084 "output_shape": results.get("output_shape", []),
2085 "input_dtype": results.get("input_dtype", "tensor(float)"),
2086 "output_dtype": results.get("output_dtype", "tensor(float)"),
2087 "peak_rss_mb": results.get("peak_rss_mb", 0.0),
2088 "num_warmup": results.get("num_warmup", 0),
2089 "num_timed_runs": results.get("num_timed_runs", 0),
2090 "quantized_supported": quantized_ok,
2091 "inference_samples": results.get("fp32_inference_samples", []),
2092 }
2093
2094 for k in ("fp32_inference_samples", "fp32_session_create_ms", "fp32_inference_latency_p50_ms",
2095 "fp32_inference_latency_p95_ms", "fp32_inference_latency_p99_ms",
2096 "fp32_avg_latency_ms", "fp32_throughput_inf_per_sec",
2097 "quantized_session_create_ms", "quantized_inference_latency_p50_ms",
2098 "quantized_avg_latency_ms", "quantized_throughput_inf_per_sec",
2099 "quantized_inference_samples", "quantized_ops",
2100 "quantized_error", "fp32_error"):
2101 if k in results:
2102 metrics[k] = results[k]
2103
2104 success = True
2105 return success, metrics, errors
2106
2107
2108BENCHES = {
2109 "gpu-info": _gpu_info,
2110 "memory-bandwidth": _memory_bandwidth,
2111 "tensor-core": _tensor_core,
2112 "gemm": _gemm,
2113 "pytorch": _pytorch,
2114 "flash-attention": _flash_attention,
2115 "vllm": _vllm,
2116 "deepspeed": _deepspeed,
2117 "megatron": _megatron,
2118 "onnx": _onnx,
2119 "llama-cpp": _llama_cpp,
2120}
2121
2122
2123def main():
2124 parser = argparse.ArgumentParser()
2125 parser.add_argument("bench", help=f"One of: {', '.join(BENCHES.keys())}")
2126 parser.add_argument("--json", action="store_true", help="Emit JSON")
2127 args = parser.parse_args()
2128
2129 fn = BENCHES.get(args.bench)
2130 if fn is None:
2131 payload = {
2132 "name": args.bench,
2133 "success": False,
2134 "execution_time_ms": 0,
2135 "metrics": {},
2136 "errors": [f"Unknown benchmark {args.bench}"],
2137 }
2138 print(json.dumps(payload, indent=2))
2139 sys.exit(1)
2140
2141 start = time.perf_counter()
2142 success, metrics, errors = fn()
2143 elapsed_ms = int((time.perf_counter() - start) * 1000)
2144 payload = {
2145 "name": args.bench,
2146 "success": success,
2147 "execution_time_ms": elapsed_ms,
2148 "metrics": metrics,
2149 "errors": errors,
2150 }
2151
2152 # Print a clear marker to help the Rust wrapper find the JSON block
2153 print("\n---BENCHMARK_RESULTS_START---")
2154 print(json.dumps(payload, indent=2))
2155 sys.exit(0 if success else 1)
2156
2157
2158if __name__ == "__main__":
2159 main()
2160"#;