mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
0fd59d64ab
* add fmha repo Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> * fix format Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> * fix code style Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> * fix header Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> * fix header kernel_traits.h Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> * add .gitignore file Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> * add SLIDING_WINDOW_ATTENTION Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> * fix style Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> * fix format Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> * update setup.py Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> * update build_wheel.py Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> --------- Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com> Signed-off-by: qsang-nv <200703406+qsang-nv@users.noreply.github.com>
32 lines
1.3 KiB
Python
32 lines
1.3 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2023-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: NVIDIA TensorRT Source Code License Agreement
|
|
#
|
|
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
|
|
# property and proprietary rights in and to this material, related
|
|
# documentation and any modifications thereto. Any use, reproduction,
|
|
# disclosure or distribution of this material and related documentation
|
|
# without an express license agreement from NVIDIA CORPORATION or
|
|
# its affiliates is strictly prohibited.
|
|
|
|
import subprocess
|
|
|
|
import pytest
|
|
|
|
|
|
def gpu_compute_cap_str():
|
|
output = subprocess.check_output(
|
|
['nvidia-smi', "--query-gpu=compute_cap", "--format=csv"])
|
|
csv_header, csv_value, *other_csv_values = output.splitlines()
|
|
return 'sm' + str(int(float(csv_value) * 10))
|
|
|
|
|
|
def pytest_collection_modifyitems(config, items):
|
|
for item in items:
|
|
# print("[debug] " + item.nodeid)
|
|
# skip tests that are not runnable on local machine
|
|
if 'arch' in item.nodeid and gpu_compute_cap_str() not in item.nodeid:
|
|
item.add_marker(pytest.mark.skip(reason='arch-dependent test'))
|
|
# mark tests as unit when under 'unit'
|
|
if 'unit' in item.nodeid:
|
|
item.add_marker(pytest.mark.unit)
|