Repository navigation
feat(conversion): run aten._grouped_mm in TensorRT with a grouped GEMM kernel - #5021
Open
shoumikhin wants to merge 4 commits into
Open
shoumikhin wants to merge 4 commits into
shoumikhin wants to merge 4 commits into
Conversation
…he dense lowering The aten._grouped_mm converter multiplies every row by the weight of every group and then masks away the unwanted results. A mixture-of-experts layer with 128 experts therefore does 128 times the work and builds a [groups, rows, N] temporary tensor. For large layers this is far slower than PyTorch's grouped matmul, which only computes each row's own group, and it can need more GPU memory than the device has. A new setting, grouped_mm_dense_max_macs (default 2**38), makes grouped_mm_validator decline the op when the dense lowering would do more multiply-adds than that (groups * rows * K * N). Above 2**38 the dense engine was slower than PyTorch on every GPU measured (H100 and A100). Below it nothing changes, so small MoE graphs still stay in one engine. The op is still converted when it has one group (no wasted work), under require_full_compilation=True, and when the setting is None. convert_exported_program_to_serialized_trt_engine builds one engine with no PyTorch fallback, so it always converts the op. Settings saved before this change load with None, so refit rebuilds the same engines. A declined op runs in PyTorch under torch_tensorrt.compile and torch.compile, and one warning per node names the setting. The setting rejects booleans and non-positive values. When saving for ExecuTorch it needs a partitioner that takes it, such as the CUDA partitioner. PyTorch's grouped matmul copies the offsets to the host when it has no fast kernel (float16, float32, GPUs before Hopper), so such a module cannot be captured whole by CUDA graphs. The docs name None for both cases.
…imit CUDA lets a kernel use more than 48 KiB of dynamic shared memory only after the launcher opts the kernel in with cuFuncSetAttribute. TensorRT launches AOT plugin kernels without doing that (pytorch#5010). So a triton_op or cutedsl_op kernel that needs more than 48 KiB builds into an engine without complaint, and then every run of the engine fails with a generic plugin assertion that gives no reason. Large Triton tiles with software pipelining often need more than 48 KiB, so this is easy to hit. The capability validators of triton_op and cutedsl_op now decline such a kernel, the same way they decline hardware-compatible engines. The op and its eager_fn stay registered, TensorRT leaves the op to PyTorch, and a warning names the cause and the fix (fewer pipeline stages or smaller tiles). An op without an eager_fn could not run anywhere, so for it registration fails with the same explanation instead. Kernels at or below 48 KiB are unchanged. ptx_op and cuda_kernel_op callers with their own aot_fn are not checked: TensorRT launches their kernels the same way, so their shared_mem must stay at or below 48 KiB too. The docs now say so for both. Signed-off-by: Anthony Shoumikhin <shoumikhin@meta.com>
…ent hint When Triton launches a kernel the normal way, it looks at each pointer argument and tells the compiler which ones are 16-byte aligned. That hint is what lets the compiler use 16-byte vector loads, asynchronous copies and software pipelining. triton_op compiles the kernel ahead of time, before any pointer exists, so it never passes the hint, and every triton_op kernel gets scalar memory accesses only. Such a kernel can run several times slower inside a TensorRT engine than the same kernel launched by Triton. triton_op now takes pointer_alignment, with the same meaning and default as in cutedsl_op: the byte alignment the kernel may assume for every pointer, or a mapping from pointer names to alignments. Any alignment above 1 becomes Triton's divisibility hint for that pointer, capped at 16 as Triton's own launcher does, and names that are not pointer parameters are rejected. The default of 1 adds no hint, so existing registrations compile exactly as before. Whoever registers the op guarantees the declared alignment for every buffer TensorRT binds to it, engine inputs included; an offset view whose byte offset is not a multiple of it does not meet it, and the resulting misaligned access is a sticky CUDA error. Values must be powers of two from 1 to 256. Both frontends now share one pointer_alignment parser. Signed-off-by: Anthony Shoumikhin <shoumikhin@meta.com>
…M kernel Mixture-of-experts layers sort their rows by expert and call aten._grouped_mm, where each row is multiplied only by its own expert's weight. The existing converter lowers this densely: every row is multiplied by every expert's weight and the unwanted results are masked away. With 128 experts that is 128 times the work, so large layers are left to PyTorch, and each one splits the graph into more engines. The dense converter also needs a static row count, so with a dynamic number of tokens the op always runs in PyTorch. This adds a second converter for the op, opt-in through a new setting, grouped_mm_kernel (default False), and tried first when set. It inserts a TensorRT plugin that runs a Triton grouped matrix multiply. The kernel reads the group offsets on the GPU, so nothing is copied to the host, and it computes each row only against its own group's weight, with float32 accumulation as in PyTorch. The kernel is compiled ahead of time when the engine is built, through torch_tensorrt.kernels, and its PTX is stored in the engine. Running the engine needs no Triton. Rows past the last offset belong to no group, and the kernel writes zeros there, as the dense converter does, so the output never holds stale memory. With grouped_mm_kernel=True the new converter is used when all of these hold: bfloat16 or float16 matrices, int32 offsets, the (rows, K) x (groups, K, N) form without bias, out_dtype unset or equal to the input dtype, more than one group, static K, N and group count (the row count may be dynamic), a GPU with compute capability 8.0 or newer, TensorRT plugins, Triton 3.5 or newer and cuda-python installed, an engine that is not hardware compatible, and no Windows cross-compilation. cuda-python reads the driver's PTX limit; without it the PTX could be too new for the driver, so the kernel declines. The setting is off by default because an engine with the kernel needs a CUDA driver at least as new as the build machine's. Otherwise the dense converter and grouped_mm_dense_max_macs apply as before. When the setting is on but the GPU, the installed packages or other settings rule the kernel out, one warning says why. The dense converter's warning suggests the setting only when it is off and the kernel would likely take the op, checked without starting Triton. cross_compile_for_windows does not get the setting, since the kernel never runs there. The validator compiles the kernel, so a kernel that cannot be built for a plugin launch declines the op instead of failing the engine build. The result, success or failure, is cached per shape and GPU, and a failure is logged once as a warning. Offsets outside mat1, past its end or negative, are invalid, but the kernel never reads or writes outside mat1 or the output for them. The kernel uses 128 x 128 x 32 tiles. TensorRT gives a plugin kernel at most 48 KiB of shared memory today, so it uses 3 pipeline stages where they fit and 2 otherwise. From sm_100 on, Triton's default code for tl.dot uses tensor memory, which a plugin launch cannot set up, so there the kernel is compiled with sm_80 instructions (mma.sync). The PTX still targets the device, and compile_triton_to_ptx gets an arch argument for this. On those GPUs the kernel can be slower than PyTorch's grouped matrix multiply at large row counts, and the docs say to measure it. The kernel declares mat1, mat2 and the output 16-byte aligned (pointer_alignment=16), so Triton uses 16-byte loads and pipelining. TensorRT allocates its own tensors 256-byte aligned and requires engine inputs and outputs to be bound 256-byte aligned. The offsets get no hint: a load of a few offsets gains nothing from one. The tests hold one case per behaviour. The converter tests check which converter took each op. The validator tests check that the kernel is opt-in, the dtype, bias, out_dtype, offsets and one-group conditions, the warning when Triton, the driver's PTX limit, the GPU or the settings rule the kernel out, and when the dense warning suggests the setting. The kernel tests skip at run time whenever the validator would decline for the environment. Kernel tests compile it for sm_90 and sm_100 on any GPU, check the pipeline-stage fallback and the shared memory error, and cover a dynamic row count in one engine, a compile failure that falls back and is tried once, and engine-only exports run in a fresh process that loads only TensorRT and PyTorch: one checks the result, and one checks that offsets past the row count or negative never read or write outside mat1 or the output and that rows past the last offset are zero. An ExecuTorch test runs a .pte file with the kernel in a fresh process that loads neither Triton nor Torch-TensorRT. Signed-off-by: Anthony Shoumikhin <shoumikhin@meta.com>
github-actions
Bot
added
documentation
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 13, 2026
…s in TensorRT Torch-TensorRT runs `aten._grouped_mm`, the grouped matrix multiply of mixture-of-experts layers, inside TensorRT only with its opt-in `grouped_mm_kernel` setting (pytorch/TensorRT#5021). Without it, each such layer stays outside TensorRT, and Praxis-1's `velocity` program splits into many TensorRT engines with CUDA backend parts between them. A recipe had no way to pass a setting to the backend's compiler. `Part` gets `compile_options`, settings of the backend's compiler for that part alone, empty by default. They go to the compile step as JSON, so a part refuses values that JSON cannot hold. The trace step hands them to the compile step of the part's program, which passes them to the backend's `compile` as the keyword argument `compile_options`, and only for a part that has some. So a backend whose `compile` takes only the program and the path, in LeRobot or in a `lerobot_backend_<name>` package, keeps working unchanged. Only `executorch_tensorrt` takes them. It adds them to its own Torch-TensorRT settings. It refuses a setting that the installed Torch-TensorRT does not have, since Torch-TensorRT would ignore it without a word. It also refuses one that it already passes, from its own settings, the user's `--backend` flags or the arguments of its call, so a recipe never replaces them silently. The Praxis-1 recipe sets `grouped_mm_kernel` on `velocity` and `decoder`, the parts that hold the mixture-of-experts layers, only when the backend is `executorch_tensorrt`. With that backend it also checks, before any weight loads, that Torch-TensorRT is installed and that its `CompilationSettings` has the setting. It checks the setting itself rather than a version. A Torch-TensorRT without the setting is refused with a ValueError that names the installed version, pytorch/TensorRT#5021 and Torch-TensorRT 2.15.1. A missing Torch-TensorRT gets the same install hint that the compile step gives, which moves into `ExportBackendConfig.install_error` for both. Tested on CPU with PyTorch 2.15 and ExecuTorch 1.6: - a backend whose `compile` takes only the program and the path still exports; - a part's settings reach, through the trace and compile steps, a backend that takes them; in an export of a tiny Praxis-1, only `velocity` and `decoder` get them, and the other programs get no keyword argument; - the TensorRT backend hands a part's settings to a fake Torch-TensorRT with its own, and refuses one that the fake does not have, and one that it already passes; - a tiny Praxis-1's `velocity` and `decoder` ask for `grouped_mm_kernel` with `executorch_tensorrt`, and no part asks for anything with `executorch_cuda`; - with `executorch_tensorrt`, the recipe refuses before the policy loads a fake Torch-TensorRT without the setting and a missing one, and goes on with one that has it; with `executorch_cuda` it checks nothing; - a part refuses compile options that JSON cannot hold. Each of these tests fails with its part of the change removed. The export and Praxis-1 tests pass.
shoumikhin
force-pushed
the
grouped-mm-kernel
branch
from
November 14, 2026 17:24
c91fe1a to
595bc7e
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Depends on #5014, #5016 and #5017.
What is wrong today
In a mixture-of-experts (MoE) layer,
aten._grouped_mmmultiplies each row by its ownexpert's weight. Torch-TensorRT lowers it densely, every row times every expert's weight,
so large layers and dynamic row counts go to PyTorch and split the graph.
What this change does
A second converter, opt-in and tried first, inserts a TensorRT plugin that runs a Triton grouped
matrix multiply. It computes only each row's own group, in float32, and reads
offson the GPU.Its PTX is stored in the engine, so running it needs no Triton. A kernel that cannot be built
falls back to the dense converter.
New setting:
grouped_mm_kernel, defaultFalse.Truekeeps large and dynamic-row MoElayers in TensorRT, so a MoE program can build as one engine.
The cost: the engine needs a driver at least as new as the build machine's.
It is used for bfloat16 or float16 with int32
offs, no bias, more than one group, fixedK,Nand group count (rows may be dynamic), sm_80 or newer, and Triton 3.5+ and cuda-pythonat build time. When the GPU, the packages, a
hardware-compatible engine or Windows cross-compilation rule it out, one warning says why.
On sm_100+ it runs sm_80 instructions and can be slower than PyTorch, so measure it.
Results
One MoE layer (128 experts, dim 3072, top-8, bfloat16) exported to ExecuTorch, kernel on,
median of 50 runs, on an earlier version of this change:
Tests
test_moe_aten.py: both converters, the kernel's conditions and warnings.test_000_grouped_mm_kernel.py: dynamic rows, empty groups, bad offsets, a failedcompile, shared memory limits, sm_90 and sm_100 code, and the TensorRT runtime alone.
executorch/test_grouped_mm_kernel.py: a.pterun without Triton by ExecuTorch.Ran on H100 and A100 from this branch. The ExecuTorch test ran on H100 with a release/2.15
build of the same change; its runtime package is not built for main yet.