diff --git a/benchmarks/kernels/benchmark_selective_state_update.py b/benchmarks/kernels/benchmark_selective_state_update.py index a8b73da2aa..ddf1efd018 100644 --- a/benchmarks/kernels/benchmark_selective_state_update.py +++ b/benchmarks/kernels/benchmark_selective_state_update.py @@ -34,6 +34,7 @@ from vllm.model_executor.layers.mamba.ops.mamba_ssm import ( override_ssm_config, selective_state_update, ) +from vllm.platforms import current_platform from vllm.triton_utils import triton # bf16 shares configs with fp16 - same bit width. @@ -93,10 +94,10 @@ def _make_inputs( ngroups: int, dtype: torch.dtype, state_dtype: torch.dtype | None = None, - device: str = "cuda", ): if state_dtype is None: state_dtype = dtype + device = current_platform.device_type state = torch.randn(batch, nheads, dim, dstate, dtype=state_dtype, device=device) x = torch.randn(batch, nheads, dim, dtype=dtype, device=device) dt = torch.randn(batch, nheads, dim, dtype=dtype, device=device) @@ -127,9 +128,9 @@ def benchmark_config( Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update. Returns elapsed time in microseconds, or None on error. - Uses CUDA graph capture-and-replay to isolate kernel time from Python - eager-mode dispatch / kwarg-resolution overhead, mirroring the timing - methodology in benchmarks/kernels/benchmark_moe.py. + Uses accelerator graph capture-and-replay to isolate kernel time from + Python eager-mode dispatch / kwarg-resolution overhead, mirroring the + timing methodology in benchmarks/kernels/benchmark_moe.py. """ state, x, dt, A, B, C, D, dt_bias, out = _make_inputs( batch, nheads, dim, dstate, ngroups, dtype, state_dtype=state_dtype @@ -157,10 +158,15 @@ def benchmark_config( _call_kernel() torch.accelerator.synchronize() - # Capture graph_batch_size invocations into a CUDA graph so the + # Capture graph_batch_size invocations into a device graph so the # timed region runs without Python dispatch overhead per call. - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): + # Capture via graph(), not the graph object: CUDA needs a side stream. + graph = ( + torch.cuda.CUDAGraph() + if current_platform.is_cuda_alike() + else torch.xpu.XPUGraph() + ) + with current_platform.graph(graph): for _ in range(graph_batch_size): _call_kernel() torch.accelerator.synchronize() @@ -170,8 +176,8 @@ def benchmark_config( graph.replay() torch.accelerator.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) latencies: list[float] = [] for _ in range(num_iters): start.record() @@ -671,8 +677,8 @@ def main(): dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 state_dtype = _SSM_CACHE_DTYPE_MAP[args.mamba_ssm_cache_dtype] device_name = get_ssm_device_name() - cap = torch.cuda.get_device_capability() - is_blackwell = cap[0] >= 10 + cap = current_platform.get_device_capability() + is_blackwell = cap is not None and cap[0] >= 10 # Mirror all output to a results file (like Unix tee). buf = StringIO() @@ -690,7 +696,8 @@ def main(): sys.stdout = _Tee() # type: ignore[assignment] try: - print(f"Device : {device_name} (sm_{cap[0]}{cap[1]})") + cap_str = f"sm_{cap[0]}{cap[1]}" if cap is not None else "n/a" + print(f"Device : {device_name} ({cap_str})") print(f"Blackwell: {is_blackwell}") print(f"dtype : {args.dtype}") print(f"ssm_cache_dtype: {args.mamba_ssm_cache_dtype}") diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json new file mode 100644 index 0000000000..edc77b17c5 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json @@ -0,0 +1,51 @@ +{ + "triton_version": "3.7.1", + "8": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json new file mode 100644 index 0000000000..6e01b1b9a6 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json @@ -0,0 +1,51 @@ +{ + "triton_version": "3.7.1", + "8": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "16": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "256": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "512": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "1024": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json new file mode 100644 index 0000000000..8735bb1c6f --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.7.1", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "32768": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "65536": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "131072": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json new file mode 100644 index 0000000000..a382f024f9 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.7.1", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "65536": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "131072": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + } +} \ No newline at end of file