[Kernels] Bitwise-exact fused RMSNorm, RoPE and SiLU*up for CUDA graph decode

tests/test_normalization.py

@@ -69,3 +69,114 @@ def test_rmsnorm_explicit_epsilon():
     # Variance of ones is 1.0; 1.0 / sqrt(1.0 + 1.0) = 1.0 / sqrt(2) ≈ 0.7071
     expected_large = 1.0 / (2.0**0.5)
     assert torch.allclose(out_large, torch.full_like(x, expected_large), atol=1e-4)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("width", [2048, 5280, 10560])
+def test_fused_rmsnorm_matches_reference_within_rounding(dtype, width):
+    from vllm_rlt.kernels.norm import fused_rmsnorm, fused_rmsnorm_supported
+
+    torch.manual_seed(0)
+    norm = RMSNorm(width).to(device="cuda", dtype=dtype)
+    with torch.no_grad():
+        norm.weight.copy_(torch.randn(width, device="cuda") * 0.5 + 1.0)
+    hidden = (torch.randn(37, width, device="cuda") * 3).to(dtype)
+    assert fused_rmsnorm_supported(hidden)
+    fused = fused_rmsnorm(hidden, norm.weight, norm.variance_epsilon)
+    reference = norm.reference(hidden)
+    assert fused.dtype == dtype and fused.shape == reference.shape
+    # Same arithmetic and rounding points; only the FP32 reduction order differs. A
+    # flipped normalized value can move the weighted output by up to two ulps
+    # (two chained roundings); the vast majority of elements are identical.
+    if dtype == torch.float32:
+        assert torch.allclose(fused, reference, rtol=1e-5, atol=1e-6)
+    else:
+        diff = (fused.float() - reference.float()).abs()
+        ulp = torch.finfo(dtype).eps * reference.float().abs().clamp_min(torch.finfo(dtype).tiny)
+        assert (diff <= 2.01 * ulp).all()
+        assert (diff == 0).float().mean().item() > 0.99
+
+
+@pytest.mark.gpu
+def test_fused_rmsnorm_residual_epilogue_rounds_once():
+    """The approx kernel's residual epilogue computes round(residual + w * normalized).
+
+    torch's separate add rounds the weight product first. The two oracles differ
+    in ~18% of elements, by more than one ulp of the result under cancellation, so
+    the check is membership: every element equals the single-rounded oracle (or,
+    in rare fp16 cases, the double-rounded one).
+    """
+    from vllm_rlt.kernels.norm import fused_rmsnorm
+
+    torch.manual_seed(1)
+    norm = RMSNorm(2048).to(device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        # Unit weights would make the product exact and hide the epilogue's rounding.
+        norm.weight.copy_(torch.randn(2048, device="cuda") * 0.5 + 1.0)
+    hidden = (torch.randn(9, 2048, device="cuda") * 3).to(torch.bfloat16)
+    residual = (torch.randn(9, 2048, device="cuda") * 3).to(torch.bfloat16)
+    ones = torch.ones_like(norm.weight)
+    normalized = fused_rmsnorm(hidden, ones, norm.variance_epsilon)  # round(x * inv), exact weight
+    plain = fused_rmsnorm(hidden, norm.weight, norm.variance_epsilon)
+    fused = fused_rmsnorm(hidden, norm.weight, norm.variance_epsilon, residual)
+    product = norm.weight.float() * normalized.float()
+    single = (residual.float() + product).to(torch.bfloat16)
+    double = residual + plain
+    assert torch.equal(plain, product.to(torch.bfloat16))
+    assert ((fused == single) | (fused == double)).all()
+    assert (fused == single).float().mean().item() > 0.99
+    assert not torch.equal(single, double)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("width", [2048, 5280, 10560])
+@pytest.mark.parametrize("with_residual", [False, True])
+@pytest.mark.parametrize("fuse_inverse", [False, True])
+def test_fused_rmsnorm_exact_is_bitwise_identical_to_reference(
+    dtype, width, with_residual, fuse_inverse
+):
+    from vllm_rlt.kernels.norm import fused_rmsnorm_exact
+
+    torch.manual_seed(2)
+    norm = RMSNorm(width).to(device="cuda", dtype=dtype)
+    with torch.no_grad():
+        norm.weight.copy_(torch.randn(width, device="cuda") * 0.5 + 1.0)
+    hidden = (torch.randn(37, width, device="cuda") * 3).to(dtype)
+    residual = (torch.randn(37, width, device="cuda") * 3).to(dtype) if with_residual else None
+    reference = norm.reference(hidden)
+    if residual is not None:
+        reference = residual + reference
+    fused = fused_rmsnorm_exact(
+        hidden, norm.weight, norm.variance_epsilon, residual, fuse_inverse=fuse_inverse
+    )
+    assert fused.dtype == dtype and fused.shape == reference.shape
+    assert torch.equal(fused, reference)
+    # The row statistic is torch's own, so the exact mode stays bitwise under any width.
+    assert torch.equal(
+        fused_rmsnorm_exact(
+            hidden[:1], norm.weight, norm.variance_epsilon, fuse_inverse=fuse_inverse
+        ),
+        norm.reference(hidden[:1]),
+    )
+
+
+def test_fused_scope_graph_uses_reference_outside_capture(monkeypatch):
+    """Default scope: fused kernels only during CUDA graph capture; eager uses the reference."""
+    import vllm_rlt.kernels.dispatch as dispatch
+    import vllm_rlt.layers.normalization as normalization
+
+    calls = []
+    monkeypatch.setattr(normalization, "FUSED_NORM", "exact")
+    monkeypatch.setattr(
+        normalization.RMSNorm, "reference", lambda self, h: calls.append("ref") or h
+    )
+    norm = RMSNorm(8)
+    hidden = torch.ones(2, 8)
+    monkeypatch.setattr(dispatch, "FUSED_SCOPE", "graph")
+    assert dispatch.fused_allowed() is False  # no capture on CPU
+    norm(hidden)
+    assert calls == ["ref"]
+    monkeypatch.setattr(dispatch, "FUSED_SCOPE", "always")
+    assert dispatch.fused_allowed() is True

tests/test_ouro.py

@@ -299,3 +299,41 @@ def test_dtype_conversion_preserves_float32_rotary_frequencies(model):
     angles = torch.cat([angles, angles], dim=-1)
     torch.testing.assert_close(cos[:, 0], angles.cos().bfloat16(), atol=0, rtol=0)
     torch.testing.assert_close(sin[:, 0], angles.sin().bfloat16(), atol=0, rtol=0)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_fused_silu_mul_tracks_reference(dtype):
+    import torch.nn.functional as F
+
+    from vllm_rlt.kernels.activation import fused_silu_mul, fused_silu_mul_supported
+
+    torch.manual_seed(0)
+    gate = (torch.randn(64, 5632, device="cuda") * 4).to(dtype)
+    up = (torch.randn(64, 5632, device="cuda") * 2).to(dtype)
+    assert fused_silu_mul_supported(gate, up)
+    fused = fused_silu_mul(gate, up)
+    reference = F.silu(gate) * up
+    diff = (fused.float() - reference.float()).abs()
+    ulp = torch.finfo(dtype).eps * reference.float().abs().clamp_min(torch.finfo(dtype).tiny)
+    assert (diff <= 1.01 * ulp).all()
+    assert (diff == 0).float().mean().item() > 0.999
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+def test_fused_silu_mul_exact_matches_reference(dtype):
+    """Forced rounding points: the only possible difference left is the transcendental."""
+    import torch.nn.functional as F
+
+    from vllm_rlt.kernels.activation import fused_silu_mul
+
+    torch.manual_seed(0)
+    gate = (torch.randn(64, 5632, device="cuda") * 4).to(dtype)
+    up = (torch.randn(64, 5632, device="cuda") * 2).to(dtype)
+    fused = fused_silu_mul(gate, up, exact=True)
+    reference = F.silu(gate) * up
+    mismatches = int((fused != reference).sum())
+    # libdevice exp is the CUDA math library PyTorch links, so this is expected to be
+    # zero; the bound keeps the test informative if a platform's exp differs.
+    assert mismatches <= 2, mismatches

tests/test_rotary_embedding.py

@@ -90,3 +90,67 @@ def test_rotary_embedding_forward_precision_and_shapes(dtype):
     assert sin.shape == (len(positions), 1, head_dim)
     assert cos.dtype == dtype
     assert sin.dtype == dtype
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("heads,dim", [(16, 128), (55, 96)])
+def test_fused_rope_approx_rounds_the_sum_once(dtype, heads, dim):
+    from vllm_rlt.kernels.rope import fused_rope_, fused_rope_supported
+    from vllm_rlt.layers.rotary_embedding import apply_rotary_pos_emb_reference
+
+    if (dim // 2) & (dim // 2 - 1):
+        pytest.skip("fused RoPE needs a power-of-two half dimension")
+    torch.manual_seed(0)
+    rope = RotaryEmbedding(head_dim=dim).to("cuda")
+    rows = 37
+    hidden = torch.zeros(rows, 8, device="cuda", dtype=dtype)
+    positions = torch.randint(0, 60000, (rows,), device="cuda")
+    cos, sin = rope(hidden, positions)
+    q = (torch.randn(rows, heads, dim, device="cuda") * 2).to(dtype)
+    k = (torch.randn(rows, heads, dim, device="cuda") * 2).to(dtype)
+    assert fused_rope_supported(q, cos)
+    expected_q, expected_k = apply_rotary_pos_emb_reference(q, k, cos, sin)
+    fused_q, fused_k = fused_rope_(q.clone(), cos, sin), fused_rope_(k.clone(), cos, sin)
+    # The compiler is free to keep either product unrounded (FP32, or fused into an FMA
+    # with the other term), so each element is one of four roundings of
+    # q*cos + rotate_half(q)*sin: both products exact, either one rounded first, or
+    # both rounded (the reference). GH200 showed ~9% of elements outside the first and
+    # last, i.e. a mixed form. Under cancellation these differ by more than any fixed
+    # number of result ulps, so the contract is membership in that set.
+    from vllm_rlt.layers.rotary_embedding import rotate_half
+
+    for value, fused, expected in ((q, fused_q, expected_q), (k, fused_k, expected_k)):
+        p1 = value.float() * cos.float()
+        p2 = rotate_half(value).float() * sin.float()
+        r1, r2 = p1.to(dtype).float(), p2.to(dtype).float()
+        candidates = [(p1 + p2), (r1 + p2), (p1 + r2)]
+        member = fused == expected
+        for c in candidates:
+            member |= fused == c.to(dtype)
+        assert member.float().mean().item() > 0.999
+        assert (fused == expected).float().mean().item() > 0.7
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("heads,dim", [(16, 128), (4, 64)])
+def test_fused_rope_exact_is_bitwise_identical_to_reference(dtype, heads, dim):
+    from vllm_rlt.kernels.rope import fused_rope_, fused_rope_supported
+    from vllm_rlt.layers.rotary_embedding import apply_rotary_pos_emb_reference
+
+    torch.manual_seed(0)
+    rope = RotaryEmbedding(head_dim=dim).to("cuda")
+    rows = 37
+    hidden = torch.zeros(rows, 8, device="cuda", dtype=dtype)
+    positions = torch.randint(0, 60000, (rows,), device="cuda")
+    cos, sin = rope(hidden, positions)
+    q = (torch.randn(rows, heads, dim, device="cuda") * 2).to(dtype)
+    k = (torch.randn(rows, heads, dim, device="cuda") * 2).to(dtype)
+    assert fused_rope_supported(q, cos)
+    expected_q, expected_k = apply_rotary_pos_emb_reference(q, k, cos, sin)
+    fused_q = fused_rope_(q.clone(), cos, sin, exact=True)
+    fused_k = fused_rope_(k.clone(), cos, sin, exact=True)
+    # Forced rounding points reproduce the reference's three roundings exactly.
+    assert torch.equal(fused_q, expected_q)
+    assert torch.equal(fused_k, expected_k)

vllm_rlt/kernels/activation.py

@@ -1,4 +1,17 @@
-"""MLP activation operators and their unfused correctness reference."""
+"""MLP activation operators and their unfused correctness reference.
+
+``silu_and_mul`` dispatches to the optional native CUDA kernel or the Torch
+reference. ``fused_silu_mul`` is the fused Triton SiLU(gate) * up used while
+fusion is allowed (see ``vllm_rlt.kernels.dispatch``).
+
+The reference ``F.silu(gate) * up`` is two kernels, each rounding to the activation dtype.
+This kernel computes ``silu`` in FP32 as ``x / (1 + exp(-x))`` with libdevice ``exp`` and
+IEEE division (the same formula PyTorch's CUDA kernel uses), rounds, multiplies by ``up``
+and rounds again. The approximate kernel (VLLM_RLT_FUSED_ACT=1) lets the compiler elide the
+intermediate rounding; the exact kernel (VLLM_RLT_FUSED_ACT=exact) forces it through inline
+PTX so the only remaining difference could be the transcendental's last fp32 ulp (libdevice
+``exp`` is the CUDA math library PyTorch links). The GPU test reports the identical fraction.
+"""
 
 import torch
 import torch.nn.functional as F
@@ -37,3 +50,59 @@ def silu_and_mul_reference(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor
 _library.impl("silu_and_mul", silu_and_mul_reference, "CompositeExplicitAutograd")
 
 silu_and_mul = torch.ops.vllm_rlt.silu_and_mul
+
+try:
+    import triton
+    import triton.language as tl
+    from triton.language.extra import libdevice
+
+    from vllm_rlt.kernels.rounding import mul_round, round_to
+except ImportError:  # pragma: no cover - CPU-only installs
+    triton = None
+
+
+if triton is not None:
+
+    @triton.jit
+    def _silu_mul_kernel(gate_ptr, up_ptr, out_ptr, numel, BLOCK: tl.constexpr):
+        offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+        mask = offsets < numel
+        dtype = out_ptr.dtype.element_ty
+        g = tl.load(gate_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        u = tl.load(up_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        act = tl.math.div_rn(g, 1.0 + libdevice.exp(-g)).to(dtype).to(tl.float32)
+        tl.store(out_ptr + offsets, (act * u).to(dtype), mask=mask)
+
+    @triton.jit
+    def _silu_mul_exact_kernel(gate_ptr, up_ptr, out_ptr, numel, BLOCK: tl.constexpr):
+        offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+        mask = offsets < numel
+        dtype = out_ptr.dtype.element_ty
+        g = tl.load(gate_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        u = tl.load(up_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        # silu rounded to dtype (PyTorch's F.silu output), then the dtype product.
+        act = round_to(tl.math.div_rn(g, 1.0 + libdevice.exp(-g)), dtype)
+        tl.store(out_ptr + offsets, mul_round(act, u, dtype).to(dtype), mask=mask)
+
+
+def fused_silu_mul_supported(gate: torch.Tensor, up: torch.Tensor) -> bool:
+    return (
+        triton is not None
+        and gate.is_cuda
+        and gate.dtype in (torch.bfloat16, torch.float16, torch.float32)
+        and gate.dtype == up.dtype
+        and gate.shape == up.shape
+        and gate.is_contiguous()
+        and up.is_contiguous()
+    )
+
+
+def fused_silu_mul(gate: torch.Tensor, up: torch.Tensor, *, exact: bool = False) -> torch.Tensor:
+    out = torch.empty_like(gate)
+    numel = gate.numel()
+    if numel == 0:
+        return out
+    block = 1024
+    kernel = _silu_mul_exact_kernel if exact else _silu_mul_kernel
+    kernel[((numel + block - 1) // block,)](gate, up, out, numel, BLOCK=block, num_warps=4)
+    return out

vllm_rlt/kernels/dispatch.py

@@ -0,0 +1,34 @@
+"""Where fused elementwise kernels run: inside CUDA graph capture only, or everywhere.
+
+A fused Triton kernel replaces several PyTorch kernels, but each Triton launch from Python
+costs more host time than a PyTorch op. Inside a captured CUDA graph that cost is paid
+once at capture; in eager execution (prefill, Ouro prelude/coda, sync eager decode) it is
+paid every call, and on GH200's Grace CPU it outweighs the saved launches: sync eager
+Ouro-1.4B decode ran at 0.40x and prefill steps took 33-46% longer with the fused set on.
+
+VLLM_RLT_FUSED_SCOPE=graph (default) uses the fused kernels only while a CUDA graph is
+being captured; =always also uses them eagerly. The exact kernels are bitwise identical
+to the eager reference, so the scope never changes numerics in exact mode.
+"""
+
+import os
+
+import torch
+
+FUSED_SCOPE = os.environ.get("VLLM_RLT_FUSED_SCOPE", "graph")
+if FUSED_SCOPE not in ("graph", "always"):
+    raise ValueError("VLLM_RLT_FUSED_SCOPE must be 'graph' or 'always'")
+
+
+_CUDA = None
+
+
+def fused_allowed() -> bool:
+    """Whether a fused kernel may replace the eager reference at this call site."""
+    global _CUDA
+    if FUSED_SCOPE == "always":
+        return True
+    if _CUDA is None:
+        # CPU-only builds have no capture query at all; nothing is ever captured there.
+        _CUDA = torch.cuda.is_available()
+    return _CUDA and torch.cuda.is_current_stream_capturing()