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

[kernel] Share the warp vectorized copy and enforce its alignment (#36176)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: BBuf <1182563586@qq.com>
This commit is contained in:
DarkSharpness
2026-09-18 22:40:54 +08:00
committed by GitHub
co-authored by Claude Opus 5 BBuf
parent 9784d5f979
commit 81363bf8cb
34 changed files with 1172 additions and 991 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/benchmark/kvcache/bench_hicache.py
+23 -82
View File
@@ -15,19 +15,15 @@ Note: Uses do_bench instead of do_bench_cudagraph since CUDA graph
capture doesn't support CPU-GPU memory transfers.
"""
import itertools
import os
from dataclasses import dataclass
from typing import Tuple
import torch
import triton
import triton.testing
from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer
from sglang.kernels.jit.benchmark.utils import DEFAULT_QUANTILES, get_benchmark_range
from sglang.kernels.jit.benchmark import marker
from sglang.kernels.jit.benchmark.utils import get_benchmark_range
from sglang.kernels.ops.kvcache.hicache import (
can_use_hicache_jit_kernel,
transfer_hicache_all_layer,
transfer_hicache_one_layer,
)
@@ -39,7 +35,7 @@ register_cuda_ci(
register_amd_ci(est_time=29, stage="jit-kernel-benchmark", runner_config="amd")
DISABLE_TORCH = os.environ.get("DISABLE_TORCH", "0") == "1"
PAGE_SIZE = 1
PAGE_SIZE = int(os.environ.get("PAGE_SIZE", "1"))
ENABLE_SORT = True
GPU_CACHE_SIZE = 256 * 1024 # 256K tokens on GPU
HOST_CACHE_SIZE = 512 * 1024 # 512K tokens on CPU
@@ -187,20 +183,14 @@ def pytorch_transfer(
# Benchmark configuration
BS_RANGE = get_benchmark_range(
full_range=[2**n for n in range(0, 16)],
ci_range=[16],
)
ELEMENT_SIZE_RANGE = get_benchmark_range(
full_range=[64, 128, 256, 512, 1024],
ci_range=[1024],
)
LINE_VALS = ["aot", "jit", "torch"]
LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"]
STYLES = [("orange", "-"), ("blue", "--"), ("red", ":")]
CONFIGS = list(itertools.product(ELEMENT_SIZE_RANGE, BS_RANGE))
if DISABLE_TORCH:
LINE_VALS.remove("torch")
# =============================================================================
@@ -208,22 +198,10 @@ CONFIGS = list(itertools.product(ELEMENT_SIZE_RANGE, BS_RANGE))
# =============================================================================
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["element_size", "batch_size"],
x_vals=CONFIGS,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="hicache-one-layer-h2d",
args={},
)
)
def benchmark_one_layer_h2d(
element_size: int, batch_size: int, provider: str
) -> Tuple[float, float, float]:
@marker.parametrize("element_size", ELEMENT_SIZE_RANGE)
@marker.parametrize("batch_size", marker.range(14, pattern="pow2"), [16])
@marker.benchmark("provider", LINE_VALS, unit="ms")
def benchmark_one_layer_h2d(element_size: int, batch_size: int, provider: str):
"""One Layer: Host (CPU) -> Device (GPU)."""
global cache
cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size)
@@ -281,19 +259,10 @@ def benchmark_one_layer_h2d(
],
}
if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
return (float("nan"), float("nan"), float("nan"))
if DISABLE_TORCH and provider in ["torch"]:
return (float("nan"), float("nan"), float("nan"))
ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore
FN_MAP[provider], quantiles=DEFAULT_QUANTILES, warmup=5, rep=25
)
return (
1000 * ms / NUM_LAYERS,
1000 * max_ms / NUM_LAYERS,
1000 * min_ms / NUM_LAYERS,
return marker.do_bench(
FN_MAP[provider],
use_cuda_graph=False,
extra_memory_footprint=NUM_LAYERS * batch_size * (2 * element_bytes),
)
@@ -311,22 +280,10 @@ def _create_ptr_tensor(tensors, device="cuda"):
)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["element_size", "batch_size"],
x_vals=CONFIGS,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="hicache-all-layer-d2h",
args={},
)
)
def benchmark_all_layer_d2h(
element_size: int, batch_size: int, provider: str
) -> Tuple[float, float, float]:
@marker.parametrize("element_size", ELEMENT_SIZE_RANGE)
@marker.parametrize("batch_size", marker.range(14, pattern="pow2"), [16])
@marker.benchmark("provider", LINE_VALS, unit="ms")
def benchmark_all_layer_d2h(element_size: int, batch_size: int, provider: str):
"""All Layer: Device (GPU) -> Host (CPU)."""
global cache
cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size)
@@ -385,19 +342,10 @@ def benchmark_all_layer_d2h(
],
}
if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
return (float("nan"), float("nan"), float("nan"))
if DISABLE_TORCH and provider in ["torch"]:
return (float("nan"), float("nan"), float("nan"))
ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore
FN_MAP[provider], quantiles=DEFAULT_QUANTILES, warmup=5, rep=25
)
return (
1000 * ms / NUM_LAYERS,
1000 * max_ms / NUM_LAYERS,
1000 * min_ms / NUM_LAYERS,
return marker.do_bench(
FN_MAP[provider],
use_cuda_graph=False,
extra_memory_footprint=NUM_LAYERS * batch_size * (2 * element_bytes),
)
@@ -413,12 +361,5 @@ if __name__ == "__main__":
v_cache_host=torch.empty(HOST_SHAPE, dtype=torch.bfloat16, pin_memory=True),
)
print("=" * 60)
print("One Layer: Host -> Device (CPU -> GPU)")
print("=" * 60)
benchmark_one_layer_h2d.run(print_data=True)
print("\n" + "=" * 60)
print("All Layer: Device -> Host (GPU -> CPU) [per-layer avg]")
print("=" * 60)
benchmark_all_layer_d2h.run(print_data=True)
benchmark_one_layer_h2d.run(print_prefix="Per Layer: Host -> Device (CPU -> GPU)")
benchmark_all_layer_d2h.run(print_prefix="All Layer: Device -> Host (GPU -> CPU)")
test/registered/kernels/benchmark/kvcache/bench_set_mla_kv_buffer.py
+20 -55
View File
@@ -7,18 +7,13 @@ Compares three providers across a batch-size sweep:
- ``triton``: the BLOCK-tiled Triton kernel (SM<90 fallback path).
"""
import itertools
from typing import Tuple
import torch
import triton
import triton.testing
from sglang.kernels.jit.benchmark import marker
from sglang.kernels.jit.benchmark.utils import (
DEFAULT_DEVICE,
DEFAULT_DTYPE,
DEFAULT_QUANTILES,
get_benchmark_range,
)
from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.kernels.ops.kvcache.set_mla_kv_buffer import set_mla_kv_buffer as jit_set
@@ -50,62 +45,38 @@ def _triton_baseline(kv_buffer, loc, cache_k_nope, cache_k_rope):
cache_k_rope.stride(0),
nope_dim,
rope_dim,
BLOCK=BLOCK,
DCP_RANK=0,
DCP_WORLD_SIZE=1,
**pdl_kwargs,
BLOCK=BLOCK, # type: ignore
DCP_RANK=0, # type: ignore
DCP_WORLD_SIZE=1, # type: ignore
**pdl_kwargs, # type: ignore
)
NUM_LAYERS = 8
CACHE_SIZE = (2 * 1024 * 1024) // NUM_LAYERS
# 2M elements
CACHE_SIZE = 2 * 1024 * 1024
NOPE_DIM = 512
ROPE_DIM = 64
BS_RANGE = get_benchmark_range(
full_range=[1, 8, 32, 128, 512, 1024, 2048, 4096, 8192, 16384],
ci_range=[1, 128, 2048, 4096, 8192],
)
LINE_VALS = ["wrapper", "jit_tma", "triton"]
LINE_NAMES = ["Wrapper (auto)", "JIT TMA bulk-store", "Triton (BLOCK=128 baseline)"]
STYLES = [("blue", "-"), ("green", "--"), ("red", "-.")]
X_NAMES = ["batch_size"]
CONFIGS = list(itertools.product(BS_RANGE))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=X_NAMES,
x_vals=CONFIGS,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
styles=STYLES,
ylabel="us",
plot_name="set-mla-kv-buffer-performance",
args={},
)
)
def benchmark(batch_size: int, provider: str) -> Tuple[float, float, float]:
@marker.parametrize("batch_size", marker.range(15, pattern="pow2"), [1, 128, 8192])
@marker.benchmark("provider", ["wrapper", "jit_tma", "triton"])
def benchmark(batch_size: int, provider: str):
cache_k_nope = torch.randn(
(NUM_LAYERS, batch_size, 1, NOPE_DIM),
(batch_size, 1, NOPE_DIM),
dtype=DEFAULT_DTYPE,
device=DEFAULT_DEVICE,
)
cache_k_rope = torch.randn(
(NUM_LAYERS, batch_size, 1, ROPE_DIM),
(batch_size, 1, ROPE_DIM),
dtype=DEFAULT_DTYPE,
device=DEFAULT_DEVICE,
)
kv_buffer = torch.randn(
(NUM_LAYERS, CACHE_SIZE, 1, NOPE_DIM + ROPE_DIM),
(CACHE_SIZE, 1, NOPE_DIM + ROPE_DIM),
dtype=DEFAULT_DTYPE,
device=DEFAULT_DEVICE,
)
loc = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE)[:batch_size]
torch.cuda.synchronize()
FN_MAP = {
"wrapper": sglang_wrapper,
@@ -113,20 +84,14 @@ def benchmark(batch_size: int, provider: str) -> Tuple[float, float, float]:
"triton": _triton_baseline,
}
def fn():
impl = FN_MAP[provider]
for i in range(NUM_LAYERS):
impl(kv_buffer[i], loc, cache_k_nope[i], cache_k_rope[i])
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
fn, quantiles=DEFAULT_QUANTILES
)
return (
1000 * ms / NUM_LAYERS,
1000 * max_ms / NUM_LAYERS,
1000 * min_ms / NUM_LAYERS,
return marker.do_bench(
FN_MAP[provider],
input_args=(kv_buffer, loc, cache_k_nope, cache_k_rope),
graph_clone_args=(1, 2, 3),
memory_args=(loc, cache_k_nope, cache_k_rope),
memory_output=(cache_k_nope, cache_k_rope),
)
if __name__ == "__main__":
benchmark.run(print_data=True)
benchmark.run()
test/registered/kernels/ops/kvcache/test_store_cache.py
+2 -93
View File
@@ -17,7 +17,7 @@ BS_LIST = [2**n for n in range(0, 15)]
BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 16399])
HIDDEN_DIMS = get_ci_test_range(
[64, 128, 256, 512, 1024, 96, 98, 100], [64, 512, 1024, 98]
[64, 128, 256, 512, 1024, 96, 97, 100], [64, 512, 1024, 97]
)
CACHE_SIZE = 1024 * 1024
DTYPE = torch.bfloat16
@@ -35,7 +35,6 @@ def test_store_cache(batch_size: int, element_dim: int) -> None:
v_cache = torch.randn((CACHE_SIZE, element_dim), dtype=DTYPE, device=DEVICE)
indices = torch.randperm(CACHE_SIZE - 1, device=DEVICE)[:batch_size] + 1
# AOT store cache
store_cache(k, v, k_cache, v_cache, indices)
assert torch.all(k_cache[indices] == k)
@@ -89,10 +88,7 @@ def test_store_cache_int32_indices(batch_size: int, element_dim: int) -> None:
@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64])
@pytest.mark.parametrize("num_split", [1, 2, 4])
def test_store_cache_reserved_skip_index(
index_dtype: torch.dtype, num_split: int
) -> None:
def test_store_cache_reserved_skip_index(index_dtype: torch.dtype) -> None:
element_dim = 1024
k = torch.randn((4, element_dim), dtype=DTYPE, device=DEVICE)
v = torch.randn((4, element_dim), dtype=DTYPE, device=DEVICE)
@@ -112,7 +108,6 @@ def test_store_cache_reserved_skip_index(
k_cache,
v_cache,
indices,
num_split=num_split,
)
torch.testing.assert_close(k_cache[0], reserved_k_before, rtol=0.0, atol=0.0)
@@ -137,43 +132,6 @@ def test_store_cache_zero_index_can_be_written_when_skip_disabled() -> None:
torch.testing.assert_close(v_cache[0], v[0], rtol=0.0, atol=0.0)
def _valid_num_splits(element_dim: int, dtype: torch.dtype) -> list:
"""Return the list of valid num_split values for a given element_dim/dtype."""
row_bytes = element_dim * dtype.itemsize
splits = [1]
if row_bytes % (2 * 128) == 0:
splits.append(2)
if row_bytes % (4 * 128) == 0:
splits.append(4)
return splits
_NUM_SPLIT_CASES = [
(_dim, _ns, _dtype)
for _dtype in [torch.float16, torch.bfloat16, torch.float32]
for _dim in REPR_DIMS
for _ns in _valid_num_splits(_dim, _dtype)
]
@pytest.mark.parametrize("element_dim,num_split,dtype", _NUM_SPLIT_CASES)
def test_store_cache_num_split(
element_dim: int, num_split: int, dtype: torch.dtype
) -> None:
batch_size = 128
k = torch.randn((batch_size, element_dim), dtype=dtype, device=DEVICE)
v = torch.randn((batch_size, element_dim), dtype=dtype, device=DEVICE)
k_cache = torch.randn((SMALL_CACHE, element_dim), dtype=dtype, device=DEVICE)
v_cache = torch.randn((SMALL_CACHE, element_dim), dtype=dtype, device=DEVICE)
indices = torch.randperm(SMALL_CACHE - 1, device=DEVICE)[:batch_size] + 1
# Verify each num_split kernel path (1, 2, 4) produces correct results
store_cache(k, v, k_cache, v_cache, indices, num_split=num_split)
assert torch.all(k_cache[indices] == k)
assert torch.all(v_cache[indices] == v)
# Asymmetric K/V (head_dim != v_head_dim): different row widths AND cache strides.
# MiMoV2 is 192/128. Both orderings, since nothing may assume K is the wider one.
ASYM_DIM_PAIRS = get_ci_test_range(
@@ -208,55 +166,6 @@ def test_store_cache_asymmetric(k_dim: int, v_dim: int, dtype: torch.dtype) -> N
assert torch.all(v_cache[untouched] == v_before[untouched])
def _valid_asym_num_splits(k_dim: int, v_dim: int, dtype: torch.dtype) -> list:
"""num_split values valid for BOTH rows; a split must divide each of them."""
k_bytes, v_bytes = k_dim * dtype.itemsize, v_dim * dtype.itemsize
splits = [1]
if k_bytes % (2 * 128) == 0 and v_bytes % (2 * 128) == 0:
splits.append(2)
if k_bytes % (4 * 128) == 0 and v_bytes % (4 * 128) == 0:
splits.append(4)
return splits
def _default_num_split(k_dim: int, v_dim: int, dtype: torch.dtype) -> int:
"""Mirrors the heuristic in store_cache(); the default is already exercised
by test_store_cache_asymmetric, which does not pass num_split."""
k_bytes, v_bytes = k_dim * dtype.itemsize, v_dim * dtype.itemsize
if k_bytes % 2048 == 0 and v_bytes % 2048 == 0:
return 4
if k_bytes % 1024 == 0 and v_bytes % 1024 == 0:
return 2
return 1
# Only splits the default heuristic would NOT pick: the split gate is two-sided
# (K and V must both align), so the off-default branches are what needs pinning.
_ASYM_NUM_SPLIT_CASES = [
(_k, _v, _ns)
for _k, _v in ASYM_DIM_PAIRS
for _ns in _valid_asym_num_splits(_k, _v, DTYPE)
if _ns != _default_num_split(_k, _v, DTYPE)
]
@pytest.mark.parametrize("k_dim,v_dim,num_split", _ASYM_NUM_SPLIT_CASES)
def test_store_cache_asymmetric_num_split(
k_dim: int, v_dim: int, num_split: int
) -> None:
batch_size = 128
k = torch.randn((batch_size, k_dim), dtype=DTYPE, device=DEVICE)
v = torch.randn((batch_size, v_dim), dtype=DTYPE, device=DEVICE)
k_cache = torch.randn((SMALL_CACHE, k_dim), dtype=DTYPE, device=DEVICE)
v_cache = torch.randn((SMALL_CACHE, v_dim), dtype=DTYPE, device=DEVICE)
indices = torch.randperm(SMALL_CACHE - 1, device=DEVICE)[:batch_size] + 1
store_cache(k, v, k_cache, v_cache, indices, num_split=num_split)
assert torch.all(k_cache[indices] == k)
assert torch.all(v_cache[indices] == v)
def test_can_use_store_cache() -> None:
assert can_use_store_cache(128)
assert can_use_store_cache(256)