Add TIRx inference kernels and B200 denoising timings

dceb455 verified 5 days ago

1.57 kB

# Copyright (c) 2026 SandAI. 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.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .dit_model import get_dit
from .dit_module import (
    DiTModel,
    SparseMoEMLP,
    collect_moe_aux_loss,
    collect_moe_routing_stats,
    dense_mlp_intermediate_size,
    load_dense_weights_into_sparse_moe_model,
    remap_dense_mlp_state_to_sparse_moe,
    resolve_virtual_group,
    set_sparse_moe_diffusion_timestep,
    split_expert_intermediate,
    token_sparse_moe_mlp_layers,
)
from .moe_trace import MoeRoutingAccumulator, collect_moe_modality_counts

__all__ = [
    "DiTModel",
    "get_dit",
    "SparseMoEMLP",
    "collect_moe_aux_loss",
    "collect_moe_routing_stats",
    "dense_mlp_intermediate_size",
    "load_dense_weights_into_sparse_moe_model",
    "remap_dense_mlp_state_to_sparse_moe",
    "resolve_virtual_group",
    "set_sparse_moe_diffusion_timestep",
    "split_expert_intermediate",
    "token_sparse_moe_mlp_layers",
    "MoeRoutingAccumulator",
    "collect_moe_modality_counts",
]

	# Copyright (c) 2026 SandAI. 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.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	from .dit_model import get_dit
	from .dit_module import (
	DiTModel,
	SparseMoEMLP,
	collect_moe_aux_loss,
	collect_moe_routing_stats,
	dense_mlp_intermediate_size,
	load_dense_weights_into_sparse_moe_model,
	remap_dense_mlp_state_to_sparse_moe,
	resolve_virtual_group,
	set_sparse_moe_diffusion_timestep,
	split_expert_intermediate,
	token_sparse_moe_mlp_layers,
	)
	from .moe_trace import MoeRoutingAccumulator, collect_moe_modality_counts

	__all__ = [
	"DiTModel",
	"get_dit",
	"SparseMoEMLP",
	"collect_moe_aux_loss",
	"collect_moe_routing_stats",
	"dense_mlp_intermediate_size",
	"load_dense_weights_into_sparse_moe_model",
	"remap_dense_mlp_state_to_sparse_moe",
	"resolve_virtual_group",
	"set_sparse_moe_diffusion_timestep",
	"split_expert_intermediate",
	"token_sparse_moe_mlp_layers",
	"MoeRoutingAccumulator",
	"collect_moe_modality_counts",
	]