Merge pull request #27 from flagos-ai/cuda_update

merge

Author: zyq1105331849

src/flagsparse/sparse_operations/_common.py

@@ -1,5 +1,6 @@
 """Shared imports, dtypes, and helpers for FlagSparse sparse ops."""
 
+import statistics
 import time
 
 try:
@@ -53,6 +54,7 @@
     "_prepare_inputs",
     "_prepare_scatter_inputs",
     "_benchmark_cuda_op",
+    "_benchmark_cuda_graph_op",
     "cp",
     "cpx_sparse",
     "time",
@@ -332,3 +334,50 @@ def _benchmark_cuda_op(op, warmup, iters):
         cp.cuda.runtime.deviceSynchronize()
     elapsed_ms = (time.perf_counter() - start_time) * 1000.0 / iters
     return output, elapsed_ms
+
+
+def _benchmark_cuda_graph_op(
+    op,
+    *,
+    graph_batch=100,
+    warmup=20,
+    repeats=10,
+    capture_setup=None,
+):
+    """Measure allocation-free CUDA work without Python/FFI launch gaps."""
+    graph_batch = max(1, int(graph_batch))
+    warmup = max(0, int(warmup))
+    repeats = max(1, int(repeats))
+
+    torch.cuda.synchronize()
+    capture_stream = torch.cuda.Stream()
+    graph = torch.cuda.CUDAGraph()
+
+    with torch.cuda.stream(capture_stream):
+        if capture_setup is not None:
+            capture_setup()
+        op()
+    capture_stream.synchronize()
+
+    with torch.cuda.graph(graph, stream=capture_stream):
+        for _ in range(graph_batch):
+            op()
+    capture_stream.synchronize()
+
+    with torch.cuda.stream(capture_stream):
+        for _ in range(warmup):
+            graph.replay()
+    capture_stream.synchronize()
+
+    samples_ms = []
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    for _ in range(repeats):
+        with torch.cuda.stream(capture_stream):
+            start.record()
+            graph.replay()
+            end.record()
+        end.synchronize()
+        samples_ms.append(float(start.elapsed_time(end)) / graph_batch)
+
+    return statistics.median(samples_ms)

src/flagsparse/sparse_operations/benchmarks.py

@@ -7,9 +7,11 @@
     _PreparedCusparseNativeGather,
     SUPPORTED_SCATTER_VALUE_DTYPES,
     _scatter_dtype_error_message,
+    _cusparse_native_gather_skip_reason,
     _cusparse_spmv,
     _make_scatter_selector_matrix,
     _pytorch_scatter_impl,
+    _set_cusparse_stream,
     _triton_gather_impl,
     _triton_scatter_impl,
 )
@@ -33,6 +35,9 @@
 from .spsm import benchmark_spsm_case
 
 
+_GATHER_GRAPH_BATCH = 100
+
+
 def _cupy_spmv_op_matrix(matrix, op_code):
     if op_code == SPMV_OP_NON:
         return matrix
@@ -80,21 +85,43 @@ def benchmark_gather_case(
     block_size=DEFAULT_GATHER_BLOCK_SIZE,
     run_cusparse=True,
 ):
-    """Benchmark Triton vs PyTorch indexing vs native cuSPARSE gather."""
+    """Benchmark gather device time with CUDA Graph replay and CUDA events."""
     device = torch.device("cuda")
     dense_vector = _build_random_dense(dense_size, value_dtype, device)
     indices = _build_indices(nnz, dense_size, index_dtype, device, unique=False)
 
     dense_vector, indices, kernel_indices = _prepare_inputs(dense_vector, indices)
     expected = dense_vector[indices]
 
-    pytorch_op = lambda: dense_vector[indices]
+    pytorch_out = torch.empty_like(expected)
+    triton_out = torch.empty_like(expected)
+    pytorch_op = lambda: torch.index_select(
+        dense_vector, 0, indices, out=pytorch_out
+    )
     triton_op = lambda: _triton_gather_impl(
-        dense_vector, kernel_indices, block_size=block_size
+        dense_vector, kernel_indices, out=triton_out, block_size=block_size
     )
 
