[None][test] AutoDeploy: Fix standalone linear simple on B200 (#15675)

Signed-off-by: gramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

Author: govind-ramnarayan

tensorrt_llm/_torch/auto_deploy/custom_ops/linear/linear.py

@@ -15,13 +15,21 @@
 
 """Custom ops for linear layers."""
 
-from typing import List, Optional
+from typing import Callable, List, Optional
 
 import torch
 
 from ..._compat import get_sm_version
 
 
+def _get_trtllm_cublas_mm() -> Optional[Callable[..., torch.Tensor]]:
+    """Return TRT-LLM's cublas_mm op when it is registered in this process."""
+    try:
+        return torch.ops.trtllm.cublas_mm
+    except (AttributeError, RuntimeError):
+        return None
+
+
 @torch.library.custom_op("auto_deploy::torch_linear_simple", mutates_args=())
 def simple(
     input: torch.Tensor,
@@ -72,11 +80,22 @@ def simple(
     # split-K + reduce + zero-fill for small-M (decode) projection GEMMs.
     # (Same trick PT introduced for GPT-OSS via use_custom_cublas_mm in
     # modeling_gpt_oss.py; we apply it model-agnostically based on dtype + SM.)
-    if get_sm_version() >= 100 and input.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16:
+    if (
+        input.is_cuda
+        and weight.is_cuda
+        and (bias is None or bias.is_cuda)
+        and input.dtype == torch.bfloat16
+        and weight.dtype == torch.bfloat16
+        and get_sm_version() >= 100
+    ):
+        trtllm_cublas_mm = _get_trtllm_cublas_mm()
+        if trtllm_cublas_mm is None:
+            return torch.ops.aten.linear(input, weight, bias)
+
         # cublas_mm requires 2D mat_a/mat_b. Flatten leading dims and unflatten on exit.
         in_shape = input.shape
         input_2d = input.reshape(-1, in_shape[-1])
-        out_2d = torch.ops.trtllm.cublas_mm(
+        out_2d = trtllm_cublas_mm(
             input_2d,
             weight.t(),
             bias,

tests/unittest/auto_deploy/standalone/test_standalone_package.py

@@ -141,6 +141,40 @@ def test_ops_registered(self, standalone_package):
         )
         assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}"
 
+    def test_linear_simple_works_in_standalone_on_b200(self, standalone_package):
+        """The standalone linear op must work on the B200 bf16 path."""
+        result = _run_isolated(
+            standalone_package,
+            """
+            from unittest import mock
+
+            import torch
+            import llmc
+            from llmc.custom_ops.linear import linear
+
+            if torch.cuda.is_available():
+                device = "cuda"
+                dtype = torch.bfloat16
+            else:
+                device = "cpu"
+                dtype = torch.float32
+
+            x = torch.randn(2, 3, 5, device=device, dtype=dtype)
+            weight = torch.randn(7, 5, device=device, dtype=dtype)
+            bias = torch.randn(7, device=device, dtype=dtype)
+
+            # Force the Blackwell bf16 branch even when this standalone test
+            # runs on H100 CI instead of B200 CI.
+            with mock.patch.object(linear, "get_sm_version", return_value=100):
+                out = torch.ops.auto_deploy.torch_linear_simple(x, weight, bias)
+
+            ref = torch.ops.aten.linear(x, weight, bias)
+            torch.testing.assert_close(out, ref)
+            print("OK")
+            """,
+        )
+        assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}"
+
     def test_compat_types(self, standalone_package):
         """Standalone _compat types should work."""
         result = _run_isolated(