[None][perf] Fuse MiniMax-M3 prefill projections

cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2025-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,8 @@
 #pragma once
 
 #include "tensorrt_llm/common/config.h"
+
+#include <cstdint>
 #include <cuda_runtime.h>
 
 TRTLLM_NAMESPACE_BEGIN
@@ -71,6 +73,24 @@ void launchFusedQKNormRopeOut(void const* qkv_in, // BF16 input [num_tokens, tot
     bool const interleave, int const* position_ids, float factor, float low, float high, float attention_factor,
     cudaStream_t stream, bool is_qk_norm, bool use_gemma, bool use_mrope, int mrope_section1, int mrope_section2);
 
+// MiniMax-M3-specific producer for eager pure prefill. It returns contiguous
+// FP8 Q and inserts normalized/RoPE'd FP8 K plus copy-cast FP8 V directly into
+// a paged HND pool [num_pages, 2, num_heads, page_size, head_dim].
+void launchMinimaxM3Fp8QKNormRopeKVInsert(void const* qkv_input, void* q_output, void* kv_cache,
+    int const* out_cache_loc, int64_t page_stride, int64_t plane_stride, int64_t head_stride, int64_t token_stride,
+    int page_size, int num_tokens, int num_heads_q, int num_heads_k, int num_heads_v, int head_dim, int rotary_dim,
+    float eps, void const* q_weight, void const* k_weight, float base, int const* position_ids, cudaStream_t stream);
+
+// MiniMax-M3 sparse producer for the packed [Q|K|V|index-Q|index-K]
+// projection. It uses a precomputed FP32 RoPE table, emits compact FP8 Q and
+// index-Q, and inserts main K/V plus index-K into their paged FP8 HND caches.
+void launchMinimaxM3Fp8QKVIndexerNormRopeKVInsert(void const* packed_input, void* q_output, void* index_q_output,
+    void* kv_cache, void* index_k_cache, int const* out_cache_loc, int64_t kv_page_stride, int64_t kv_plane_stride,
+    int64_t kv_head_stride, int64_t kv_token_stride, int64_t index_page_stride, int64_t index_token_stride,
+    int page_size, int num_tokens, int num_heads_q, int num_heads_kv, int num_heads_index, int head_dim, int rotary_dim,
+    float eps, void const* q_weight, void const* k_weight, void const* index_q_weight, void const* index_k_weight,
+    float const* rotary_cos_sin, int const* position_ids, cudaStream_t stream);
+
 } // namespace kernels
 
 TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h

@@ -29,7 +29,7 @@ namespace kernels
 
 void invokeMinimaxM3SelectBlocks(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride,
     int32_t const* nValidBlocks, int32_t* output, int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries,
-    int32_t initBlocks, int32_t localBlocks, cudaStream_t stream);
+    int32_t initBlocks, int32_t localBlocks, bool headMajorOutput, cudaStream_t stream);
 
 } // namespace kernels
 

cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2025-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,11 +20,38 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <torch/extension.h>
 
