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

Fix TRT-LLM MLA backend applying k_scale to BF16 KV cache in BMM1 (#18396)

Browse Source
This commit is contained in:
debo3
2026-02-08 23:11:16 +08:00
committed by GitHub
parent a41aff1243
commit 031a652b93
1 changed files with 36 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+36 -12
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
Support attention backend for TRTLLM MLA kernels from flashinfer. Support attention backend for TRTLLM MLA kernels from flashinfer.
""" """
import logging
import math import math
from dataclasses import dataclass from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional, Union from typing import TYPE_CHECKING, Optional, Union
@@ -35,6 +36,8 @@ if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.speculative.spec_info import SpecInput
logger = logging.getLogger(__name__)
_is_cuda = is_cuda() _is_cuda = is_cuda()
if _is_cuda: if _is_cuda:
@@ -909,15 +912,26 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
# The final BMM1 scale is computed as: q_scale * k_scale * softmax_scale # The final BMM1 scale is computed as: q_scale * k_scale * softmax_scale
# Scale components: # Scale components:
# - q_scale: Query scaling factor (set to 1.0 for both FP16/FP8 paths) # - q_scale: Query scaling factor (set to 1.0 for both FP16/FP8 paths)
# - k_scale: Key scaling factor from model checkpoint (defaults to 1.0 if not available) # - k_scale: Key scaling factor from model checkpoint. Only applied when KV cache
# stores FP8-quantized values, to compensate for the quantization scaling.
# For BF16/FP16 KV cache, k_scale must be 1.0 since values are unscaled.
# - softmax_scale: Attention softmax scaling = 1/sqrt(head_dim), pre-computed as layer.scaling # - softmax_scale: Attention softmax scaling = 1/sqrt(head_dim), pre-computed as layer.scaling
# This unified approach works for both FP16 and FP8 quantized attention paths.
q_scale = 1.0 q_scale = 1.0
k_scale = ( if self.data_type == torch.float8_e4m3fn:
layer.k_scale_float k_scale = (
if getattr(layer, "k_scale_float", None) is not None layer.k_scale_float
else 1.0 if getattr(layer, "k_scale_float", None) is not None
) else 1.0
)
else:
if getattr(layer, "k_scale_float", None) is not None:
logger.warning_once(
"Checkpoint has k_scale but KV cache dtype is not FP8. "
"Ignoring k_scale for BMM1 (k_scale=%.4f, kv_dtype=%s).",
layer.k_scale_float,
self.data_type,
)
k_scale = 1.0
bmm1_scale = q_scale * k_scale * layer.scaling bmm1_scale = q_scale * k_scale * layer.scaling
@@ -1033,11 +1047,21 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
kv_cache = k_cache.view(-1, self.page_size, self.kv_cache_dim).unsqueeze(1) kv_cache = k_cache.view(-1, self.page_size, self.kv_cache_dim).unsqueeze(1)
q_scale = 1.0 q_scale = 1.0
k_scale = ( if self.data_type == torch.float8_e4m3fn:
layer.k_scale_float k_scale = (
if getattr(layer, "k_scale_float", None) is not None layer.k_scale_float
else 1.0 if getattr(layer, "k_scale_float", None) is not None
) else 1.0
)
else:
if getattr(layer, "k_scale_float", None) is not None:
logger.warning_once(
"Checkpoint has k_scale but KV cache dtype is not FP8. "
"Ignoring k_scale for BMM1 (k_scale=%.4f, kv_dtype=%s).",
layer.k_scale_float,
self.data_type,
)
k_scale = 1.0
q = q.to(self.data_type) q = q.to(self.data_type)
bmm1_scale = q_scale * k_scale * layer.scaling bmm1_scale = q_scale * k_scale * layer.scaling