From 236a7c237002250b148c79bd93780d870b8b50d2 Mon Sep 17 00:00:00 2001 From: b8zhong Date: Mon, 1 Dec 2025 22:16:25 -0800 Subject: [PATCH] fix trtllm mla spec (#13738) Co-authored-by: Brayden Zhong --- python/sglang/srt/layers/attention/trtllm_mla_backend.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index c3fbaa422..0c8f832ed 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -24,7 +24,6 @@ from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import is_cuda, is_flashinfer_available, is_float4_e2m1fn_x2 -from sglang.srt.utils.common import cached_triton_kernel if is_flashinfer_available(): import flashinfer @@ -51,7 +50,6 @@ DEFAULT_WORKSPACE_SIZE_MB = 128 # Memory workspace size in MB TRTLLM_BLOCK_CONSTRAINT = 128 -@cached_triton_kernel(lambda _, kwargs: (kwargs["BLOCK_SIZE"])) @triton.jit def pad_draft_extend_query_kernel( q_ptr, # Input query tensor [total_seq_len, num_heads, head_dim] @@ -125,7 +123,6 @@ def pad_draft_extend_query_kernel( ) -@cached_triton_kernel(lambda _, kwargs: (kwargs["BLOCK_SIZE"])) @triton.jit def unpad_draft_extend_output_kernel( raw_out_ptr, # Input raw output tensor (batch_size, token_per_batch, tp_q_head_num, v_head_dim)