[None][fix] Address DSv4 primitive review comments

mingyangHao · mingyangHao · commit c75b5fc42a87 · 2026-06-15T20:30:07.000-07:00
Signed-off-by: Mingyang Hao &lt;200044211+mingyangHao@users.noreply.github.com&gt;
diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt
@@ -200,11 +200,13 @@ set(TRTLLM_LINK_LIBS
     layers_src
     runtime_src
     testing_src
-    mhcKernels_src
     compressorKernels_src
     userbuffers_src
     ${DECODER_SHARED_TARGET_0}
     ${DECODER_SHARED_TARGET_1})
+if(BUILD_DEEP_GEMM)
+  set(TRTLLM_LINK_LIBS ${TRTLLM_LINK_LIBS} mhcKernels_src)
+endif()
 if(USING_OSS_CUTLASS_LOW_LATENCY_GEMM)
   set(TRTLLM_LINK_LIBS ${TRTLLM_LINK_LIBS} low_latency_gemm_src)
   message(STATUS "USING_OSS_CUTLASS_LOW_LATENCY_GEMM")
diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt
@@ -30,7 +30,11 @@ add_subdirectory(dsv3MinLatencyKernels)
 add_subdirectory(causalConv1d)
 add_subdirectory(fusedGatedRMSNormQuant)
 add_subdirectory(mamba2MTPSSMCache)
-add_subdirectory(mhcKernels)
+# The mHC fused MMA path depends on DeepGEMM helper headers, which are only
+# fetched when BUILD_DEEP_GEMM is enabled.
+if(BUILD_DEEP_GEMM)
+  add_subdirectory(mhcKernels)
+endif()
 add_subdirectory(compressorKernels)
 
 file(GLOB_RECURSE SRC_CPP *.cpp)
diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt
@@ -1,6 +1,6 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION &
-# AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved. SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
 # use this file except in compliance with the License. You may obtain a copy of
diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 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.
diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 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.
diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt
@@ -1,6 +1,6 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION &
-# AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved. SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
 # use this file except in compliance with the License. You may obtain a copy of
diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 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.
diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 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.
diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 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.
diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 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.
diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 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.
diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt
@@ -34,6 +34,11 @@ set_property(TARGET th_utils PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
 target_link_libraries(th_utils PUBLIC ${TORCH_LIBRARIES} ${CUBLAS_LIB}
                                       ${CURAND_LIB})
 
+set(TRTLLM_TH_COMMON_EXTRA_SRCS)
+if(BUILD_DEEP_GEMM)
+  list(APPEND TRTLLM_TH_COMMON_EXTRA_SRCS mhcOp.cpp)
+endif()
+
 # TODO This does not compile with internal cutlass MOE gemm
 add_library(
   th_common SHARED
@@ -124,7 +129,7 @@ add_library(
   convertReqIndexToGlobalOp.cpp
   trtllmGenQKVProcessOp.cpp
   inplaceSliceCopyOp.cpp
-  mhcOp.cpp
+  ${TRTLLM_TH_COMMON_EXTRA_SRCS}
   compressorOp.cpp)
 set_property(TARGET th_common PROPERTY POSITION_INDEPENDENT_CODE ON)
 target_link_libraries(
diff --git a/cpp/tensorrt_llm/thop/compressorOp.cpp b/cpp/tensorrt_llm/thop/compressorOp.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 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.
diff --git a/cpp/tensorrt_llm/thop/mhcOp.cpp b/cpp/tensorrt_llm/thop/mhcOp.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 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.
diff --git a/tests/unittest/_torch/modules/test_mhc.py b/tests/unittest/_torch/modules/test_mhc.py
@@ -8,6 +8,7 @@
 import torch
 import torch.nn.functional as F
 from torch.profiler import ProfilerActivity, profile
+from utils.util import skip_pre_blackwell
 
 from tensorrt_llm._torch.modules.mhc.hyper_connection import HCHead, mHC
 
@@ -614,6 +615,7 @@ def test_mhc_fused_hc_fused_norm(n: int, hidden_size: int, hc_mult: int):
 @pytest.mark.parametrize("n", list(_BACKEND_TACTICS_BY_M.keys()))
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
 @pytest.mark.parametrize("hc_mult", [4])
+@skip_pre_blackwell
 def test_mhc_fused_hc_backends(n: int, hidden_size: int, hc_mult: int):
     """Every wired fused_hc backend sees bit-identical input and is checked
     against one shared torch reference and one golden backend output.
@@ -788,9 +790,9 @@ def make_runner_inputs():
 @pytest.mark.parametrize(
     "tactic",
     [
-        ("fused_half_mma", 0, 1, 256, 1),
+        pytest.param(("fused_half_mma", 0, 1, 256, 1), marks=skip_pre_blackwell),
         ("fused_half_fma", 2, 1, 256, 1),
-        ("fused_all_mma", 0, 1, 0, 1),
+        pytest.param(("fused_all_mma", 0, 1, 0, 1), marks=skip_pre_blackwell),
         ("fused_all_fma", 2, 1, 0, 1),
     ],
 )
@@ -862,6 +864,7 @@ def test_mhc_fused_hc_realistic_scale_regression(tactic, hidden_size: int):
 @pytest.mark.parametrize("n", [128, 2048])
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
 @pytest.mark.parametrize("hc_mult", [4])
+@skip_pre_blackwell
 def test_mhc_fused_hc_cuda_graph(n: int, hidden_size: int, hc_mult: int):
     """CUDA-graph capture/replay of mHC.fused_hc.
 

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`#`
`2`		`-# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION &`
`3`		`-# AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0`
	`2`	`+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.`
	`3`	`+# All rights reserved. SPDX-License-Identifier: Apache-2.0`
`4`	`4`	`#`
`5`	`5`	`# Licensed under the Apache License, Version 2.0 (the "License"); you may not`
`6`	`6`	`# use this file except in compliance with the License. You may obtain a copy of`
Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`	`1`	`/*`
`2`		`- * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved.`
	`2`	`+ * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.`
`3`	`3`	`*`
`4`	`4`	`* Licensed under the Apache License, Version 2.0 (the "License");`
`5`	`5`	`* you may not use this file except in compliance with the License.`