chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Refactor KTransformers heterogeneous compute with unified GPU-quantization backend (#12834)

Browse Source 
Co-authored-by: Chen Hongtao <56470055+chenht2022@users.noreply.github.com>
Co-authored-by: chenht2022 <cht22@mails.tsinghua.edu.cn>
Co-authored-by: skqliao <skqliao@gmail.com>
Co-authored-by: ovowei <1913953267@qq.com>
This commit is contained in:
Atream
2025-11-10 13:06:32 +08:00
committed by GitHub
parent d1be60c3c5
commit ddd1440d0f
10 changed files with 494 additions and 507 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/model_executor/cuda_graph_runner.py
View File

@@ -69,7 +69,7 @@ from sglang.srt.utils.patch_torch import monkey_patch_torch_compile
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
try:
from kt_kernel import AMXMoEWrapper
from kt_kernel import KTMoEWrapper
KTRANSFORMERS_AVAILABLE = True
except ImportError:
@@ -259,7 +259,7 @@ class CudaGraphRunner:
self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner)
log_info_on_rank0(logger, f"Capture cuda graph bs {self.capture_bs}")
if KTRANSFORMERS_AVAILABLE:
AMXMoEWrapper.set_capture_batch_sizes(self.capture_bs)
KTMoEWrapper.set_capture_batch_sizes(self.capture_bs)
self.capture_forward_mode = ForwardMode.DECODE
self.capture_hidden_mode = CaptureHiddenMode.NULL
self.num_tokens_per_bs = 1