67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
|
#
|
|
# Use of this software is governed by the terms and conditions of the
|
|
# NVIDIA End User License Agreement (EULA), available at:
|
|
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
|
|
#
|
|
# Any use, reproduction, disclosure, or distribution of this software
|
|
# and related documentation outside the scope permitted by the EULA
|
|
# is strictly prohibited.
|
|
|
|
from .static_persistent_tile_scheduler import (
|
|
WorkTileInfo,
|
|
PersistentTileSchedulerParams,
|
|
StaticPersistentTileScheduler,
|
|
)
|
|
|
|
from .hardware_info import (
|
|
HardwareInfo,
|
|
)
|
|
|
|
from .blackwell_helpers import (
|
|
compute_epilogue_tile_shape,
|
|
get_smem_store_op,
|
|
get_tmem_load_op,
|
|
get_num_tmem_alloc_cols,
|
|
make_smem_layout_a,
|
|
make_smem_layout_b,
|
|
make_smem_layout_epi,
|
|
make_trivial_tiled_mma,
|
|
)
|
|
|
|
from .hopper_helpers import (
|
|
sm90_get_smem_store_op,
|
|
)
|
|
|
|
from .grouped_gemm_tile_scheduler_helper import (
|
|
GroupSearchResult,
|
|
GroupedGemmGroupSearchState,
|
|
GroupedGemmTileSchedulerHelper,
|
|
create_initial_search_state,
|
|
)
|
|
|
|
from .tensormap_manager import (
|
|
TensorMapUpdateMode,
|
|
TensorMapManager,
|
|
)
|
|
|
|
from .smem_allocator import SmemAllocator
|
|
|
|
from .layout import LayoutEnum
|
|
|
|
__all__ = [
|
|
"SmemAllocator",
|
|
"LayoutEnum",
|
|
"WorkTileInfo",
|
|
"PersistentTileSchedulerParams",
|
|
"StaticPersistentTileScheduler",
|
|
"TensorMapUpdateMode",
|
|
"TensorMapManager",
|
|
"GroupSearchResult",
|
|
"GroupedGemmGroupSearchState",
|
|
"create_initial_search_state",
|
|
"GroupedGemmTileSchedulerHelper",
|
|
"HardwareInfo",
|
|
]
|