From cb7b57955d7bce40655ebc979a91869fe1071776 Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Sat, 23 May 2026 03:24:33 -0700 Subject: [PATCH] fix tokenspeed_mla attn kernel jit (#26170) --- python/sglang/srt/layers/attention/tokenspeed_mla_backend.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py index 5e513f081..af3da62c3 100644 --- a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py +++ b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py @@ -134,8 +134,9 @@ class TokenspeedMLABackend(TRTLLMMLABackend): # branch, which always asks for the LSE. if is_causal is False and return_lse is False: continue + # Runtime feeds fp8_e4m3fn q/k/v config = ( - torch.bfloat16, + torch.float8_e4m3fn, head_dim_qk, self.v_head_dim, is_causal, @@ -146,7 +147,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend): if config in _compiled_kernels: continue _compiled_kernels[config] = _compile_prefill_kernel( - torch.bfloat16, + torch.float8_e4m3fn, head_dim_qk, self.v_head_dim, is_causal,