-    pytorch_values, pytorch_ms = _benchmark_cuda_op(pytorch_op, warmup=warmup, iters=iters)
-    triton_values, triton_ms = _benchmark_cuda_op(triton_op, warmup=warmup, iters=iters)
+    try:
+        pytorch_ms = _benchmark_cuda_graph_op(
+            pytorch_op,
+            graph_batch=_GATHER_GRAPH_BATCH,
+            warmup=warmup,
+            repeats=iters,
+        )
+    except Exception as exc:
+        raise RuntimeError(f"PyTorch CUDA Graph timing failed: {exc}") from exc
+    try:
+        triton_ms = _benchmark_cuda_graph_op(
+            triton_op,
+            graph_batch=_GATHER_GRAPH_BATCH,
+            warmup=warmup,
+            repeats=iters,
+        )
+    except Exception as exc:
+        raise RuntimeError(f"Triton CUDA Graph timing failed: {exc}") from exc
+    pytorch_values = pytorch_out
+    triton_values = triton_out
 
     atol, rtol = _tolerance_for_dtype(value_dtype)
     triton_match = torch.allclose(triton_values, expected, atol=atol, rtol=rtol)
@@ -109,7 +136,7 @@ def benchmark_gather_case(
     cusparse_max_error = None
     cusparse_reason = None
     if run_cusparse:
-        skip_reason = _cusparse_baseline_skip_reason(value_dtype)
+        skip_reason = _cusparse_native_gather_skip_reason(value_dtype)
         if skip_reason:
             cusparse_reason = skip_reason
         else:
@@ -120,9 +147,21 @@ def benchmark_gather_case(
                     dense_vector, indices, out=cusparse_out
                 )
                 cusparse_op = lambda: cusparse_plan.run()
-                cusparse_values, cusparse_ms = _benchmark_cuda_op(
-                    cusparse_op, warmup=warmup, iters=iters
-                )
+                try:
+                    cusparse_ms = _benchmark_cuda_graph_op(
+                        cusparse_op,
+                        graph_batch=_GATHER_GRAPH_BATCH,
+                        warmup=warmup,
+                        repeats=iters,
+                        capture_setup=lambda: _set_cusparse_stream(
+                            cusparse_plan.lib, cusparse_plan.handle, strict=True
+                        ),
+                    )
+                except Exception as exc:
+                    raise RuntimeError(
+                        f"cuSPARSE CUDA Graph timing failed: {exc}"
+                    ) from exc
+                cusparse_values = cusparse_out
                 cusparse_match = torch.allclose(
                     cusparse_values, expected, atol=atol, rtol=rtol
                 )
@@ -154,13 +193,15 @@ def benchmark_gather_case(
             "index_dtype": str(index_dtype),
             "warmup": warmup,
             "iters": iters,
+            "kernel_graph_batch": _GATHER_GRAPH_BATCH,
         },
         "performance": {
             "pytorch_ms": pytorch_ms,
             "triton_ms": triton_ms,
             "cusparse_ms": cusparse_ms,
             "triton_speedup_vs_pytorch": triton_speedup_vs_pytorch,
             "triton_speedup_vs_cusparse": triton_speedup_vs_cusparse,
+            "kernel_timing_method": "cuda_graph_event_amortized_device_estimate",
         },
         "verification": {
             "triton_match_pytorch": triton_match,

src/flagsparse/sparse_operations/gather_scatter.py

@@ -298,6 +298,12 @@ def _validate_gather_value_dtype(dense_vector, op_name):
 _CUSPARSE_LIB_LOAD_ERROR = None
 
 
+def _cusparse_native_gather_skip_reason(value_dtype):
+    if value_dtype == torch.bfloat16:
+        return "bfloat16 is not supported by native cusparseGather; skipped"
+    return None
+
+
 def _cuda_data_type_from_torch(torch_dtype):
     mapping = {
         torch.float16: _CUDA_R_16F,
@@ -388,27 +394,33 @@ def _check_cusparse_status(status, op_name):
         raise RuntimeError(f"{op_name} failed with cuSPARSE status {int(status)}")
 
 
-def _set_cusparse_stream(lib, handle):
+def _set_cusparse_stream(lib, handle, *, strict=False):
     if not hasattr(lib, "cusparseSetStream"):
+        if strict:
+            raise RuntimeError("Loaded cuSPARSE library does not export cusparseSetStream")
         return
     try:
         stream = torch.cuda.current_stream()
         stream_ptr = getattr(stream, "cuda_stream", None)
         if stream_ptr is None:
+            if strict:
+                raise RuntimeError("Could not obtain the current CUDA stream pointer")
             return
         _check_cusparse_status(
             lib.cusparseSetStream(handle, ctypes.c_void_p(int(stream_ptr))),
             "cusparseSetStream",
         )
     except Exception:
+        if strict:
+            raise
         return
 
 
 class _PreparedCusparseNativeGather:
     def __init__(self, dense_vector, indices, out=None):
         dense_vector, indices, _ = _prepare_inputs(dense_vector, indices)
         _validate_gather_value_dtype(dense_vector, "cusparse_native_gather")
-        skip_reason = _cusparse_baseline_skip_reason(dense_vector.dtype)
+        skip_reason = _cusparse_native_gather_skip_reason(dense_vector.dtype)
         if skip_reason:
             raise RuntimeError(skip_reason)
 

tests/test_gather.py

@@ -2,6 +2,14 @@
 import csv
 import math
 import os
+import sys
+from pathlib import Path
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+SRC_ROOT = PROJECT_ROOT / "src"
+if SRC_ROOT.is_dir():
+    sys.path.insert(0, str(SRC_ROOT))
 
 import torch
 
@@ -18,6 +26,7 @@
 DEFAULT_INDEX_DTYPES = "int32,int64"
 WARMUP = 20
 ITERS = 200
+KERNEL_GRAPH_BATCH = 100
 
 
 def _fmt_ms(value):
@@ -82,8 +91,8 @@ def _parse_cases(raw):
         left, right = item.split(":", 1)
         dense_size = int(left)
         nnz = int(right)
-        if dense_size < 0 or nnz < 0:
-            raise ValueError(f"case values must be non-negative: {item}")
+        if dense_size <= 0 or nnz <= 0:
+            raise ValueError(f"case values must be positive: {item}")
         pairs.append((dense_size, nnz))
     if not pairs:
         raise ValueError("case list is empty")
@@ -172,7 +181,7 @@ def _print_header():
     print("-" * 196)
     print(
         f"{'ValueReq':>14} {'ValueEff':>18} {'Index':>6} {'Dense':>10} {'NNZ':>10} "
-        f"{'IFB':>4} {'PT(ms)':>10} {'FS(ms)':>10} {'CS(ms)':>10} "
+        f"{'IFB':>4} {'FS(ms)':>10} {'PT(ms)':>10} {'CS(ms)':>10} "
         f"{'FS/PT':>8} {'FS/CS':>8} {'Status':>6} {'Err(FS)':>12} {'Err(CS)':>12}"
     )
     print("-" * 196)
@@ -182,7 +191,7 @@ def _print_row(row):
     print(
         f"{row['value_dtype_req']:>14} {row['value_dtype_compute']:>18} {row['index_dtype']:>6} "
         f"{row['dense_size']:>10,d} {row['nnz']:>10,d} {str(row['index_fallback_applied']):>4} "
-        f"{_fmt_ms(row['pytorch_ms']):>10} {_fmt_ms(row['triton_ms']):>10} {_fmt_ms(row['cusparse_ms']):>10} "
+        f"{_fmt_ms(row['triton_ms']):>10} {_fmt_ms(row['pytorch_ms']):>10} {_fmt_ms(row['cusparse_ms']):>10} "
         f"{_fmt_speedup(row['triton_speedup_vs_pytorch']):>8} {_fmt_speedup(row['triton_speedup_vs_cusparse']):>8} "
         f"{row['status']:>6} {_fmt_err(row['triton_max_error']):>12} {_fmt_err(row['cusparse_max_error']):>12}"
     )
@@ -201,9 +210,11 @@ def run_cli(args):
     print("=" * 180)
     print("FLAGSPARSE GATHER BENCHMARK/VALIDATION")
     print("=" * 180)
+    print(f"FlagSparse source: {Path(ast.__file__).resolve()}")
     print(f"GPU: {torch.cuda.get_device_name(0)}")
     print(
         f"Warmup: {args.warmup} | Iterations: {args.iters} | "
+        f"Kernel graph batch: {KERNEL_GRAPH_BATCH} | "
         f"index_fallback_policy: {args.index_fallback_policy}"
     )
     print()
@@ -242,6 +253,18 @@ def run_cli(args):
                     verify = result["verification"]
                     params = result["parameters"]
                     backend = result["backend_status"]
+                    timing_method = perf.get("kernel_timing_method")
+                    graph_batch = params.get("kernel_graph_batch")
+                    if timing_method != "cuda_graph_event_amortized_device_estimate":
+                        raise RuntimeError(
+                            "loaded benchmark_gather_case does not provide CUDA Graph timing; "
+                            f"flagsparse was loaded from {Path(ast.__file__).resolve()}"
+                        )
+                    if graph_batch != KERNEL_GRAPH_BATCH:
+                        raise RuntimeError(
+                            "unexpected gather graph batch: "
+                            f"expected {KERNEL_GRAPH_BATCH}, got {graph_batch}"
+                        )
                     status = _status_from_result(verify)
                     if status != "PASS":
                         failed_cases += 1
@@ -263,6 +286,8 @@ def run_cli(args):
                         "cusparse_ms": perf.get("cusparse_ms"),
                         "triton_speedup_vs_pytorch": perf.get("triton_speedup_vs_pytorch"),
                         "triton_speedup_vs_cusparse": perf.get("triton_speedup_vs_cusparse"),
+                        "kernel_timing_method": timing_method,
+                        "kernel_graph_batch": graph_batch,
                         "triton_match_pytorch": verify.get("triton_match_pytorch"),
                         "cusparse_match_pytorch": verify.get("cusparse_match_pytorch"),
                         "triton_max_error": verify.get("triton_max_error"),
@@ -285,6 +310,8 @@ def run_cli(args):
                         )
                 except Exception as exc:
                     failed_cases += 1
+                    error_text = f"{exc.__class__.__name__}: {exc}"
+                    print(f"\nERROR [{case_id}]: {error_text}")
                     row = {
                         "case_id": case_id,
                         "gpu": torch.cuda.get_device_name(0),
@@ -301,12 +328,14 @@ def run_cli(args):
                         "cusparse_ms": None,
                         "triton_speedup_vs_pytorch": None,
                         "triton_speedup_vs_cusparse": None,
+                        "kernel_timing_method": "cuda_graph_event_amortized_device_estimate",
+                        "kernel_graph_batch": KERNEL_GRAPH_BATCH,
                         "triton_match_pytorch": None,
                         "cusparse_match_pytorch": None,
                         "triton_max_error": None,
                         "cusparse_max_error": None,
-                        "cusparse_unavailable_reason": str(exc),
-                        "index_fallback_reason": str(exc),
+                        "cusparse_unavailable_reason": error_text,
+                        "index_fallback_reason": error_text,
                         "status": "ERROR",
                     }
                     summary_rows.append(row)
@@ -334,6 +363,8 @@ def run_cli(args):
             "cusparse_ms",
             "triton_speedup_vs_pytorch",
             "triton_speedup_vs_cusparse",
+            "kernel_timing_method",
+            "kernel_graph_batch",
             "triton_match_pytorch",
             "cusparse_match_pytorch",
             "triton_max_error",