main
py 24 lines 705 Bytes
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """AMD/ROCm fused Triton ops for MiniMax-M3.
4
5 These replace per-element PyTorch fallbacks (FlashInfer / fused HIP kernels are
6 unavailable on ROCm) with single-pass Triton kernels to cut launch overhead and
7 intermediate-tensor traffic during decode.
8 """
9
10 from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import (
11 gemma_fused_add_rmsnorm,
12 gemma_rmsnorm,
13 )
14 from vllm.models.minimax_m3.amd.ops.swiglu_oai import (
15 swiglu_oai_quantize_mxfp8,
16 swiglu_oai_split,
17 )
18
19 __all__ = [
20 "gemma_rmsnorm",
21 "gemma_fused_add_rmsnorm",
22 "swiglu_oai_split",
23 "swiglu_oai_quantize_mxfp8",
24 ]