tensorrt_llm/_torch/modules/swiglu.py (1)

54-58: All existing swiglu calls use torch.dtype – optional guard for Triton dtypes
All current invocations of swiglu(x, quant_scale, quant_type) pass a torch.dtype (or None), so no runtime dtype mismatches will occur. If you’d like to future-proof against someone accidentally passing a Triton dtype (e.g. tl.float8e4nv or tl.float8e5), you can optionally add:

• File: tensorrt_llm/_torch/modules/swiglu.py (around lines 51–58)
• Wrap the custom-op call’s dtype argument in a small converter:

@@ def swiglu(x, quant_scale: torch.Tensor = None, quant_type=None):
-        return torch.ops.trtllm.silu_and_mul(
-            x,
-            scale=quant_scale,
-            dtype=quant_type,
-        )
+        return torch.ops.trtllm.silu_and_mul(
+            x,
+            scale=quant_scale,
+            dtype=_as_torch_dtype(quant_type),
+        )

And add this helper at the bottom of the module:

def _as_torch_dtype(dt):
    import triton.language as tl, torch
    if dt is None or isinstance(dt, torch.dtype):
        return dt
    # Map Triton FP8→PyTorch FP8
    if dt is getattr(tl, "float8e4nv", None):
        return getattr(torch, "float8_e4m3fn", torch.float16)
    if dt is getattr(tl, "float8e5", None):
        return getattr(torch, "float8_e5m2", torch.bfloat16)
    return dt

No changes are required for correctness today—this is purely an optional refactor.

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (3)

913-940: Handle non-contiguous inputs or enforce contiguity explicitly.

The Triton kernel indexes the last dimension as if stride(1) == 1. If x is a non-contiguous view (e.g., sliced/transposed), the loads will read wrong elements. Either assert contiguity or make a safe fallback.

Apply this guard to avoid silent misreads (rare path does a copy; common path stays zero-copy):

 def silu_and_mul(x: torch.Tensor,
                  scale: Optional[torch.Tensor] = None,
                  dtype: Optional[torch.dtype] = None) -> torch.Tensor:
     b, n = x.shape

     assert n % 2 == 0
     d = n // 2

+    # Kernel assumes contiguous last-dimension indexing.
+    if x.stride(1) != 1:
+        x = x.contiguous()
+
     o_dtype = dtype or x.dtype
     o = torch.empty((b, d), dtype=o_dtype, device=x.device)

If you prefer to avoid the copy entirely, make it a hard check:

if x.stride(1) != 1:
    raise ValueError("trtllm::silu_and_mul expects x to be contiguous along the last dimension")

---

919-921: Prefer a ValueError over assert for shape checks.

Using assert for input validation can be stripped with Python optimizations, leading to hard-to-debug failures downstream. Raise an explicit ValueError instead.

-    assert n % 2 == 0
+    if (n % 2) != 0:
+        raise ValueError(f"trtllm::silu_and_mul expects even last dim; got n={n}")

---

913-955: File header missing: add NVIDIA copyright header.

Our guidelines require the NVIDIA copyright header (current year) on all source files.

Add at the top of the file:

# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)

925-927: Python 3.8 compatibility: use typing.Tuple instead of built-in tuple generics.

This is a SyntaxError on 3.8.

-    def grid(meta: Mapping[str, int]) -> tuple[int, int]:
+    def grid(meta: Mapping[str, int]) -> Tuple[int, int]:

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (4)

1-1: Add NVIDIA copyright header (2025) per repo guidelines.

The file is missing the required NVIDIA copyright header.

Apply at the very top:

+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

---

5-5: Optional: drop top-level Triton import by avoiding triton.cdiv in grid.

This keeps import-time dependencies lighter (helpful for fake/meta paths), and the integer expression is equivalent.

