From e279b0bf72aa0ab9150253a9cddb7f6ea7d8c169 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sat, 30 May 2026 22:25:19 +0800 Subject: [PATCH] Optimize large add_constant tensors (#24755) Co-authored-by: Codex Co-authored-by: BBuf --- .../benchmark/bench_add_constant.py | 59 ++++++++++++++++++ .../sglang/jit_kernel/csrc/add_constant.cuh | 60 ++++++++++++++++--- .../jit_kernel/tests/test_add_constant.py | 21 ++++++- 3 files changed, 130 insertions(+), 10 deletions(-) create mode 100644 python/sglang/jit_kernel/benchmark/bench_add_constant.py diff --git a/python/sglang/jit_kernel/benchmark/bench_add_constant.py b/python/sglang/jit_kernel/benchmark/bench_add_constant.py new file mode 100644 index 000000000..fbadf112c --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_add_constant.py @@ -0,0 +1,59 @@ +import torch +import triton +import triton.testing + +from sglang.jit_kernel.add_constant import _jit_add_constant_module, add_constant +from sglang.jit_kernel.benchmark.utils import ( + DEFAULT_DEVICE, + get_benchmark_range, + run_benchmark_no_cudagraph, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=15, suite="base-b-kernel-benchmark-1-gpu-large") + +CONSTANT = 7 +SIZE_LIST = get_benchmark_range( + full_range=[128, 1024, 1025, 4096, 4097, 65536, 2**20, 2**22, 2**24], + ci_range=[4096, 2**20], +) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["size"], + x_vals=SIZE_LIST, + line_arg="provider", + line_vals=["jit_module", "jit_wrapper", "torch"], + line_names=["JIT module", "JIT wrapper", "PyTorch"], + styles=[("blue", "-"), ("orange", "-"), ("green", "--")], + ylabel="us", + plot_name="add-constant-performance", + args={}, + ) +) +def benchmark(size: int, provider: str): + src = torch.arange(size, dtype=torch.int32, device=DEFAULT_DEVICE) + + if provider == "jit_module": + dst = torch.empty_like(src) + module = _jit_add_constant_module(CONSTANT) + + def fn(): + module.add_constant(dst, src) + + elif provider == "jit_wrapper": + + def fn(): + add_constant(src, CONSTANT) + + else: + + def fn(): + src + CONSTANT + + return run_benchmark_no_cudagraph(fn) + + +if __name__ == "__main__": + benchmark.run(print_data=True) diff --git a/python/sglang/jit_kernel/csrc/add_constant.cuh b/python/sglang/jit_kernel/csrc/add_constant.cuh index 6c723a761..f42be433c 100644 --- a/python/sglang/jit_kernel/csrc/add_constant.cuh +++ b/python/sglang/jit_kernel/csrc/add_constant.cuh @@ -2,6 +2,7 @@ #include // For div_ceil, RuntimeCheck #include // For LaunchKernel +#include #include #include @@ -11,6 +12,17 @@ namespace { +constexpr size_t kBlockSize = 256; +constexpr size_t kVectorizedMinElements = 1 << 20; +constexpr size_t kVectorBytes = device::kMaxVecBytes; +static_assert(kVectorBytes % sizeof(int32_t) == 0, "Vector byte width must contain whole int32_t elements"); +constexpr size_t kElementsPerVector = kVectorBytes / sizeof(int32_t); + +template +bool is_aligned_for_vector(const int32_t* ptr) { + return reinterpret_cast(ptr) % alignof(Vector) == 0; +} + template __global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; @@ -19,7 +31,28 @@ __global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t len } } -constexpr size_t kBlockSize = 256; +template +__global__ void add_constant_vectorized_kernel(int32_t* dst, const int32_t* src, size_t length) { + using Vector = device::AlignedVector; + + const size_t work_idx = blockIdx.x * blockDim.x + threadIdx.x; + const size_t vector_count = length / kElementsPerVector; + const size_t tail_start = vector_count * kElementsPerVector; + + if (work_idx < vector_count) { + auto values = device::load_as(src, work_idx); +#pragma unroll + for (size_t i = 0; i < kElementsPerVector; ++i) { + values[i] += kConstant; + } + device::store_as(dst, values, work_idx); + } else { + const size_t tail_idx = tail_start + work_idx - vector_count; + if (tail_idx < length) { + dst[tail_idx] = src[tail_idx] + kConstant; + } + } +} // You can also use struct with static method as an alternative template @@ -37,7 +70,6 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { // 2. Extract required parameters, prepare for kernel launch const size_t num_elements = N.unwrap(); - const size_t grid_size = div_ceil(num_elements, kBlockSize); const DLDevice device = device_.unwrap(); [[maybe_unused]] // optional, can be omitted const size_t dynamic_smem = 0; @@ -46,14 +78,24 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { // some extra runtime checks using host::RuntimeCheck RuntimeCheck(num_elements > 0, "We only support non-empty tensors, got num_elements = ", num_elements); + const auto* src_ptr = static_cast(src.data_ptr()); + auto* dst_ptr = static_cast(dst.data_ptr()); + using Vector = device::AlignedVector; + const bool is_vector_aligned = is_aligned_for_vector(src_ptr) && is_aligned_for_vector(dst_ptr); + // 3. Launch the kernel. Error code will be automatically checked. - LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)( - // kernel function - add_constant_kernel, - // kernel arguments - static_cast(dst.data_ptr()), - static_cast(src.data_ptr()), - num_elements); + if (num_elements >= kVectorizedMinElements && is_vector_aligned) { + const size_t vector_count = num_elements / kElementsPerVector; + const size_t tail_count = num_elements - vector_count * kElementsPerVector; + const size_t work_items = vector_count + tail_count; + const size_t grid_size = div_ceil(work_items, kBlockSize); + LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)( + add_constant_vectorized_kernel, dst_ptr, src_ptr, num_elements); + } else { + const size_t grid_size = div_ceil(num_elements, kBlockSize); + LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)( + add_constant_kernel, dst_ptr, src_ptr, num_elements); + } } } // namespace diff --git a/python/sglang/jit_kernel/tests/test_add_constant.py b/python/sglang/jit_kernel/tests/test_add_constant.py index 77cb31fb9..6a4ecf178 100644 --- a/python/sglang/jit_kernel/tests/test_add_constant.py +++ b/python/sglang/jit_kernel/tests/test_add_constant.py @@ -10,7 +10,7 @@ register_cuda_ci(est_time=45, suite="base-b-kernel-unit-1-gpu-large") register_cuda_ci(est_time=180, suite="nightly-kernel-1-gpu", nightly=True) -@pytest.mark.parametrize("size", [1, 2, 127, 128, 1024, 1025]) +@pytest.mark.parametrize("size", [1, 2, 127, 128, 1024, 1025, 4096, 4097]) @pytest.mark.parametrize("constant", [0, 1, 7, 1024, -3]) def test_add_constant(size: int, constant: int) -> None: src = torch.arange(0, size, dtype=torch.int32, device="cuda") @@ -18,5 +18,24 @@ def test_add_constant(size: int, constant: int) -> None: assert torch.all(dst == src + constant) +def test_add_constant_unaligned_input() -> None: + src = torch.arange(0, 4098, dtype=torch.int32, device="cuda")[1:] + dst = add_constant(src, 7) + assert torch.all(dst == src + 7) + + +@pytest.mark.parametrize("size", [2**20, 2**20 + 3]) +def test_add_constant_large_aligned_input(size: int) -> None: + src = torch.arange(0, size, dtype=torch.int32, device="cuda") + dst = add_constant(src, -3) + assert torch.all(dst == src - 3) + + +def test_add_constant_large_unaligned_input() -> None: + src = torch.arange(0, 2**20 + 4, dtype=torch.int32, device="cuda")[1:] + dst = add_constant(src, 7) + assert torch.all(dst == src + 7) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", "-s"]))