fix tokenspeed_mla attn kernel jit (#26170)

2026-05-23 03:24:33 -07:00
parent 89ff2bc111
commit cb7b57955d
1 changed files with 3 additions and 2 deletions
@@ -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,