⚡️ Speed up function check_cuda_result by 84%

[codeflash-ai]

📄 84% (0.84x) speedup for check_cuda_result in python/sglang/srt/utils/common.py

⏱️ Runtime : 106 microseconds → 57.9 microseconds (best of 162 runs)

📝 Explanation and details

The optimization achieves an 83% speedup by eliminating expensive repeated module imports and attribute lookups that were occurring on every function call.

Key Changes:

• Moved import to module scope: The import cuda.bindings.runtime as cuda_rt was moved from inside the function to the top-level module scope
• Cached constant lookup: The cuda_rt.cudaError_t.cudaSuccess value is now pre-computed and stored in _CUDA_SUCCESS at import time

Why This Is Faster:
The line profiler shows the dramatic impact - the original version spent 99.8% of its time (76ms out of 77ms total) just importing the cuda runtime module on every call. The optimized version eliminates this entirely, reducing total runtime from 106μs to 58μs.

Python's import system has significant overhead when repeatedly importing modules, even when they're already cached. Additionally, the attribute chain lookup cuda_rt.cudaError_t.cudaSuccess involves multiple dictionary lookups that are now avoided.

Impact on Workloads:
Based on the function reference, check_cuda_result is called from CUDA memory allocation operations in hot paths like _malloc_raw. Since memory operations are frequently called during model inference and training, this optimization provides meaningful benefits for GPU-intensive workloads.

Test Case Performance:
All test cases show consistent speedups ranging from 31% to 175%, with the largest gains on simple success cases (no exceptions) and the smallest gains on large-scale operations where the relative cost of the import becomes less significant compared to data processing overhead.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 43 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import sys
# We need to mock 'cuda.bindings.runtime' and its members for testing.
import types

# imports
import pytest  # used for our unit tests
import torch.distributed
from sglang.srt.utils.common import check_cuda_result

# ------------------- Basic Test Cases -------------------

def test_success_no_results():
    # Test: CUDA success, no results
    codeflash_output = check_cuda_result([0]); result = codeflash_output # 1.91μs -> 876ns (118% faster)

def test_success_single_result():
    # Test: CUDA success, one result
    codeflash_output = check_cuda_result([0, 42]); result = codeflash_output # 1.91μs -> 789ns (142% faster)

def test_success_multiple_results():
    # Test: CUDA success, multiple results
    codeflash_output = check_cuda_result([0, 1, 2, 3]); result = codeflash_output # 1.81μs -> 849ns (113% faster)

def test_success_varied_types():
    # Test: CUDA success, results with varied types
    codeflash_output = check_cuda_result([0, "foo", 123, 4.56, None]); result = codeflash_output # 1.84μs -> 786ns (134% faster)

# ------------------- Edge Test Cases -------------------

def test_error_raises_exception():
    # Test: CUDA error (non-success code) triggers exception
    with pytest.raises(Exception) as excinfo:
        check_cuda_result([999]) # 2.33μs -> 1.39μs (67.4% faster)

def test_error_with_results_raises_exception():
    # Test: CUDA error with additional results still raises
    with pytest.raises(Exception) as excinfo:
        check_cuda_result([2, "should", "not", "matter"]) # 2.38μs -> 1.43μs (65.9% faster)

def test_success_with_none_result():
    # Test: CUDA success with None as result
    codeflash_output = check_cuda_result([0, None]); result = codeflash_output # 1.80μs -> 845ns (113% faster)

def test_success_with_empty_string_and_zero():
    # Test: CUDA success with empty string and zero as results
    codeflash_output = check_cuda_result([0, "", 0]); result = codeflash_output # 1.77μs -> 768ns (131% faster)

def test_success_with_nested_list_result():
    # Test: CUDA success with a nested list as a result
    nested = [1, 2, [3, 4]]
    codeflash_output = check_cuda_result([0, nested]); result = codeflash_output # 1.81μs -> 762ns (137% faster)

def test_error_code_is_not_int():
    # Test: CUDA error code is not an int (simulate bad input)
    class WeirdCode:
        def __eq__(self, other):
            return False
        def __str__(self):
            return "WeirdCode"
    with pytest.raises(Exception) as excinfo:
        check_cuda_result([WeirdCode()]) # 3.57μs -> 2.49μs (43.3% faster)

