`aten._grouped_mm` converter raises `TypeError` when the offsets are a constant

Open Beginner friendly
  • Is something wrong?
#5,038 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 3 days

Nobody has claimed this yet.

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
74/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python
Domain
compilers, machine-learning

Research direction

Start in the grouped_mm converter in torch_tensorrt/dynamo/conversion/impl/moe.py, where offs is passed to cast_trt_tensor in converter_utils.py, and check how enforce_tensor_types handles argument 2. Run the repro in the issue with a constant offs buffer and confirm the TypeError. Done when the constant-offsets case compiles and matches PyTorch within the stated tolerance, and the existing grouped_mm tests still pass. Coordinate with PR #5021 so the two fixes do not conflict.

Written by the indexing model from the issue text.

Description

What is wrong

In mixture-of-experts models the group offsets of torch._grouped_mm are sometimes fixed, for example a buffer that says how many rows each expert gets. After freezing, the offsets reach the converter as a constant torch.Tensor, not a TensorRT tensor. The converter only forces its first 2 arguments to TensorRT tensors (enforce_tensor_types({0: ..., 1: ...})), then passes offs to cast_trt_tensor, which expects an ITensor. The conversion fails with a TypeError inside that cast, so the grouped matmul does not run in TensorRT.

What this blocks

cosmos3 (pytorch/executorch#24353) hits this: its grouped matmul stays out of TensorRT, so the graph is split around it into more TensorRT engines and CUDA backend programs.
A fix keeps the grouped matmul inside the engine, so mixture-of-experts layers with fixed groups run fully in TensorRT.

Repro
import torch
import torch_tensorrt


class FixedGroups(torch.nn.Module):
    def __init__(self, groups=4, k=64, n=32, rows_per_group=8):
        super().__init__()
        self.w = torch.nn.Parameter(torch.randn(groups, k, n))
        ends = torch.arange(1, groups + 1, dtype=torch.int32) * rows_per_group
        self.register_buffer("offs", ends)

    def forward(self, x):
        return torch._grouped_mm(x, self.w, offs=self.offs)


model = FixedGroups().cuda().bfloat16().eval()
x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)
ep = torch.export.export(model, (x,))
trt_gm = torch_tensorrt.dynamo.compile(ep, arg_inputs=[x], min_block_size=1)
torch.testing.assert_close(trt_gm(x), model(x), rtol=2e-2, atol=2e-2)
Expected

The grouped matmul converts and matches PyTorch.

Actual

dynamo.compile raises inside the converter (end of the traceback):

  File ".../torch_tensorrt/dynamo/conversion/impl/moe.py", line 93, in grouped_mm
    offs_i32 = cast_trt_tensor(
  File ".../torch_tensorrt/dynamo/conversion/converter_utils.py", line 276, in cast_trt_tensor
    cast_layer = ctx.net.add_cast(input_val, trt_dtype)
TypeError: add_cast(): incompatible function arguments. The following argument types are supported:
    1. (self: tensorrt_bindings.tensorrt.INetworkDefinition, input: tensorrt_bindings.tensorrt.ITensor, to_type: tensorrt_bindings.tensorrt.DataType) -> tensorrt_bindings.tensorrt.ICastLayer
Invoked with: <tensorrt_bindings.tensorrt.INetworkDefinition object at 0x7f4dc10b3d30>, tensor([ 8, 16, 24, 32], dtype=torch.int32), <DataType.INT32: 3>
While executing %_grouped_mm : [num_users=1] = call_function[target=torch.ops.aten._grouped_mm.default](args = (%x, %w, %offs), kwargs = {})
Versions

torch 2.15.0 (CUDA 13.2), torch_tensorrt 2.15.0, TensorRT 11.4.0.106, H100, driver 580.126.09.

What was tried

Reading the converter. A likely fix is to add 2: (TRTTensor,) to the converter's enforce_tensor_types, or to call get_trt_tensor on offs before the cast. The new grouped GEMM kernel path in the open PR #5021 already does the second; the default dense path does not.

Dominant language
Python
Stars
3k
Forks
421
Avg merge
4d 7h
Merged PRs (30d)
62

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from pytorch/TensorRT

All issues in pytorch/TensorRT

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.