+#include <tuple>
+
 TRTLLM_NAMESPACE_BEGIN
 
 namespace torch_ext
 {
 
+namespace
+{
+
+void checkMinimaxM3HndKVPool(torch::Tensor const& kvCache, int64_t numHeads, int64_t headDim)
+{
+    TORCH_CHECK(kvCache.is_cuda(), "kv_cache must be a CUDA tensor");
+    TORCH_CHECK(kvCache.scalar_type() == at::ScalarType::Float8_e4m3fn, "kv_cache must use torch.float8_e4m3fn");
+    TORCH_CHECK(kvCache.dim() == 5, "kv_cache must be HND [num_pages, 2, num_heads, page_size, head_dim]");
+    TORCH_CHECK(kvCache.size(0) > 0 && kvCache.size(3) > 0, "kv_cache must have positive num_pages and page_size");
+    TORCH_CHECK(kvCache.size(1) == 2, "kv_cache plane dimension must contain K and V");
+    TORCH_CHECK(kvCache.size(2) == numHeads, "kv_cache num_heads mismatch");
+    TORCH_CHECK(kvCache.size(4) == headDim, "kv_cache head_dim mismatch");
+    TORCH_CHECK(kvCache.stride(4) == 1 && kvCache.stride(3) == headDim,
+        "kv_cache must have contiguous head_dim rows in HND layout");
+    TORCH_CHECK(kvCache.stride(2) == kvCache.size(3) * kvCache.stride(3),
+        "kv_cache must have contiguous [page_size, head_dim] blocks in HND layout");
+    TORCH_CHECK(kvCache.stride(1) >= kvCache.size(2) * kvCache.stride(2), "kv_cache K and V planes must not overlap");
+    TORCH_CHECK(kvCache.stride(0) >= kvCache.size(1) * kvCache.stride(1),
+        "kv_cache page stride must not overlap adjacent HND pages");
+    TORCH_CHECK(kvCache.stride(0) % 4 == 0 && kvCache.stride(1) % 4 == 0,
+        "kv_cache page and plane strides must preserve 32-bit FP8 store alignment");
+}
+
+} // namespace
+
 // Function for fused QK Norm and RoPE
 // This operator applies RMS normalization and RoPE to Q and K tensors in a single CUDA kernel.
 // The OP performs operations in-place on the input qkv tensor.
@@ -149,6 +176,160 @@ torch::Tensor fused_qk_norm_rope_to_fp8_meta(torch::Tensor const& qkv, int64_t n
     return torch::empty({num_tokens, total_heads * head_dim}, qkv.options().dtype(torch::kFloat8_e4m3fn));
 }
 
+torch::Tensor minimaxM3Fp8QKNormRopeKVInsert(torch::Tensor const& qkv, torch::Tensor& kvCache,
+    torch::Tensor const& outCacheLoc, int64_t numHeadsQ, int64_t numHeadsK, int64_t numHeadsV, int64_t headDim,
+    int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, double base, bool isNeox,
+    torch::Tensor const& positionIds)
+{
+    constexpr int64_t kHeadDim = 128;
+    constexpr int64_t kRotaryDim = 64;
+    constexpr int64_t kPageSize = 128;
+    TORCH_CHECK(numHeadsQ > 0, "MiniMax-M3 FP8 main Q/K/V producer requires num_heads_q > 0");
+    TORCH_CHECK(numHeadsK > 0 && numHeadsV > 0, "MiniMax-M3 FP8 main Q/K/V producer requires K and V heads");
+    TORCH_CHECK(numHeadsK == numHeadsV, "MiniMax-M3 FP8 main Q/K/V producer requires equal K and V head counts");
+    TORCH_CHECK(headDim == kHeadDim, "MiniMax-M3 FP8 main Q/K/V producer requires head_dim=128");
+    TORCH_CHECK(rotaryDim == kRotaryDim, "MiniMax-M3 FP8 main Q/K/V producer requires rotary_dim=64");
+    TORCH_CHECK(isNeox, "MiniMax-M3 FP8 main Q/K/V producer requires NeoX RoPE");
+    TORCH_CHECK(eps >= 0.0, "MiniMax-M3 FP8 main Q/K/V producer requires eps >= 0");
+    TORCH_CHECK(base > 0.0, "MiniMax-M3 FP8 main Q/K/V producer requires base > 0");
+
+    TORCH_CHECK(qkv.dim() == 2, "QKV tensor must be 2D: [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]");
+    TORCH_CHECK(outCacheLoc.dim() == 1, "out_cache_loc must be one-dimensional");
+    TORCH_CHECK(positionIds.dim() == 1, "position_ids must be one-dimensional");
+    TORCH_CHECK(qWeight.dim() == 1 && kWeight.dim() == 1, "Q/K norm weights must be one-dimensional");
+
+    CHECK_INPUT(qkv, torch::kBFloat16);
+    CHECK_INPUT(outCacheLoc, torch::kInt32);
+    CHECK_INPUT(positionIds, torch::kInt32);
+    CHECK_INPUT(qWeight, torch::kBFloat16);
+    CHECK_INPUT(kWeight, torch::kBFloat16);
+    checkMinimaxM3HndKVPool(kvCache, numHeadsK, headDim);
+    TORCH_CHECK(kvCache.size(3) == kPageSize, "MiniMax-M3 FP8 main Q/K/V producer requires page_size=128");
+
+    int64_t const numTokens = qkv.size(0);
+    int64_t const totalHeads = numHeadsQ + numHeadsK + numHeadsV;
+    TORCH_CHECK(qkv.size(1) == totalHeads * headDim,
+        "QKV tensor width must equal (num_heads_q + num_heads_k + num_heads_v) * head_dim");
+    TORCH_CHECK(outCacheLoc.numel() >= numTokens, "out_cache_loc is shorter than num_tokens");
+    TORCH_CHECK(positionIds.numel() == numTokens, "position_ids length must equal num_tokens");
+    TORCH_CHECK(qWeight.numel() == headDim && kWeight.numel() == headDim, "Q/K norm weight width must equal head_dim");
+    TORCH_CHECK(qkv.get_device() == kvCache.get_device() && qkv.get_device() == outCacheLoc.get_device()
+            && qkv.get_device() == positionIds.get_device() && qkv.get_device() == qWeight.get_device()
+            && qkv.get_device() == kWeight.get_device(),
+        "All MiniMax-M3 FP8 main Q/K/V producer tensors must be on the same CUDA device");
+
+    auto qOut = torch::empty({numTokens, numHeadsQ, headDim}, qkv.options().dtype(at::ScalarType::Float8_e4m3fn));
+    if (numTokens == 0)
+    {
+        return qOut;
+    }
+
+    auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
+    tensorrt_llm::kernels::launchMinimaxM3Fp8QKNormRopeKVInsert(qkv.data_ptr(), qOut.data_ptr(), kvCache.data_ptr(),
+        outCacheLoc.data_ptr<int>(), kvCache.stride(0), kvCache.stride(1), kvCache.stride(2), kvCache.stride(3),
+        static_cast<int>(kvCache.size(3)), static_cast<int>(numTokens), static_cast<int>(numHeadsQ),
+        static_cast<int>(numHeadsK), static_cast<int>(numHeadsV), static_cast<int>(headDim),
+        static_cast<int>(rotaryDim), static_cast<float>(eps), qWeight.data_ptr(), kWeight.data_ptr(),
+        static_cast<float>(base), positionIds.data_ptr<int>(), stream);
+    return qOut;
+}
+
+torch::Tensor minimaxM3Fp8QKNormRopeKVInsertMeta(torch::Tensor const& qkv, torch::Tensor& /*kvCache*/,
+    torch::Tensor const& /*outCacheLoc*/, int64_t numHeadsQ, int64_t /*numHeadsK*/, int64_t /*numHeadsV*/,
+    int64_t headDim, int64_t /*rotaryDim*/, double /*eps*/, torch::Tensor const& /*qWeight*/,
+    torch::Tensor const& /*kWeight*/, double /*base*/, bool /*isNeox*/, torch::Tensor const& /*positionIds*/)
+{
+    return torch::empty({qkv.size(0), numHeadsQ, headDim}, qkv.options().dtype(at::ScalarType::Float8_e4m3fn));
+}
+
+std::tuple<torch::Tensor, torch::Tensor> minimaxM3Fp8QKVIndexerNormRopeKVInsert(torch::Tensor const& packed,
+    torch::Tensor& kvCache, torch::Tensor& indexKCache, torch::Tensor const& outCacheLoc, int64_t numHeadsQ,
+    int64_t numHeadsKV, int64_t numHeadsIndex, int64_t headDim, int64_t rotaryDim, double eps,
+    torch::Tensor const& qWeight, torch::Tensor const& kWeight, torch::Tensor const& indexQWeight,
+    torch::Tensor const& indexKWeight, torch::Tensor const& rotaryCosSin, torch::Tensor const& positionIds)
+{
+    constexpr int64_t kHeadDim = 128;
+    constexpr int64_t kRotaryDim = 64;
+    constexpr int64_t kPageSize = 128;
+    TORCH_CHECK(numHeadsQ > 0 && numHeadsKV > 0 && numHeadsIndex > 0,
+        "MiniMax-M3 horizontal producer requires Q, KV, and index heads");
+    TORCH_CHECK(numHeadsKV == numHeadsIndex, "MiniMax-M3 horizontal producer requires index heads to equal KV heads");
+    TORCH_CHECK(headDim == kHeadDim, "MiniMax-M3 horizontal producer requires head_dim=128");
+    TORCH_CHECK(rotaryDim == kRotaryDim, "MiniMax-M3 horizontal producer requires rotary_dim=64");
+    TORCH_CHECK(eps >= 0.0, "MiniMax-M3 horizontal producer requires eps >= 0");
+
+    TORCH_CHECK(packed.dim() == 2, "Packed QKV+index tensor must be two-dimensional");
+    TORCH_CHECK(outCacheLoc.dim() == 1, "out_cache_loc must be one-dimensional");
+    TORCH_CHECK(positionIds.dim() == 1, "position_ids must be one-dimensional");
+    CHECK_INPUT(packed, torch::kBFloat16);
+    CHECK_INPUT(outCacheLoc, torch::kInt32);
+    CHECK_INPUT(positionIds, torch::kInt32);
+    CHECK_INPUT(qWeight, torch::kBFloat16);
+    CHECK_INPUT(kWeight, torch::kBFloat16);
+    CHECK_INPUT(indexQWeight, torch::kBFloat16);
+    CHECK_INPUT(indexKWeight, torch::kBFloat16);
+    CHECK_INPUT(rotaryCosSin, torch::kFloat32);
+    checkMinimaxM3HndKVPool(kvCache, numHeadsKV, headDim);
+    TORCH_CHECK(kvCache.size(3) == kPageSize, "MiniMax-M3 horizontal producer requires page_size=128");
+    TORCH_CHECK(indexKCache.is_cuda() && indexKCache.scalar_type() == at::ScalarType::Float8_e4m3fn,
+        "Index-K cache must be CUDA torch.float8_e4m3fn");
+    TORCH_CHECK(indexKCache.dim() == 4 && indexKCache.size(1) == 1 && indexKCache.size(2) == kPageSize
+            && indexKCache.size(3) == kHeadDim,
+        "Index-K cache must be HND [num_pages, 1, 128, 128]");
+    TORCH_CHECK(indexKCache.stride(3) == 1 && indexKCache.stride(2) == kHeadDim,
+        "Index-K cache must have contiguous token rows");
+    TORCH_CHECK(rotaryCosSin.dim() == 3 && rotaryCosSin.size(1) == 2 && rotaryCosSin.size(2) == kRotaryDim / 2,
+        "rotary_cos_sin must be [max_positions, 2, rotary_dim/2]");
+
+    int64_t const numTokens = packed.size(0);
+    int64_t const totalHeads = numHeadsQ + 2 * numHeadsKV + numHeadsIndex + 1;
+    TORCH_CHECK(
+        packed.size(1) == totalHeads * headDim, "Packed tensor width must equal (Q + 2*KV + index-Q + 1) * head_dim");
+    TORCH_CHECK(outCacheLoc.numel() >= numTokens, "out_cache_loc is shorter than num_tokens");
+    TORCH_CHECK(positionIds.numel() == numTokens, "position_ids length must equal num_tokens");
+    TORCH_CHECK(qWeight.numel() == headDim && kWeight.numel() == headDim && indexQWeight.numel() == headDim
+            && indexKWeight.numel() == headDim,
+        "All norm weights must contain head_dim elements");
+    TORCH_CHECK(packed.get_device() == kvCache.get_device() && packed.get_device() == indexKCache.get_device()
+            && packed.get_device() == outCacheLoc.get_device() && packed.get_device() == positionIds.get_device()
+            && packed.get_device() == qWeight.get_device() && packed.get_device() == kWeight.get_device()
+            && packed.get_device() == indexQWeight.get_device() && packed.get_device() == indexKWeight.get_device()
+            && packed.get_device() == rotaryCosSin.get_device(),
+        "All MiniMax-M3 horizontal producer tensors must be on the same CUDA device");
+
+    auto qOut = torch::empty({numTokens, numHeadsQ, headDim}, packed.options().dtype(at::ScalarType::Float8_e4m3fn));
+    auto indexQOut
+        = torch::empty({numTokens, numHeadsIndex, headDim}, packed.options().dtype(at::ScalarType::Float8_e4m3fn));
+    if (numTokens == 0)
+    {
+        return {qOut, indexQOut};
+    }
+
+    auto stream = at::cuda::getCurrentCUDAStream(packed.get_device());
+    tensorrt_llm::kernels::launchMinimaxM3Fp8QKVIndexerNormRopeKVInsert(packed.data_ptr(), qOut.data_ptr(),
+        indexQOut.data_ptr(), kvCache.data_ptr(), indexKCache.data_ptr(), outCacheLoc.data_ptr<int>(),
+        kvCache.stride(0), kvCache.stride(1), kvCache.stride(2), kvCache.stride(3), indexKCache.stride(0),
+        indexKCache.stride(2), static_cast<int>(kvCache.size(3)), static_cast<int>(numTokens),
+        static_cast<int>(numHeadsQ), static_cast<int>(numHeadsKV), static_cast<int>(numHeadsIndex),
+        static_cast<int>(headDim), static_cast<int>(rotaryDim), static_cast<float>(eps), qWeight.data_ptr(),
+        kWeight.data_ptr(), indexQWeight.data_ptr(), indexKWeight.data_ptr(), rotaryCosSin.data_ptr<float>(),
+        positionIds.data_ptr<int>(), stream);
+    return {qOut, indexQOut};
+}
+
+std::tuple<torch::Tensor, torch::Tensor> minimaxM3Fp8QKVIndexerNormRopeKVInsertMeta(torch::Tensor const& packed,
+    torch::Tensor& /*kvCache*/, torch::Tensor& /*indexKCache*/, torch::Tensor const& /*outCacheLoc*/, int64_t numHeadsQ,
+    int64_t /*numHeadsKV*/, int64_t numHeadsIndex, int64_t headDim, int64_t /*rotaryDim*/, double /*eps*/,
+    torch::Tensor const& /*qWeight*/, torch::Tensor const& /*kWeight*/, torch::Tensor const& /*indexQWeight*/,
+    torch::Tensor const& /*indexKWeight*/, torch::Tensor const& /*rotaryCosSin*/, torch::Tensor const& /*positionIds*/)
+{
+    auto options = packed.options().dtype(at::ScalarType::Float8_e4m3fn);
+    return {
+        torch::empty({packed.size(0), numHeadsQ, headDim}, options),
+        torch::empty({packed.size(0), numHeadsIndex, headDim}, options),
+    };
+}
+
 // Register the PyTorch operators
 TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
@@ -163,19 +344,32 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         "rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, float "
         "factor, float low, float high, float attention_factor, bool is_qk_norm, bool use_gemma, bool use_mrope, int "
         "mrope_section1, int mrope_section2) -> Tensor");