-import triton  # type: ignore[import]
@@
-    def grid(meta: Mapping[str, int]) -> Tuple[int, int]:
-        return (b, triton.cdiv(d, meta["BLOCK_SIZE"]))
+    def grid(meta: Mapping[str, int]) -> Tuple[int, int]:
+        # ceil_div(d, BLOCK_SIZE)
+        return (b, (d + meta["BLOCK_SIZE"] - 1) // meta["BLOCK_SIZE"])

Also applies to: 925-927

---

13-13: Prefer module-namespace import to avoid symbol-level coupling.

Keeps imports consistent with the guideline to maintain module namespaces, and reduces risk of future name collisions.

-from ..modules.swiglu import silu_and_mul_kernel
+from ..modules import swiglu as _swiglu
@@
-    silu_and_mul_kernel[grid](
+    _swiglu.silu_and_mul_kernel[grid](
         o_ptr=o,
         o_stride=o.stride(0),
         o_scale_ptr=scale,
         x_ptr=x,
         x_stride=x.stride(0),
         d=d,
         BLOCK_SIZE=1024,
         HAS_O_SCALE=scale is not None,
     )

Also applies to: 928-938

---

942-955: Meta (fake) impl: align with real op’s validation and docstring.

Mirror minimal validation and add a short docstring. This helps shape errors surface consistently under torch.compile.

 @silu_and_mul.register_fake
 def _(
     x: torch.Tensor,
     scale: Optional[torch.Tensor] = None,
     dtype: Optional[torch.dtype] = None,
 ) -> torch.Tensor:
-    b, n = x.shape
-
-    assert n % 2 == 0
-    d = n // 2
+    """Meta for silu_and_mul: returns an empty tensor with correct shape/dtype."""
+    if x.dim() != 2:
+        raise ValueError(f"silu_and_mul (fake): expected 2D input [B, N], got {tuple(x.shape)}")
+    b, n = x.shape
+    if n % 2 != 0:
+        raise ValueError(f"silu_and_mul (fake): expected even N, got N={n}")
+    d = n // 2
     o_dtype = dtype or x.dtype
     return x.new_empty((b, d), dtype=o_dtype)

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)

907-909: Python 3.8 compatibility: use typing.Tuple instead of built-in tuple generics.

tuple[int, int] requires Python 3.9+. We target 3.8+, so switch to Tuple[int, int] (already imported).

-    def grid(meta: Mapping[str, int]) -> tuple[int, int]:
+    def grid(meta: Mapping[str, int]) -> Tuple[int, int]:
         return (b, triton.cdiv(d, meta["BLOCK_SIZE"]))

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (4)

2-2: Minor: You import Tuple but still use built-in tuple generics later.

After fixing the grid annotation (Lines 997-999), Tuple becomes “used”; otherwise many linters will flag it as unused.

---

985-1011: Harden silu_and_mul: replace asserts with explicit validation, enforce CUDA + last-dim contiguity, require scalar scale, fix Python 3.8 tuple generics, and add a concise docstring.

• Asserts can be stripped under -O and give vague errors; raise clear exceptions instead.
• The kernel reads rows with pointer arithmetic and assumes x.stride(1) == 1; validate to avoid silent misreads or implicit copies.
• Kernel loads a single o_scale via tl.load(o_scale_ptr), so scale must be scalar (0-dim or 1-element).
• Python 3.8 target: change tuple[int, int] to Tuple[int, int] (Tuple already imported).
• Add a short Google-style docstring per guidelines.

Apply:

 @torch.library.custom_op("trtllm::silu_and_mul", mutates_args=())
 def silu_and_mul(x: torch.Tensor,
                  scale: Optional[torch.Tensor] = None,
                  dtype: Optional[torch.dtype] = None) -> torch.Tensor:
-    b, n = x.shape
-
-    assert n % 2 == 0
-    d = n // 2
-
-    o_dtype = dtype or x.dtype
-    o = torch.empty((b, d), dtype=o_dtype, device=x.device)
-
-    def grid(meta: Mapping[str, int]) -> tuple[int, int]:
-        return (b, triton.cdiv(d, meta["BLOCK_SIZE"]))
+    """SwiGLU activation backed by a Triton kernel with optional scalar scaling.
+
+    Args:
+        x: 2D CUDA tensor [B, N]. N must be even. Last dimension must be contiguous.
+        scale: Optional scalar CUDA tensor (0-dim or 1-element).
+        dtype: Optional output dtype. Defaults to x.dtype.
+
+    Returns:
+        Tensor [B, D] where D = N // 2, dtype = dtype or x.dtype.
+
+    Raises:
+        ValueError: on non-2D input, odd N, or non-contiguous last dim.
+        RuntimeError: on non-CUDA input or device mismatch for scale.
+        TypeError: if scale is provided but not scalar.
+    """
+    if x.dim() != 2:
+        raise ValueError(f"silu_and_mul: expected 2D input [B, N], got {tuple(x.shape)}")
+    if x.device.type != "cuda":
+        raise RuntimeError("silu_and_mul: expected CUDA tensor input.")
+    if x.stride(1) != 1:
+        raise ValueError("silu_and_mul: expected contiguous last dimension (x.stride(1) == 1).")
+    b, n = x.shape
+    if n % 2 != 0:
+        raise ValueError(f"silu_and_mul: expected even N, got N={n}")
+    d = n // 2
+    if scale is not None:
+        if scale.device != x.device:
+            raise RuntimeError("silu_and_mul: `scale` must be on the same device as `x`.")
+        if scale.numel() != 1:
+            raise TypeError("silu_and_mul: `scale` must be a scalar (0-dim or 1-element) tensor.")
+
+    o_dtype = dtype or x.dtype
+    o = torch.empty((b, d), dtype=o_dtype, device=x.device)
+
+    def grid(meta: Mapping[str, int]) -> Tuple[int, int]:
+        return (b, triton.cdiv(d, meta["BLOCK_SIZE"]))
 
     silu_and_mul_kernel[grid](
         o_ptr=o,
         o_stride=o.stride(0),
         o_scale_ptr=scale,
         x_ptr=x,
         x_stride=x.stride(0),
         d=d,
         BLOCK_SIZE=1024,
         HAS_O_SCALE=scale is not None,
     )
 
     return o

To quickly confirm there are no other Python 3.9+ generic annotations that could break 3.8:

#!/bin/bash
# Find Python 3.9+ built-in generic annotations that would break 3.8
rg -nP --type=py '\b(?:list|dict|set|tuple)\s*\[' -g '!**/third_party/**' -C2

---

997-999: Python 3.8 compatibility: change tuple[int, int] to typing.Tuple[int, int].

This file imports Tuple on Line 2 but still uses the 3.9+ form here. This will break 3.8 and may explain recurring pipeline failures. Apply:

-    def grid(meta: Mapping[str, int]) -> tuple[int, int]:
+    def grid(meta: Mapping[str, int]) -> Tuple[int, int]:

---

1014-1026: Mirror validation in the fake path; avoid bare asserts and provide a clear error.

Fake kernels participate in shape/dtype propagation during torch.compile, and early, cheap validation helps catch misuse without launching the real kernel.

Apply:

 @silu_and_mul.register_fake
 def _(
     x: torch.Tensor,
     scale: Optional[torch.Tensor] = None,
     dtype: Optional[torch.dtype] = None,
 ) -> torch.Tensor:
-    b, n = x.shape
-
-    assert n % 2 == 0
-    d = n // 2
+    if x.dim() != 2:
+        raise ValueError(f"silu_and_mul(fake): expected 2D input [B, N], got {tuple(x.shape)}")
+    b, n = x.shape
+    if n % 2 != 0:
+        raise ValueError(f"silu_and_mul(fake): expected even N, got N={n}")
+    d = n // 2
 
     o_dtype = dtype or x.dtype
     return x.new_empty((b, d), dtype=o_dtype)

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)

15-15: Reduce cross-module coupling: consider relocating the kernel into a dedicated kernels submodule.

Importing the kernel from ..modules.swiglu creates a subtle layering dependency between “custom_ops” and “modules”. A kernels subpackage (e.g., tensorrt_llm/_torch/kernels/swiglu.py) would make ownership clearer and avoid future import tangles.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.