[Perf] [OmniVoice] Triton kernel fusion + CUDA Graph acceleration (#3336)

Signed-off-by: univa-HARRY <harry@univa.co.kr>
Signed-off-by: Harry <harry@univa.co.kr>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/model_executor/models/omnivoice/test_cuda_graph_generator.py

@@ -0,0 +1,263 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for OmniVoice CUDA Graph generator wrapper numerical equivalence.
+
+Verifies that _OmniVoiceCUDAGraphForward produces results equivalent to
+eager mode across three scenarios:
+  - Exact-size inputs (no padding) → bit-identical
+  - Padded inputs (padded to nearest bucket) → correct slicing, exact match
+    for position-independent models
+  - Oversized inputs (fallback to eager) → bit-identical
+
+Uses a lightweight synthetic generator to keep warmup fast and avoid
+loading actual model weights.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+DEVICE = torch.device("cuda:0")
+NUM_CB = 8
+VOCAB = 1025
+HIDDEN = 64
+CAPTURE_SIZES = [32, 64, 128]
+
+
+# ---------------------------------------------------------------------------
+# Synthetic generator: matches _OmniVoiceCUDAGraphForward's interface
+# ---------------------------------------------------------------------------
+
+
+class _SyntheticGenerator:
+    """Position-independent Embedding+Linear replacing OmniVoiceGenerator's 28-layer transformer.
+
+    Each token position is processed independently (no attention), so padded
+    runs produce bit-identical results to non-padded runs on the same positions.
+    This makes the padding/slicing logic easy to verify exactly.
+    """
+
+    class _Cfg:
+        num_audio_codebook = NUM_CB
+
+    def __init__(self, device: torch.device):
+        self.config = self._Cfg()
+        self.text_embedding = nn.Embedding(1000, HIDDEN).to(device).eval()
+        self._linear = nn.Linear(HIDDEN, NUM_CB * VOCAB, bias=False).to(device).eval()
+        self._rope_cos: torch.Tensor | None = None
+        self._rope_sin: torch.Tensor | None = None
+
+    def _ensure_rope(self, seq_len: int, device: torch.device) -> None:
+        if self._rope_cos is None or self._rope_cos.shape[0] < seq_len:
+            max_len = max(seq_len, 512)
+            self._rope_cos = torch.zeros(max_len, 64, device=device)
+            self._rope_sin = torch.zeros(max_len, 64, device=device)
+
+    def _step_forward(
+        self,
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        two_b, _, S = input_ids.shape
+        x = self.text_embedding(input_ids[:, 0, :].clamp(0, 999))  # [two_b, S, H]
+        logits = self._linear(x)  # [two_b, S, 8*1025]
+        return logits.view(two_b, S, NUM_CB, VOCAB).permute(0, 2, 1, 3)  # [two_b, 8, S, 1025]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _eager(gen: _SyntheticGenerator, ids: torch.Tensor, mask: torch.Tensor, attn: torch.Tensor) -> torch.Tensor:
+    """Run _step_forward with zero RoPE tensors (synthetic gen ignores them)."""
+    S = ids.shape[-1]
+    cos = torch.zeros(S, 64, device=ids.device)
+    sin = torch.zeros(S, 64, device=ids.device)
+    return gen._step_forward(ids, mask, attn, cos, sin)
+
+
+def _make_inputs(seq_len: int, device: torch.device = DEVICE):
+    two_b = 2
+    ids = torch.randint(0, 100, (two_b, NUM_CB, seq_len), dtype=torch.long, device=device)
+    mask = torch.ones(two_b, seq_len, dtype=torch.bool, device=device)
+    attn = torch.ones(two_b, 1, seq_len, seq_len, dtype=torch.bool, device=device)
+    return ids, mask, attn
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def gen():
+    torch.manual_seed(42)
+    return _SyntheticGenerator(DEVICE)
+
+
+@pytest.fixture(scope="module")
+def wrapper(gen):
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+        _OmniVoiceCUDAGraphForward,
+    )
+
+    w = _OmniVoiceCUDAGraphForward(gen, capture_sizes=CAPTURE_SIZES)
+    w.warmup(DEVICE)
+    return w
+
+
+# ---------------------------------------------------------------------------
+# 1. Exact-size inputs → bit-identical
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("seq_len", CAPTURE_SIZES)
+def test_exact_size_bit_identical(gen, wrapper, seq_len):
+    """When input exactly matches a captured bucket, output must be bit-identical to eager."""
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        eager_out = _eager(gen, ids, mask, attn)
+        graph_out = wrapper(ids, mask, attn)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 2. Padded inputs → correct shape, values match at actual positions
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("seq_len", [1, 15, 33, 60, 100])
+def test_padded_output_shape(gen, wrapper, seq_len):
+    """Graph output must be sliced back to actual seq_len, not the bucket size."""
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        graph_out = wrapper(ids, mask, attn)
+    assert graph_out.shape == (2, NUM_CB, seq_len, VOCAB)
+
+
+@pytest.mark.parametrize("seq_len", [15, 33, 60, 100])
+def test_padded_output_matches_eager(gen, wrapper, seq_len):
+    """Padded graph output must equal eager output at actual positions.
+
+    The synthetic model has no attention across positions, so zero-padding
+    does not affect non-padded positions — exact match is expected.
+    """
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        eager_out = _eager(gen, ids, mask, attn)
+        graph_out = wrapper(ids, mask, attn)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 3. Oversized inputs → fallback to eager (lazy capture), bit-identical
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("seq_len", [129, 200, 256])
+def test_fallback_eager_bit_identical(gen, wrapper, seq_len):
+    """Sequences exceeding the largest bucket fall back to lazy capture → bit-identical."""
+    ids, mask, attn = _make_inputs(seq_len)
+    with torch.no_grad():
+        eager_out = _eager(gen, ids, mask, attn)
+        graph_out = wrapper(ids, mask, attn)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 4. Determinism
+# ---------------------------------------------------------------------------
+
+
+def test_deterministic_across_calls(wrapper):
+    """Same input must produce identical output on repeated CUDA graph replays."""
+    ids, mask, attn = _make_inputs(32)
+    with torch.no_grad():
+        out1 = wrapper(ids, mask, attn)
+        out2 = wrapper(ids, mask, attn)
+    torch.testing.assert_close(out1, out2, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 5. _find_bucket logic (CPU, no CUDA graph)
+# ---------------------------------------------------------------------------
+
+
+def test_find_bucket_returns_nearest_bucket():
+    """_find_bucket must return the smallest bucket >= seq_len, or None if all are smaller."""
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+        _OmniVoiceCUDAGraphForward,
+    )
+
+    w = _OmniVoiceCUDAGraphForward.__new__(_OmniVoiceCUDAGraphForward)
+    w._capture_sizes = [32, 64, 128]
+    w._graphs = {}
+
+    assert w._find_bucket(1) == 32
+    assert w._find_bucket(32) == 32
+    assert w._find_bucket(33) == 64
+    assert w._find_bucket(64) == 64
+    assert w._find_bucket(128) == 128
+    assert w._find_bucket(129) is None
+    assert w._find_bucket(1000) is None
+
+
+# ---------------------------------------------------------------------------
+# 6. enable_cuda_graph=False produces same tokens as enable_cuda_graph=True
+# ---------------------------------------------------------------------------
+
+
+def test_cuda_graph_disabled_matches_eager_generator():
+    """OmniVoiceGenerator with enable_cuda_graph=False must produce the same
+    _step_forward output as one with enable_cuda_graph=True (after warmup).
+
+    Uses a 2-layer config so warmup completes quickly without loading weights.
+    """
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator
+    from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
+
+    cfg = OmniVoiceConfig()
+    cfg.llm_num_hidden_layers = 2
+    cfg.num_hidden_layers = 2
+
+    torch.manual_seed(0)
+    gen_eager = OmniVoiceGenerator(cfg).to(DEVICE).eval()
+    gen_eager.config.enable_cuda_graph = False
+    gen_eager._cuda_graph_fwd = None
+
+    torch.manual_seed(0)
+    gen_graph = OmniVoiceGenerator(cfg).to(DEVICE).eval()
+    gen_graph.config.cuda_graph_capture_sizes = [64]
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import _OmniVoiceCUDAGraphForward
+
+    gen_graph._cuda_graph_fwd = _OmniVoiceCUDAGraphForward(gen_graph, capture_sizes=[64])
+    gen_graph._cuda_graph_fwd.warmup(DEVICE)
+
+    seq_len = 64
+    ids = torch.zeros(2, cfg.num_audio_codebook, seq_len, dtype=torch.long, device=DEVICE)
+    mask = torch.ones(2, seq_len, dtype=torch.bool, device=DEVICE)
+    attn = torch.ones(2, 1, seq_len, seq_len, dtype=torch.bool, device=DEVICE)
+
+    gen_eager._ensure_rope(seq_len, DEVICE)
+    dtype = gen_eager.text_embedding.weight.dtype
+    cos = gen_eager._rope_cos[:seq_len].to(device=DEVICE, dtype=dtype)
+    sin = gen_eager._rope_sin[:seq_len].to(device=DEVICE, dtype=dtype)
+
+    with torch.no_grad():
+        eager_logits = gen_eager._step_forward(ids, mask, attn, cos, sin)
+        graph_logits = gen_graph._cuda_graph_fwd(ids, mask, attn)
+
+    torch.testing.assert_close(graph_logits, eager_logits, atol=0, rtol=0)

tests/model_executor/models/omnivoice/test_decoder.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Correctness and regression tests for OmniVoiceDecoder.
+
+Covers three areas:
+1. NaN regression — verifies the fp16 overflow fix:
+     acoustic decoder must operate in float32 even when fc2 weights are fp16.
+     Demonstrates that the pre-fix code path (no .float() after fc2) produces
+     NaN/Inf with large-scale weights, while the fixed path does not.
+2. Output shape / value sanity — correct [B, 1, T*upsample] shape, finite values.
+3. Acoustic decoder float32 invariant — weights must be float32 after build.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import (
+    HiggsAudioRVQ,
+    OmniVoiceDecoder,
+)
+from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+DEVICE = torch.device("cuda:0")
+UPSAMPLE = 4  # synthetic acoustic decoder stride (real DAC uses 960)
+T_FRAMES = 25  # number of audio token frames to decode
+
+
+# ---------------------------------------------------------------------------
+# Tiny acoustic decoder: ConvTranspose1d upsample by UPSAMPLE
+# ---------------------------------------------------------------------------
+
+
+class _TinyAcousticDecoder(nn.Module):
+    def __init__(self, in_channels: int = 256, upsample: int = UPSAMPLE):
+        super().__init__()
+        self.conv = nn.ConvTranspose1d(in_channels, 1, kernel_size=upsample, stride=upsample)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)  # [B, 1, T*upsample]
+
+
+# ---------------------------------------------------------------------------
+# Builder
+# ---------------------------------------------------------------------------
+
+
+def _build_decoder(
+    *,
+    fc2_dtype: torch.dtype = torch.float32,
+    acoustic_dtype: torch.dtype = torch.float32,
+    weight_scale: float = 1.0,
+) -> OmniVoiceDecoder:
+    """Build an OmniVoiceDecoder with synthetic components (no model checkpoint required).
+
+    Args:
+        fc2_dtype: dtype for the fc2 linear layer weights.
+        acoustic_dtype: dtype for the acoustic decoder weights.
+        weight_scale: constant weight value for overflow testing.
+    """
+    config = OmniVoiceConfig()
+    decoder = OmniVoiceDecoder(config)
+
+    decoder.quantizer = HiggsAudioRVQ(num_quantizers=8, codebook_size=1024, codebook_dim=64, hidden_size=1024).to(
+        DEVICE
+    )
+
+    decoder.fc2 = nn.Linear(1024, 256).to(DEVICE).to(fc2_dtype)
+
+    acoustic = _TinyAcousticDecoder(in_channels=256, upsample=UPSAMPLE).to(DEVICE).to(acoustic_dtype)
+    if weight_scale != 1.0:
+        with torch.no_grad():
+            nn.init.constant_(acoustic.conv.weight, weight_scale)
+            nn.init.constant_(acoustic.conv.bias, weight_scale)
+    decoder.acoustic_decoder = acoustic
+
+    decoder._loaded = True
+    return decoder
+
+
+# ---------------------------------------------------------------------------
+# 1. NaN regression
+# ---------------------------------------------------------------------------
+
+
+def test_no_nan_fp16_fc2_float32_acoustic():
+    """fp16 fc2 + float32 acoustic decoder (the fixed path) must not produce NaN or Inf."""
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float32)
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert not torch.isnan(audio).any(), "NaN in output — fp16 overflow may have reached acoustic decoder"
+    assert not torch.isinf(audio).any(), "Inf in output"
+
+
+def test_fp16_acoustic_decoder_produces_nan_without_fix():
+    """Pre-fix scenario: fp16 acoustic decoder with large weights overflows to NaN/Inf.
+
+    This test documents the original bug and will fail (raising AssertionError)
+    if the underlying hardware / Triton version somehow avoids overflow, but
+    with weights scaled to 300.0 the ConvTranspose1d intermediate activations
+    reliably exceed fp16 max (~65504) in practice.
+    """
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float16, weight_scale=300.0)
+    with torch.no_grad():
+        nn.init.constant_(decoder.fc2.weight, 300.0)
+        nn.init.constant_(decoder.fc2.bias, 300.0)
+        for q in decoder.quantizer.quantizers:
+            nn.init.constant_(q.codebook.weight, 300.0)
+
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+
+    # Simulate pre-fix forward: no .float() after fc2
+    codes = tokens.transpose(0, 1).long()
+    quantized = decoder.quantizer.decode(codes)
+    quantized_fp16 = decoder.fc2(quantized.transpose(1, 2).to(decoder.fc2.weight.dtype)).transpose(
+        1, 2
+    )  # stays fp16, no .float()
+    audio_broken = decoder.acoustic_decoder(quantized_fp16)
+
+    assert torch.isnan(audio_broken).any() or torch.isinf(audio_broken).any(), (
+        "Expected NaN/Inf in pre-fix fp16 path with large weights — test setup may need a higher weight_scale"
+    )
+
+
+def test_acoustic_decoder_receives_float32_input():
+    """The .float() cast after fc2 must ensure acoustic decoder always gets float32 input,
+    regardless of fc2 weight dtype."""
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float32)
+
+    captured_dtypes: list[torch.dtype] = []
+    original_fwd = decoder.acoustic_decoder.forward
+
+    def _hook(x: torch.Tensor) -> torch.Tensor:
+        captured_dtypes.append(x.dtype)
+        return original_fwd(x)
+
+    decoder.acoustic_decoder.forward = _hook  # type: ignore[method-assign]
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+    decoder(tokens)
+
+    assert captured_dtypes, "acoustic_decoder.forward was never called"
+    assert captured_dtypes[0] == torch.float32, (
+        f"Acoustic decoder received {captured_dtypes[0]}, expected torch.float32 — "
+        "the .float() upcast after fc2 may be missing"
+    )
+
+
+# ---------------------------------------------------------------------------
+# 2. Output shape and value sanity
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("batch,t_frames", [(1, 10), (1, 50), (2, 25)])
+def test_output_shape(batch, t_frames):
+    """Decoder must return [B, 1, T*upsample] waveform tensor."""
+    decoder = _build_decoder()
+    tokens = torch.randint(0, 1024, (batch, 8, t_frames), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert audio.shape == (batch, 1, t_frames * UPSAMPLE), (
+        f"Expected ({batch}, 1, {t_frames * UPSAMPLE}), got {audio.shape}"
+    )
+
+
+def test_no_nan_float32_baseline():
+    """Float32 fc2 + float32 acoustic decoder must always produce finite output."""
+    decoder = _build_decoder(fc2_dtype=torch.float32, acoustic_dtype=torch.float32)
+    tokens = torch.randint(0, 1024, (1, 8, T_FRAMES), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert not torch.isnan(audio).any()
+    assert not torch.isinf(audio).any()
+
+
+def test_output_is_3d():
+    """Output must be 3-D: [B, 1, samples]. OmniVoiceDecoder.forward unsqueezes if needed."""
+    decoder = _build_decoder()
+    tokens = torch.randint(0, 1024, (1, 8, 10), device=DEVICE)
+
+    audio = decoder(tokens)
+
+    assert audio.dim() == 3, f"Expected 3-D output, got {audio.dim()}-D"
+    assert audio.shape[1] == 1, f"Expected channel dim=1, got {audio.shape[1]}"
+
+
+# ---------------------------------------------------------------------------
+# 3. Acoustic decoder float32 invariant
+# ---------------------------------------------------------------------------
+
+
+def test_acoustic_decoder_weights_are_float32():
+    """Acoustic decoder weights must be float32 — any fp16 parameter is a regression indicator.
+
+    In production, load_weights() calls self.acoustic_decoder.float() to enforce this.
+    The synthetic builder mirrors that invariant.
+    """
+    decoder = _build_decoder(fc2_dtype=torch.float16, acoustic_dtype=torch.float32)
+
+    for name, param in decoder.acoustic_decoder.named_parameters():
+        assert param.dtype == torch.float32, (
+            f"acoustic_decoder.{name} is {param.dtype} — must be float32 (see load_weights fix)"
+        )

tests/model_executor/models/omnivoice/test_triton_kernels.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Correctness tests for OmniVoice Triton kernels.
+
+Verifies that RMSNorm, SwiGLU, and fused-add+RMSNorm Triton kernels
+produce results numerically equivalent to the PyTorch reference.
+All computations follow the kernel's internal float32 promotion rules.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+DEVICE = torch.device("cuda:0")
+EPS = 1e-6
+
+try:
+    from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+        _TRITON_AVAILABLE,
+        triton_fused_add_rms_norm,
+        triton_rms_norm,
+        triton_swiglu,
+    )
+except ImportError:
+    _TRITON_AVAILABLE = False
+
+triton_available = pytest.mark.skipif(not _TRITON_AVAILABLE, reason="Triton not available on this platform")
+
+
+# ---------------------------------------------------------------------------
+# PyTorch reference implementations (mirror the kernel's internal arithmetic)
+# ---------------------------------------------------------------------------
+
+
+def _ref_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    # Kernel promotes x to float32, computes RMS, casts back to x.dtype, then mul weight
+    x_f32 = x.float()
+    ms = x_f32.pow(2).mean(-1, keepdim=True)
+    normed = (x_f32 * torch.rsqrt(ms + eps)).to(x.dtype)
+    return normed * weight.to(normed.dtype)
+
+
+def _ref_swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    # Kernel: gate → float32, silu, cast to up.dtype, mul up
+    gate_f32 = gate.float()
+    silu = gate_f32 * torch.sigmoid(gate_f32)
+    return silu.to(up.dtype) * up
+
+
+def _ref_fused_add_rms_norm(
+    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # Kernel: s_f32 = x+r in float32, s stored as x.dtype; y = norm(s_f32)*weight
+    s_f32 = x.float() + residual.float()
+    s = s_f32.to(x.dtype)
+    ms = s_f32.pow(2).mean(-1, keepdim=True)
+    normed = (s_f32 * torch.rsqrt(ms + eps)).to(x.dtype)
+    y = normed * weight.to(normed.dtype)
+    return y, s
+
+
+# ---------------------------------------------------------------------------
+# RMSNorm
+# ---------------------------------------------------------------------------
+
+
+@triton_available
+@pytest.mark.parametrize("rows,cols", [(1, 64), (4, 512), (16, 1024), (32, 3072)])
+def test_rms_norm_float32(rows, cols):
+    torch.manual_seed(0)
+    x = torch.randn(rows, cols, device=DEVICE)
+    w = torch.ones(cols, device=DEVICE)
+
+    ref = _ref_rms_norm(x, w, EPS)
+    out = triton_rms_norm(x, w, EPS)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+
+
+@triton_available
+@pytest.mark.parametrize("shape", [(2, 16, 64), (4, 32, 512)])
+def test_rms_norm_3d_input(shape):
+    """Triton wrapper must handle 3-D [B, S, H] inputs by flattening and restoring shape."""
+    torch.manual_seed(1)
+    hidden = shape[-1]
+    x = torch.randn(*shape, device=DEVICE)
+    w = torch.ones(hidden, device=DEVICE)
+
+    ref = _ref_rms_norm(x, w, EPS)
+    out = triton_rms_norm(x, w, EPS)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+    assert out.shape == x.shape
+
+
+@triton_available
+def test_rms_norm_no_nan_on_zero_input():
+    """RMSNorm of the zero vector must not produce NaN (eps prevents divide-by-zero)."""
+    x = torch.zeros(4, 256, device=DEVICE)
+    w = torch.ones(256, device=DEVICE)
+
+    out = triton_rms_norm(x, w, EPS)
+
+    assert not torch.isnan(out).any()
+
+
+# ---------------------------------------------------------------------------
+# SwiGLU
+# ---------------------------------------------------------------------------
+
+
+@triton_available
+@pytest.mark.parametrize("rows,cols", [(1, 64), (4, 512), (16, 3072)])
+def test_swiglu_float32(rows, cols):
+    torch.manual_seed(2)
+    gate = torch.randn(rows, cols, device=DEVICE)
+    up = torch.randn(rows, cols, device=DEVICE)
+
+    ref = _ref_swiglu(gate, up)
+    out = triton_swiglu(gate, up)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+
+
+@triton_available
+def test_swiglu_negative_gate():
+    """silu(x) = x*sigmoid(x) is defined for negative x; result must not be NaN."""
+    gate = torch.full((4, 64), -10.0, device=DEVICE)
+    up = torch.ones(4, 64, device=DEVICE)
+
+    ref = _ref_swiglu(gate, up)
+    out = triton_swiglu(gate, up)
+
+    torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+    assert not torch.isnan(out).any()
+
+
+@triton_available
+def test_swiglu_output_shape_preserved():
+    gate = torch.randn(3, 8, 512, device=DEVICE)
+    up = torch.randn(3, 8, 512, device=DEVICE)
+    out = triton_swiglu(gate, up)
+    assert out.shape == gate.shape
+
+
+# ---------------------------------------------------------------------------
+# Fused Add + RMSNorm
+# ---------------------------------------------------------------------------
+
+
+@triton_available
+@pytest.mark.parametrize("rows,cols", [(1, 64), (4, 512), (16, 1024)])
+def test_fused_add_rms_norm_float32(rows, cols):
+    torch.manual_seed(3)
+    x = torch.randn(rows, cols, device=DEVICE)
+    residual = torch.randn(rows, cols, device=DEVICE)
+    weight = torch.ones(cols, device=DEVICE)
+
+    ref_y, ref_s = _ref_fused_add_rms_norm(x, residual, weight, EPS)
+    out_y, out_s = triton_fused_add_rms_norm(x, residual, weight, EPS)
+
+    torch.testing.assert_close(out_y, ref_y, atol=1e-5, rtol=1e-5)
+    # Sum output s = x + r must be bit-exact (no numerics involved, just addition)
+    torch.testing.assert_close(out_s, ref_s, atol=0, rtol=0)
+
+
+@triton_available
+def test_fused_add_rms_norm_sum_is_bitexact():
+    """The residual sum s = x + r must be bit-exact between Triton and PyTorch."""
+    torch.manual_seed(4)
+    x = torch.randn(8, 512, device=DEVICE)
+    r = torch.randn(8, 512, device=DEVICE)
+    w = torch.ones(512, device=DEVICE)
+
+    _, ref_s = _ref_fused_add_rms_norm(x, r, w, EPS)
+    _, out_s = triton_fused_add_rms_norm(x, r, w, EPS)
+
+    torch.testing.assert_close(out_s, ref_s, atol=0, rtol=0)
+
+
+@triton_available
+def test_fused_add_rms_norm_returns_two_tensors():
+    """Wrapper must return exactly two tensors: (normed_output, sum)."""
+    x = torch.randn(4, 256, device=DEVICE)
+    r = torch.randn(4, 256, device=DEVICE)
+    w = torch.ones(256, device=DEVICE)
+
+    result = triton_fused_add_rms_norm(x, r, w, EPS)
+
+    assert len(result) == 2
+    y, s = result
+    assert y.shape == x.shape
+    assert s.shape == x.shape