A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 19:00 UTC
[Perf] Add opt-in SM100 Oink RMSNorm custom-op path (#31828)
Signed-off-by: Laura Wang <3700467+Laurawly@users.noreply.github.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
Laura Wang committed
Feb 24, 2026 at 23:01 UTC
24650715105a26b41c88f4164777b5e1a0f3200b
4 files changed
+331
tests/model_executor/test_oink_integration.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import types
5
+
6
+import pytest
7
+import torch
8
+
9
+
10
+def _load_oink_ops_module():
11
+ # Import the module normally (vllm is installed as an editable package in CI).
12
+ from vllm import _oink_ops
13
+
14
+ return _oink_ops
15
+
16
+
17
+def test_oink_availability_checks(monkeypatch: pytest.MonkeyPatch):
18
+ _oink_ops = _load_oink_ops_module()
19
+
20
+ # Ensure the ops namespace exists and is mutable for tests.
21
+ monkeypatch.setattr(
22
+ torch.ops,
23
+ "oink",
24
+ types.SimpleNamespace(rmsnorm=lambda x, w, eps: x),
25
+ raising=False,
26
+ )
27
+
28
+ # Case 1: CUDA not available.
29
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
30
+ assert _oink_ops.is_oink_available_for_device(0) is False
31
+
32
+ # Case 2: CUDA available but < SM100.
33
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
34
+ monkeypatch.setattr(torch.cuda, "get_device_capability", lambda idx: (9, 0))
35
+ assert _oink_ops.is_oink_available_for_device(0) is False
36
+
37
+ # Case 3: CUDA available and SM100, rmsnorm op registered.
38
+ monkeypatch.setattr(torch.cuda, "get_device_capability", lambda idx: (10, 0))
39
+ assert _oink_ops.is_oink_available_for_device(0) is True
40
+
41
+ # fused op presence probe
42
+ assert _oink_ops.has_fused_add_rms_norm() is False
43
+ monkeypatch.setattr(
44
+ torch.ops,
45
+ "oink",
46
+ types.SimpleNamespace(
47
+ rmsnorm=lambda x, w, eps: x,
48
+ fused_add_rms_norm=lambda x, residual, w, eps: None,
49
+ ),
50
+ raising=False,
51
+ )
52
+ assert _oink_ops.has_fused_add_rms_norm() is True
53
+
54
+
55
+def test_can_view_as_2d_stride_guard():
56
+ # Import the helper from the layernorm module.
57
+ from vllm.model_executor.layers.layernorm import _can_view_as_2d
58
+
59
+ x = torch.zeros((2, 3, 4))
60
+ assert _can_view_as_2d(x) is True
61
+
62
+ # Size-1 dims should be ignored by the viewability check.
63
+ # Create a tensor where stride(0) != stride(1) * size(1) due to padding,
64
+ # but view(-1, H) is still valid because dim 1 has size 1.
65
+ base = torch.zeros((2, 10, 4))
66
+ x_singleton = base[:, :1, :]
67
+ x_singleton.view(-1, x_singleton.shape[-1])
68
+ assert _can_view_as_2d(x_singleton) is True
69
+
70
+ # Middle-dimension stride break: view(-1, hidden) should be invalid.
71
+ x2 = x[:, ::2, :]
72
+ with pytest.raises(RuntimeError):
73
+ x2.view(-1, x2.shape[-1])
74
+ assert _can_view_as_2d(x2) is False
vllm/_oink_ops.py
new
+96
@@ -0,0 +1,96 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Small helper wrappers for external Oink Blackwell custom ops.
4
+
5
+vLLM does not depend on the external Oink repository/package. When an external
6
+plugin registers torch.library.custom_op entrypoints under the `oink::`
7
+namespace (e.g. via vLLM's general_plugins mechanism) and
8
+`VLLM_USE_OINK_OPS=1` is set, vLLM can route eligible calls to those ops.
9
+
10
+This module provides:
11
+- A single place to probe Oink op availability at module init time
12
+ (outside torch.compile tracing), and
13
+- Thin wrappers around the torch.ops entrypoints for use in CUDA fast paths,
14
+ without introducing graph breaks.
15
+
16
+Important:
17
+ Do not call the availability helpers in a compiled region. They may call
18
+ functions decorated with `torch._dynamo.disable` to safely check
19
+ conditions that should not be traced.
20
+"""
21
+
22
+from __future__ import annotations
23
+
24
+from collections.abc import Callable
25
+
26
+import torch
27
+
28
+try:
29
+ from torch._dynamo import disable as _dynamo_disable # type: ignore[attr-defined]
30
+except Exception: # pragma: no cover
31
+
32
+ def _dynamo_disable(fn: Callable): # type: ignore[misc]
33
+ return fn
34
+
35
+
36
+def _has_oink_op(op_name: str) -> bool:
37
+ """Check if a specific oink op is registered."""
38
+ return hasattr(torch.ops, "oink") and hasattr(torch.ops.oink, op_name)
39
+
40
+
41
+@_dynamo_disable
42
+def is_oink_available_for_device(device_index: int) -> bool:
43
+ """Return True if Oink ops are registered and device is SM100+.
44
+
45
+ This function is intended to be called during module initialization
46
+ (e.g., in RMSNorm.__init__), not in the forward path.
47
+
48
+ External plugins are expected to gate registration on SM100+ and
49
+ VLLM_USE_OINK_OPS=1, so if the ops are present they should be usable.
50
+ """
51
+ if not torch.cuda.is_available():
52
+ return False
53
+
54
+ try:
55
+ major, minor = torch.cuda.get_device_capability(device_index)
56
+ sm = 10 * major + minor
57
+ if sm < 100:
58
+ return False
59
+ except Exception:
60
+ return False
61
+
62
+ return _has_oink_op("rmsnorm")
63
+
64
+
65
+def has_fused_add_rms_norm() -> bool:
66
+ """Return True if the in-place fused op is registered."""
67
+ return _has_oink_op("fused_add_rms_norm")
68
+
69
+
70
+def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
71
+ """Call `torch.ops.oink.rmsnorm`.
72
+
73
+ This wrapper is safe to call in torch.compile regions.
74
+ """
75
+ return torch.ops.oink.rmsnorm(x, weight, eps)
76
+
77
+
78
+def fused_add_rms_norm_(
79
+ x: torch.Tensor,
80
+ residual: torch.Tensor,
81
+ weight: torch.Tensor,
82
+ eps: float,
83
+) -> None:
84
+ """Call `torch.ops.oink.fused_add_rms_norm` (mutates x and residual)."""
85
+ torch.ops.oink.fused_add_rms_norm(x, residual, weight, eps)
86
+
87
+
88
+def fused_add_rms_norm(
89
+ x: torch.Tensor,
90
+ residual: torch.Tensor,
91
+ weight: torch.Tensor,
92
+ eps: float,
93
+) -> tuple[torch.Tensor, torch.Tensor]:
94
+ """Convenience wrapper returning (x, residual) after in-place mutation."""
95
+ fused_add_rms_norm_(x, residual, weight, eps)
96
+ return x, residual
vllm/envs.py
+6
@@ -97,6 +97,7 @@ if TYPE_CHECKING:
97
VLLM_SKIP_P2P_CHECK: bool = False
98
VLLM_DISABLED_KERNELS: list[str] = []
99
VLLM_DISABLE_PYNCCL: bool = False
100
+ VLLM_USE_OINK_OPS: bool = False
101
VLLM_ROCM_USE_AITER: bool = False
102
VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False
103
VLLM_ROCM_USE_AITER_LINEAR: bool = True
@@ -896,6 +897,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
897
"VLLM_DISABLE_PYNCCL": lambda: (
898
os.getenv("VLLM_DISABLE_PYNCCL", "False").lower() in ("true", "1")
899
),
900
+ # Optional: enable external Oink custom ops (e.g., Blackwell RMSNorm).
901
+ # Disabled by default.
902
+ "VLLM_USE_OINK_OPS": lambda: (
903
+ os.getenv("VLLM_USE_OINK_OPS", "False").lower() in ("true", "1")
904
+ ),
905
# Disable aiter ops unless specifically enabled.
906
# Acts as a parent switch to enable the rest of the other operations.
907
"VLLM_ROCM_USE_AITER": lambda: (
vllm/model_executor/layers/layernorm.py
+155
@@ -6,7 +6,9 @@ import torch
6
import torch.nn as nn
7
import torch.nn.functional as F
8
9
+from vllm import _oink_ops, envs
10
from vllm._aiter_ops import rocm_aiter_ops
11
+from vllm.logger import init_logger
12
from vllm.model_executor.custom_op import CustomOp
13
from vllm.model_executor.layers.batch_invariant import (
14
rms_norm_batch_invariant,
@@ -14,6 +16,41 @@ from vllm.model_executor.layers.batch_invariant import (
16
)
17
from vllm.platforms import current_platform
18
19
+logger = init_logger(__name__)
20
+
21
+
22
+def _can_view_as_2d(x: torch.Tensor) -> bool:
23
+ """Return True if x.view(-1, x.shape[-1]) is viewable (no copy)."""
24
+ if x.dim() < 2:
25
+ return False
26
+ if x.dim() == 2:
27
+ return True
28
+ # For a view(-1, N) to be valid, all leading dims must be contiguous with
29
+ # respect to each other (size-1 dims are ignored).
30
+ for dim in range(x.dim() - 1):
31
+ # Strides for size-1 dims are irrelevant and can be arbitrary.
32
+ if x.size(dim + 1) != 1 and x.stride(dim) != x.stride(dim + 1) * x.size(
33
+ dim + 1
34
+ ):
35
+ return False
36
+ return True
37
+
38
+
39
+def _is_oink_stride_compatible_2d(x_2d: torch.Tensor) -> bool:
40
+ """Return True if x_2d meets Oink's pointer-path stride constraints."""
41
+ if x_2d.dim() != 2:
42
+ return False
43
+ if x_2d.stride(1) != 1:
44
+ return False
45
+ # Match Oink's vectorization constraint: stride(0) divisible by 256b.
46
+ if x_2d.dtype in (torch.float16, torch.bfloat16):
47
+ divby = 16
48
+ elif x_2d.dtype == torch.float32:
49
+ divby = 8
50
+ else:
51
+ return False
52
+ return (x_2d.stride(0) % divby) == 0
53
+
54
55
def rms_norm(
56
x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
@@ -131,6 +168,57 @@ class RMSNorm(CustomOp):
168
with_fused_add=True, dtype=weight_dtype, use_aiter=aiter_rmsnorm_enabled
169
)
170
171
+ # Optional: enable Oink Blackwell RMSNorm custom-op fast path on
172
+ # compatible CUDA devices (e.g., SM100) when the external Oink
173
+ # package is available. This is detected once at construction time
174
+ # to avoid per-call device queries in the hot path.
175
+ self._use_oink_rmsnorm = False
176
+ self._use_oink_fused_add_rmsnorm = False
177
+ if (
178
+ not current_platform.is_rocm()
179
+ and torch.cuda.is_available()
180
+ and bool(getattr(envs, "VLLM_USE_OINK_OPS", False))
181
+ ):
182
+ # NOTE: vLLM disables custom ops by default when using Inductor.
183
+ # If this op is disabled, CustomOp will dispatch to forward_native,
184
+ # and the Oink path in forward_cuda will never run.
185
+ if getattr(self._forward_method, "__func__", None) is getattr(
186
+ self.forward_native, "__func__", None
187
+ ):
188
+ try:
189
+ from vllm.config import get_cached_compilation_config
190
+
191
+ custom_ops = get_cached_compilation_config().custom_ops
192
+ except Exception:
193
+ custom_ops = ["<unknown>"]
194
+ logger.warning_once(
195
+ "VLLM_USE_OINK_OPS=1 but the `rms_norm` custom op is "
196
+ "disabled (CompilationConfig.custom_ops=%s). Enable it via "
197
+ "`compilation_config={'custom_ops': ['none', '+rms_norm']}` "
198
+ "(or `['all']`) to let vLLM call into torch.ops.oink.*.",
199
+ custom_ops,
200
+ )
201
+ # Custom op disabled => forward_cuda won't run. Avoid doing any
202
+ # external Oink initialization work in this case.
203
+ else:
204
+ try:
205
+ device_index = torch.cuda.current_device()
206
+ if _oink_ops.is_oink_available_for_device(device_index):
207
+ self._use_oink_rmsnorm = True
208
+ self._use_oink_fused_add_rmsnorm = (
209
+ _oink_ops.has_fused_add_rms_norm()
210
+ )
211
+ except Exception as e:
212
+ # If anything goes wrong (no Oink install, CPU-only env, etc.),
213
+ # silently fall back to the built-in RMSNorm path.
214
+ logger.warning_once(
215
+ "VLLM_USE_OINK_OPS=1 but failed to initialize Oink "
216
+ "RMSNorm; falling back to vLLM RMSNorm. Error: %s",
217
+ e,
218
+ )
219
+ self._use_oink_rmsnorm = False
220
+ self._use_oink_fused_add_rmsnorm = False
221
+
222
@staticmethod
223
def forward_static(
224
x: torch.Tensor,
@@ -202,6 +290,73 @@ class RMSNorm(CustomOp):
290
if self.variance_size_override is not None:
291
return self.forward_native(x, residual)
292
293
+ # Optional Oink SM100 fast path (no residual). This path is
294
+ # torch.compile-friendly via torch.ops.oink.rmsnorm and preserves
295
+ # 2D layouts (including padded rows) when using the Oink
296
+ # pointer-based kernel.
297
+ if (
298
+ residual is None
299
+ and getattr(self, "_use_oink_rmsnorm", False)
300
+ and x.is_cuda
301
+ and x.dim() >= 2
302
+ and self.has_weight
303
+ and not vllm_is_batch_invariant()
304
+ and self.weight.data.dtype == x.dtype
305
+ and self.weight.data.is_contiguous()
306
+ ):
307
+ orig_shape = x.shape
308
+ hidden_size = orig_shape[-1]
309
+ if _can_view_as_2d(x):
310
+ x_2d = x.view(-1, hidden_size)
311
+ if _is_oink_stride_compatible_2d(x_2d):
312
+ y_2d = _oink_ops.rmsnorm(
313
+ x_2d,
314
+ self.weight.data,
315
+ self.variance_epsilon,
316
+ )
317
+ return y_2d.view(orig_shape)
318
+
319
+ # Optional Oink SM100 fast path (fused residual-add + RMSNorm, in-place).
320
+ # This mirrors vLLM's fused_add_rms_norm semantics by mutating both
321
+ # `x` (normalized output) and `residual` (residual-out buffer).
322
+ if (
323
+ residual is not None
324
+ and getattr(self, "_use_oink_fused_add_rmsnorm", False)
325
+ and x.is_cuda
326
+ and residual.is_cuda
327
+ and x.shape == residual.shape
328
+ and x.dtype == residual.dtype
329
+ and x.dim() >= 2
330
+ and self.has_weight
331
+ and not vllm_is_batch_invariant()
332
+ and self.weight.data.dtype == x.dtype
333
+ and self.weight.data.is_contiguous()
334
+ ):
335
+ orig_shape = x.shape
336
+ hidden_size = orig_shape[-1]
337
+ if _can_view_as_2d(x) and _can_view_as_2d(residual):
338
+ x_2d = x.view(-1, hidden_size)
339
+ res_2d = residual.view(-1, hidden_size)
340
+
341
+ # The Oink in-place pointer path supports the common vLLM
342
+ # layout where:
343
+ # - `x` may be strided/padded row-major (stride(1) == 1), and
344
+ # - `residual` is contiguous row-major ([M, N] with stride(0) == N).
345
+ # If these conditions are not met, fall back to vLLM's built-in
346
+ # fused kernel.
347
+ if (
348
+ _is_oink_stride_compatible_2d(x_2d)
349
+ and _is_oink_stride_compatible_2d(res_2d)
350
+ and res_2d.is_contiguous()
351
+ ):
352
+ _oink_ops.fused_add_rms_norm_(
353
+ x_2d,
354
+ res_2d,
355
+ self.weight.data,
356
+ self.variance_epsilon,
357
+ )
358
+ return x, residual
359
+
360
add_residual = residual is not None
361
if add_residual:
362
return fused_add_rms_norm(