minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Optimize large add_constant tensors (#24755)

Browse Source 
Co-authored-by: Codex <codex@example.com>
Co-authored-by: BBuf <xiaoyu.zhang@radixark.net>
This commit is contained in:
Xiaoyu Zhang
2026-05-30 22:25:19 +08:00
committed by GitHub
co-authored by Codex BBuf
parent b421e60eed
commit e279b0bf72
3 changed files with 130 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/tests/test_add_constant.py
+20 -1
View File
@@ -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"]))