def test_success_with_bool_result():
    # Test: CUDA success with boolean result
    codeflash_output = check_cuda_result([0, True, False]); result = codeflash_output # 2.01μs -> 851ns (136% faster)

def test_success_with_large_integers():
    # Test: CUDA success with very large integer results
    big = 2**62
    codeflash_output = check_cuda_result([0, big, -big]); result = codeflash_output # 1.97μs -> 835ns (136% faster)

def test_success_with_empty_tuple_result():
    # Test: CUDA success with empty tuple as a result
    codeflash_output = check_cuda_result([0, ()]); result = codeflash_output # 1.91μs -> 782ns (145% faster)

# ------------------- Large Scale Test Cases -------------------

def test_success_many_results():
    # Test: CUDA success with a large number of results (under 1000 elements)
    results = list(range(999))
    codeflash_output = check_cuda_result([0] + results); out = codeflash_output # 4.37μs -> 3.32μs (31.7% faster)

def test_success_large_string_result():
    # Test: CUDA success with a large string result (under 100MB)
    large_str = "a" * (10**6)  # 1MB string
    codeflash_output = check_cuda_result([0, large_str]); result = codeflash_output # 2.69μs -> 978ns (175% faster)

def test_success_large_mixed_results():
    # Test: CUDA success with a large list of mixed types
    results = []
    for i in range(500):
        results.append(i)
        results.append(str(i))
        results.append(None if i % 2 == 0 else True)
    codeflash_output = check_cuda_result([0] + results); out = codeflash_output # 5.30μs -> 4.09μs (29.4% faster)

def test_success_large_nested_structure():
    # Test: CUDA success with a large nested structure
    nested = [[j for j in range(10)] for i in range(50)]  # 50 lists of 10 ints
    codeflash_output = check_cuda_result([0, nested]); result = codeflash_output # 1.93μs -> 756ns (155% faster)

def test_success_large_empty_results():
    # Test: CUDA success with many empty results
    results = [None] * 999
    codeflash_output = check_cuda_result([0] + results); out = codeflash_output # 3.76μs -> 2.90μs (29.5% faster)

# ------------------- Additional Defensive/Mutation Tests -------------------

def test_success_does_not_raise_on_success():
    # Test: Should not raise when err == cudaSuccess
    try:
        check_cuda_result([0, "ok"])
    except Exception:
        pass

def test_raises_on_nonzero_even_if_result_is_none():
    # Test: Should raise if err != cudaSuccess, even if results are None
    with pytest.raises(Exception):
        check_cuda_result([11, None]) # 2.42μs -> 1.40μs (73.5% faster)

def test_result_identity():
    # Test: Returned list is a new list (not a reference to input)
    input_results = [0, 1, 2, 3]
    codeflash_output = check_cuda_result(input_results); out = codeflash_output # 1.84μs -> 809ns (127% faster)

def test_success_with_falsey_results():
    # Test: CUDA success with falsey results (0, '', [], {}, False)
    results = [0, '', [], {}, False]
    codeflash_output = check_cuda_result([0] + results); out = codeflash_output # 1.77μs -> 765ns (132% faster)

def test_error_message_contains_code():
    # Test: Error message should contain the error code
    error_code = 2
    with pytest.raises(Exception) as excinfo:
        check_cuda_result([error_code]) # 2.44μs -> 1.34μs (82.5% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

import sys
import types

# imports
import pytest
import torch.distributed
from sglang.srt.utils.common import check_cuda_result

# 1. BASIC TEST CASES

def test_success_no_results():
    # Test: Success error code, no results
    codeflash_output = check_cuda_result([0]); result = codeflash_output # 1.83μs -> 768ns (139% faster)

def test_success_single_result():
    # Test: Success error code, one result
    codeflash_output = check_cuda_result([0, 42]); result = codeflash_output # 1.76μs -> 703ns (150% faster)

def test_success_multiple_results():
    # Test: Success error code, multiple results
    codeflash_output = check_cuda_result([0, 'a', 3.14, None]); result = codeflash_output # 1.70μs -> 713ns (139% faster)

