From 99cb2ed98894f4bebb89c4a5da8177ef6e14d575 Mon Sep 17 00:00:00 2001 From: Hudson Xing <1277646412@qq.com> Date: Sun, 14 Dec 2025 15:22:37 -0800 Subject: [PATCH] [refactor] Move trtllm_fp8_kv_kernel to triton_ops directory (#15044) --- .../layers/attention/{ => triton_ops}/trtllm_fp8_kv_kernel.py | 0 python/sglang/srt/layers/attention/trtllm_mha_backend.py | 4 +++- test/manual/test_trtllm_fp8_kv_kernel.py | 4 +++- 3 files changed, 6 insertions(+), 2 deletions(-) rename python/sglang/srt/layers/attention/{ => triton_ops}/trtllm_fp8_kv_kernel.py (100%) diff --git a/python/sglang/srt/layers/attention/trtllm_fp8_kv_kernel.py b/python/sglang/srt/layers/attention/triton_ops/trtllm_fp8_kv_kernel.py similarity index 100% rename from python/sglang/srt/layers/attention/trtllm_fp8_kv_kernel.py rename to python/sglang/srt/layers/attention/triton_ops/trtllm_fp8_kv_kernel.py diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 5dcbfb1b2..8f2447d21 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -15,7 +15,9 @@ from sglang.srt.layers.attention.flashinfer_backend import ( FlashInferAttnBackend, FlashInferMultiStepDraftBackend, ) -from sglang.srt.layers.attention.trtllm_fp8_kv_kernel import fused_fp8_set_kv_buffer +from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import ( + fused_fp8_set_kv_buffer, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.utils import is_flashinfer_available diff --git a/test/manual/test_trtllm_fp8_kv_kernel.py b/test/manual/test_trtllm_fp8_kv_kernel.py index b713747a2..c97deaf11 100644 --- a/test/manual/test_trtllm_fp8_kv_kernel.py +++ b/test/manual/test_trtllm_fp8_kv_kernel.py @@ -6,7 +6,9 @@ import unittest import torch -from sglang.srt.layers.attention.trtllm_fp8_kv_kernel import fused_fp8_set_kv_buffer +from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import ( + fused_fp8_set_kv_buffer, +) from sglang.test.test_utils import CustomTestCase