diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 93e6a2c18e21..f853549541ad 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -291,6 +291,18 @@ class TllmGenFmhaKernel bool shouldUseNvrtc = options.mFmhaKernelType == FmhaKernelType::SwapsMmaAbForGeneration && !options.mIsMlaGen && options.mDtypeKv != tg::Dtype::E2m1 && options.mHeadDimQk != 64 && !isLlama70bFp4Tp4; + // Prefer a pre-compiled cubin over NVRTC when one is available (e.g. sm_100f cubins on sm_100 GPUs). + if (shouldUseNvrtc) + { + FmhaOptions cubinCheckOptions = options; + algoFilterForCubinPath(cubinCheckOptions); + auto [checkHashId, checkInfo] = hashFromFmhaOptions(cubinCheckOptions); + if (mFunctions.find(checkHashId) != mFunctions.end()) + { + shouldUseNvrtc = false; + } + } + if (shouldUseNvrtc) { // nvrtc path - uses mFmhaInterface member for kernel caching diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index cf194e94a936..b2f25f2cf584 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -162,6 +162,31 @@ def low_memory_overrides(config, return config +def reduced_model_kwargs(num_hidden_layers: int, + model_path: str | None = None) -> dict: + """Return model_kwargs to cap a model at ``num_hidden_layers`` layers. + + Reduces peak memory so large models fit on a single GPU for pre-merge + smoke testing. The rest of the architecture (attention, MoE, SSM) is + preserved; only the layer count is truncated. + + For models whose config derives layer count from a list attribute (e.g. + ``layers_block_type`` in NemotronH), pass ``model_path`` so the list + is also truncated — otherwise the ``num_hidden_layers`` override has + no effect. + """ + overrides = {"num_hidden_layers": num_hidden_layers} + if model_path is not None: + from transformers import AutoConfig + hf_config = AutoConfig.from_pretrained(model_path, + trust_remote_code=True) + for attr in ("layers_block_type", ): + val = getattr(hf_config, attr, None) + if val is not None: + overrides[attr] = val[:num_hidden_layers] + return {"model_kwargs": overrides} + + class TestLlama3_1_8B(LlmapiAccuracyTestHarness): MODEL_NAME = "meta-llama/Llama-3.1-8B" MODEL_PATH = hf_id_to_local_model_dir(MODEL_NAME) @@ -598,6 +623,33 @@ def test_accuracy(self, model_id, world_size, enable_attention_dp, print_memory_usage("after evaluation") + @skip_pre_hopper + @pytest.mark.skip_less_device_memory(40000) + @pytest.mark.parametrize("dtype", ["bf16", "fp8"]) + def test_functional_small(self, dtype): + """Single-GPU smoke test using a layer-reduced model. + + Overrides num_hidden_layers so the 120B model fits on one GPU, + enabling pre-merge coverage of the full kernel dispatch path + (attention, MoE, SSM) without requiring a multi-GPU machine. + No accuracy threshold is checked — the truncated model is not + expected to produce meaningful text. + """ + model_path = self.MODEL_PATHS[dtype] + kwargs = {} + kwargs.update( + reduced_model_kwargs(num_hidden_layers=16, model_path=model_path)) + with AutoDeployLLM(model=model_path, + tokenizer=model_path, + world_size=1, + yaml_extra=[self.CONFIG_YAML], + trust_remote_code=True, + **kwargs) as llm: + outputs = llm.generate( + ["Hello, how are you?"], + sampling_params=SamplingParams(max_tokens=10)) + assert len(outputs) == 1 + @pytest.mark.skip_less_device_memory(180000) @pytest.mark.parametrize("world_size", [4, 8]) def test_mtp(self, world_size): diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 08272f76ad59..e71f9cb44e8f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -298,6 +298,8 @@ l0_b200: - accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[fp8-1-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-1-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[nvfp4-1-attn_dp_off-trtllm] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[bf16] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[fp8] - unittest/auto_deploy/singlegpu/compile - unittest/auto_deploy/singlegpu/custom_ops - unittest/auto_deploy/singlegpu/models