wehub-resource-sync
59a0a3844c
PR Test AMD / finish (push) Blocked by required conditions
PR Test NVIDIA ARM / finish (push) Blocked by required conditions
PR Test NVIDIA / finish (push) Blocked by required conditions
PR Test AMD / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / scan (push) Has been skipped
PR Test NVIDIA / cancel-on-close (push) Has been skipped
PR Test AMD / scan (push) Has been skipped
PR Test AMD / ${{ matrix.name }} (${{ matrix.runner }}) (push) Waiting to run
PR Test NVIDIA ARM / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / ${{ matrix.name }} (${{ matrix.runner }}) (push) Waiting to run
PR Test NVIDIA / scan (push) Has been skipped
Release Docker Images / build (cu129-torch-2.11.0) (push) Has been skipped
PR Test NVIDIA / ${{ matrix.name }} (${{ matrix.runner }}) (push) Waiting to run
Release Docker Images / build (cu130-torch-2.11.0) (push) Has been skipped
Release PyPI / publish (push) Has been skipped
Scheduler Python Test / test (push) Successful in 27m19s
Docs / build (push) Successful in 28m8s
Scheduler C++ Test / test (push) Successful in 28m19s
Scheduler C++ Test / test-flat (push) Successful in 28m18s
Docs / deploy (push) Has been cancelled
91 lines
2.7 KiB
Python
91 lines
2.7 KiB
Python
# Copyright (c) 2026 LightSeek Foundation
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in
|
|
# all copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
# SOFTWARE.
|
|
|
|
from tokenspeed_kernel.profiling import bootstrap_profiling_from_env
|
|
|
|
bootstrap_profiling_from_env()
|
|
|
|
from tokenspeed_kernel.ops.attention import (
|
|
GdnCheckpointLayout,
|
|
GdnChunkPrefillResult,
|
|
attn_merge_state,
|
|
dsa_decode,
|
|
dsa_decode_topk,
|
|
dsa_plan,
|
|
dsa_prefill,
|
|
dsa_prefill_topk,
|
|
gdn_chunk_prefill,
|
|
mha_decode_with_kvcache,
|
|
mha_extend_with_kvcache,
|
|
mha_plan,
|
|
mha_prefill,
|
|
mla_decode_with_kvcache,
|
|
mla_prefill,
|
|
)
|
|
from tokenspeed_kernel.ops.gemm import mm
|
|
from tokenspeed_kernel.ops.moe import moe_apply, moe_plan, moe_process_weights
|
|
from tokenspeed_kernel.ops.quantization import (
|
|
quantize_fp8,
|
|
quantize_fp8_with_scale,
|
|
quantize_mxfp4,
|
|
quantize_mxfp8,
|
|
quantize_nvfp4,
|
|
)
|
|
from tokenspeed_kernel.ops.sampling import argmax
|
|
from tokenspeed_kernel.ops.transform import hadamard_transform
|
|
from tokenspeed_kernel.selection import NoKernelFoundError
|
|
|
|
__all__ = [
|
|
# exceptions
|
|
"NoKernelFoundError",
|
|
# gemm
|
|
"mm",
|
|
# attention
|
|
"mha_plan",
|
|
"mha_prefill",
|
|
"mha_extend_with_kvcache",
|
|
"mha_decode_with_kvcache",
|
|
"mla_prefill",
|
|
"mla_decode_with_kvcache",
|
|
"dsa_prefill",
|
|
"dsa_decode",
|
|
"dsa_prefill_topk",
|
|
"dsa_decode_topk",
|
|
"dsa_plan",
|
|
"attn_merge_state",
|
|
"gdn_chunk_prefill",
|
|
"GdnCheckpointLayout",
|
|
"GdnChunkPrefillResult",
|
|
# moe
|
|
"moe_apply",
|
|
"moe_plan",
|
|
"moe_process_weights",
|
|
# quantization
|
|
"quantize_fp8",
|
|
"quantize_fp8_with_scale",
|
|
"quantize_mxfp8",
|
|
"quantize_nvfp4",
|
|
"quantize_mxfp4",
|
|
# sampling
|
|
"argmax",
|
|
# transform
|
|
"hadamard_transform",
|
|
]
|