+    m.def(
+        "minimax_m3_fp8_qk_norm_rope_kv_insert(Tensor qkv, Tensor(a!) kv_cache, Tensor out_cache_loc, int "
+        "num_heads_q, int num_heads_k, int num_heads_v, int head_dim, int rotary_dim, float eps, Tensor q_weight, "
+        "Tensor k_weight, float base, bool is_neox, Tensor position_ids) -> Tensor");
+    m.def(
+        "minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert(Tensor packed, Tensor(a!) kv_cache, Tensor(b!) "
+        "index_k_cache, Tensor out_cache_loc, int num_heads_q, int num_heads_kv, int num_heads_index, int head_dim, "
+        "int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, Tensor index_q_weight, Tensor index_k_weight, "
+        "Tensor rotary_cos_sin, Tensor position_ids) -> (Tensor, Tensor)");
 }
 
 // Register the CUDA implementation
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
 {
     m.impl("fused_qk_norm_rope", &fused_qk_norm_rope);
     m.impl("fused_qk_norm_rope_to_fp8", &fused_qk_norm_rope_to_fp8);
+    m.impl("minimax_m3_fp8_qk_norm_rope_kv_insert", &minimaxM3Fp8QKNormRopeKVInsert);
+    m.impl("minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert", &minimaxM3Fp8QKVIndexerNormRopeKVInsert);
 }
 
 // Register the Meta implementation (shape/dtype inference for torch.compile).
 TORCH_LIBRARY_IMPL(trtllm, Meta, m)
 {
     m.impl("fused_qk_norm_rope_to_fp8", &fused_qk_norm_rope_to_fp8_meta);
+    m.impl("minimax_m3_fp8_qk_norm_rope_kv_insert", &minimaxM3Fp8QKNormRopeKVInsertMeta);
+    m.impl("minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert", &minimaxM3Fp8QKVIndexerNormRopeKVInsertMeta);
 }
 
 } // namespace torch_ext

cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp

@@ -27,7 +27,7 @@ namespace torch_ext
 {
 
 torch::Tensor minimaxM3SelectBlocks(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, int64_t topK,
-    int64_t initBlocks, int64_t localBlocks)
+    int64_t initBlocks, int64_t localBlocks, bool headMajorOutput)
 {
     constexpr int64_t kRequiredTopK = 16;
     constexpr int64_t kMaxBlockIndex = 65'535;
@@ -61,13 +61,16 @@ torch::Tensor minimaxM3SelectBlocks(torch::Tensor const& scores, torch::Tensor c
     TORCH_CHECK(initBlocks <= std::numeric_limits<int32_t>::max() && localBlocks <= std::numeric_limits<int32_t>::max(),
         "minimax_m3_select_blocks forcing ranges exceed int32 range");
 
-    auto output = torch::empty({scores.size(2), scores.size(0), topK}, scores.options().dtype(torch::kInt32));
+    auto const outputOptions = scores.options().dtype(torch::kInt32);
+    auto output = headMajorOutput
+        ? torch::empty({scores.size(0), scores.size(2), topK}, outputOptions).permute({1, 0, 2})
+        : torch::empty({scores.size(2), scores.size(0), topK}, outputOptions);
     auto const stream = at::cuda::getCurrentCUDAStream(scores.get_device());
     tensorrt_llm::kernels::invokeMinimaxM3SelectBlocks(scores.data_ptr<float>(), scores.stride(0), scores.stride(1),
         scores.stride(2), nValidBlocks.data_ptr<int32_t>(), output.data_ptr<int32_t>(),
         static_cast<int32_t>(scores.size(0)), static_cast<int32_t>(scores.size(1)),
         static_cast<int32_t>(scores.size(2)), static_cast<int32_t>(initBlocks), static_cast<int32_t>(localBlocks),
-        stream);
+        headMajorOutput, stream);
     return output;
 }
 
@@ -79,7 +82,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
     m.def(
         "minimax_m3_select_blocks(Tensor scores, Tensor n_valid_blocks, int topk, int init_blocks, int "
-        "local_blocks) -> Tensor");
+        "local_blocks, bool head_major_output=False) -> Tensor");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)