A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 16:00 UTC
[MLA] Fuse cat and qaunt for fp8 kv-cache (#32950)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jan 24, 2026 at 09:03 UTC
da5e7b12beb6f4496b01265025faf5a123628ae5
1 file changed
+41
-20
vllm/model_executor/layers/attention/mla_attention.py
+41
-20
@@ -202,13 +202,16 @@ from vllm._aiter_ops import rocm_aiter_ops
202
from vllm.config import ModelConfig, VllmConfig, get_current_vllm_config
203
from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank
204
from vllm.logger import init_logger
205
+from vllm.model_executor.custom_op import CustomOp
206
from vllm.model_executor.layers.batch_invariant import (
207
vllm_is_batch_invariant,
208
)
209
from vllm.model_executor.layers.linear import (
210
ColumnParallelLinear,
211
)
212
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
213
from vllm.model_executor.layers.quantization.utils.quant_utils import (
214
+ GroupShape,
215
get_and_maybe_dequant_weights,
216
)
217
from vllm.platforms import current_platform
@@ -287,6 +290,37 @@ def dynamic_per_batched_tensor_quant(
290
291
logger = init_logger(__name__)
292
293
+
294
+@CustomOp.register("mla_decode_concat_quant_fp8")
295
+class _DecodeConcatQuantFP8(QuantFP8):
296
+ """
297
+ QuantFP8 variant that concatenates decode_ql_nope and decode_q_pe before
298
+ quantization. When disabled, forward_native is compiled via torch.compile,
299
+ fusing cat/reshape/quant/view together.
300
+ """
301
+
302
+ def _make_forward(quant_fn): # noqa: N805
303
+ """Factory to create forward methods that concat before quantization."""
304
+
305
+ def forward(
306
+ self,
307
+ decode_ql_nope: torch.Tensor,
308
+ decode_q_pe: torch.Tensor,
309
+ scale: torch.Tensor,
310
+ scale_ub: torch.Tensor | None = None,
311
+ ) -> torch.Tensor:
312
+ decode_q0 = torch.cat((decode_ql_nope, decode_q_pe), dim=-1)
313
+ decode_q_flat = decode_q0.reshape(decode_q0.shape[0], -1)
314
+ decode_q, _ = quant_fn(self, decode_q_flat, scale, scale_ub)
315
+ return decode_q.view(decode_q0.shape)
316
+
317
+ return forward
318
+
319
+ forward_native = _make_forward(QuantFP8.forward_native) # type: ignore[arg-type]
320
+ forward_cuda = _make_forward(QuantFP8.forward_cuda) # type: ignore[arg-type]
321
+ forward_hip = _make_forward(QuantFP8.forward_hip) # type: ignore[arg-type]
322
+
323
+
324
CUDNN_WORKSPACE_SIZE = 12800
325
326
@@ -1398,6 +1432,11 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1432
self.cp_kv_cache_interleave_size: int = (
1433
get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size
1434
)
1435
+ self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(
1436
+ static=True,
1437
+ group_shape=GroupShape.PER_TENSOR,
1438
+ compile_native=True,
1439
+ )
1440
1441
def _flash_attn_varlen_diff_headdims(
1442
self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs
@@ -2048,29 +2087,11 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
2087
decode_ql_nope = decode_ql_nope.transpose(0, 1)
2088
2089
if fp8_attention:
2051
- ql_nope_shape = decode_ql_nope.shape
2052
- q_pe_shape = decode_q_pe.shape
2090
assert decode_ql_nope.shape[0] == decode_q_pe.shape[0]
2091
assert decode_ql_nope.shape[1] == decode_q_pe.shape[1]
2055
- decode_q_shape = (
2056
- ql_nope_shape[0],
2057
- ql_nope_shape[1],
2058
- ql_nope_shape[2] + q_pe_shape[2],
2059
- )
2060
- # Using empty and copy since torch.cat introduces significant overhead.
2061
- decode_q0 = torch.empty(
2062
- decode_q_shape,
2063
- device=decode_ql_nope.device,
2064
- dtype=decode_ql_nope.dtype,
2065
- )
2066
- decode_q0[..., : ql_nope_shape[2]].copy_(decode_ql_nope)
2067
- decode_q0[..., ql_nope_shape[2] :].copy_(decode_q_pe)
2068
-
2069
- decode_q, _ = ops.scaled_fp8_quant(
2070
- decode_q0.view(decode_q_shape[0], -1),
2071
- layer._q_scale,
2092
+ decode_q = self._decode_concat_quant_fp8_op(
2093
+ decode_ql_nope, decode_q_pe, layer._q_scale
2094
)
2073
- decode_q = decode_q.view(decode_q_shape)
2095
else:
2096
decode_q = (decode_ql_nope, decode_q_pe)
2097
if self.dcp_world_size > 1: