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

fix trtllm mla spec (#13738)

Browse Source 
Co-authored-by: Brayden Zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
b8zhong
2025-12-01 22:16:25 -08:00
committed by GitHub
co-authored by Brayden Zhong
parent 3dabd609fb
commit 236a7c2370
1 changed files with 0 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mla_backend.py
-3
View File
@@ -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.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.server_args import get_global_server_args 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 import is_cuda, is_flashinfer_available, is_float4_e2m1fn_x2
from sglang.srt.utils.common import cached_triton_kernel
if is_flashinfer_available(): if is_flashinfer_available():
import flashinfer import flashinfer
@@ -51,7 +50,6 @@ DEFAULT_WORKSPACE_SIZE_MB = 128 # Memory workspace size in MB
TRTLLM_BLOCK_CONSTRAINT = 128 TRTLLM_BLOCK_CONSTRAINT = 128
@cached_triton_kernel(lambda _, kwargs: (kwargs["BLOCK_SIZE"]))
@triton.jit @triton.jit
def pad_draft_extend_query_kernel( def pad_draft_extend_query_kernel(
q_ptr, # Input query tensor [total_seq_len, num_heads, head_dim] 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 @triton.jit
def unpad_draft_extend_output_kernel( 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) raw_out_ptr, # Input raw output tensor (batch_size, token_per_batch, tp_q_head_num, v_head_dim)