def test_error_raises_exception():
    # Test: Non-success error code raises Exception
    with pytest.raises(Exception) as excinfo:
        check_cuda_result([1, "should not be returned"]) # 2.32μs -> 1.34μs (73.0% faster)

def test_error_with_multiple_results_raises():
    # Test: Non-success error code with multiple results still raises
    with pytest.raises(Exception) as excinfo:
        check_cuda_result([2, "foo", "bar"]) # 2.36μs -> 1.35μs (75.1% faster)

# 2. EDGE TEST CASES

def test_success_with_falsey_results():
    # Test: Success with results that are falsey values
    codeflash_output = check_cuda_result([0, 0, False, '', None]); result = codeflash_output # 1.95μs -> 842ns (132% faster)

def test_success_with_nested_list():
    # Test: Success with a nested list as a result
    codeflash_output = check_cuda_result([0, [1, 2, 3]]); result = codeflash_output # 1.70μs -> 687ns (147% faster)

def test_success_with_dict_result():
    # Test: Success with a dict as a result
    d = {'key': 'value'}
    codeflash_output = check_cuda_result([0, d]); result = codeflash_output # 1.64μs -> 720ns (128% faster)

def test_success_with_tuple_result():
    # Test: Success with a tuple as a result
    t = (1, 2)
    codeflash_output = check_cuda_result([0, t]); result = codeflash_output # 1.75μs -> 720ns (142% faster)

def test_error_code_is_negative():
    # Test: Negative error code (simulate error)
    with pytest.raises(Exception) as excinfo:
        check_cuda_result([-1, "foo"]) # 2.68μs -> 1.49μs (79.5% faster)

def test_error_code_is_large():
    # Test: Large error code (simulate error)
    with pytest.raises(Exception) as excinfo:
        check_cuda_result([123456, "foo"]) # 2.56μs -> 1.40μs (82.7% faster)

def test_empty_input_list():
    # Test: Empty input list (should fail with ValueError due to unpacking)
    with pytest.raises(ValueError):
        check_cuda_result([]) # 3.20μs -> 2.35μs (36.2% faster)

def test_success_with_bytes_result():
    # Test: Success with bytes as result
    b = b'abc'
    codeflash_output = check_cuda_result([0, b]); result = codeflash_output # 3.62μs -> 1.37μs (164% faster)

# 3. LARGE SCALE TEST CASES

def test_success_many_results():
    # Test: Success with a large number of results (under 1000 elements)
    big_list = list(range(999))
    codeflash_output = check_cuda_result([0] + big_list); result = codeflash_output # 4.81μs -> 3.32μs (44.6% faster)

def test_success_large_string_result():
    # Test: Success with a large string
    s = 'x' * (10**6)  # 1MB string
    codeflash_output = check_cuda_result([0, s]); result = codeflash_output # 2.91μs -> 1.14μs (155% faster)

def test_success_large_nested_structure():
    # Test: Success with a large nested structure
    nested = [[i, i+1] for i in range(500)]
    codeflash_output = check_cuda_result([0, nested]); result = codeflash_output # 2.32μs -> 900ns (157% faster)

def test_success_large_dict_result():
    # Test: Success with a large dict
    d = {str(i): i for i in range(500)}
    codeflash_output = check_cuda_result([0, d]); result = codeflash_output # 2.09μs -> 819ns (155% faster)

def test_success_large_tuple_result():
    # Test: Success with a large tuple
    t = tuple(range(500))
    codeflash_output = check_cuda_result([0, t]); result = codeflash_output # 1.98μs -> 743ns (166% faster)

def test_success_with_mixed_types_large():
    # Test: Success with a large mixture of types
    vals = [i if i % 2 == 0 else str(i) for i in range(500)]
    codeflash_output = check_cuda_result([0] + vals); result = codeflash_output # 3.07μs -> 2.14μs (43.2% faster)

def test_error_with_large_results():
    # Test: Error code with large results still raises
    big_list = list(range(900))
    with pytest.raises(Exception):
        check_cuda_result([2] + big_list) # 4.99μs -> 3.79μs (31.8% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-check_cuda_result-mijv7arc and push.