@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / aaa901ad55
Fork 0

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, 07:30 UTC

[Attention] Move MLA `forward` from backend to layer (#33284)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed Jan 30, 2026 at 22:30 UTC aaa901ad55ee454f061649b4cc24f96c5ea04e33
13 files changed +841 -623
tests/v1/attention/test_mla_backends.py
+247 -10
@@ -274,11 +274,157 @@ class MockAttentionLayer:
274 raise NotImplementedError
275
276
277 +class MockSparseMLAAttentionLayer:
278 + """A mock sparse MLA attention layer for testing.
279 +
280 + Sparse MLA implementations only support forward_mqa (decode-style attention)
281 + for all tokens, so this class only implements that path.
282 +
283 + Unlike regular MLA impls, sparse MLA impls don't have W_UK_T and W_UV
284 + attributes. These transformations are done by the layer (MLAAttention),
285 + not the impl. This mock layer accepts these weight matrices directly.
286 + """
287 +
288 + def __init__(
289 + self,
290 + impl,
291 + num_heads: int,
292 + qk_nope_head_dim: int,
293 + qk_rope_head_dim: int,
294 + v_head_dim: int,
295 + kv_lora_rank: int,
296 + device: torch.device,
297 + W_UK: torch.Tensor,
298 + W_UV: torch.Tensor,
299 + ):
300 + self.impl = impl
301 + self.num_heads = num_heads
302 + self.qk_nope_head_dim = qk_nope_head_dim
303 + self.qk_rope_head_dim = qk_rope_head_dim
304 + self.v_head_dim = v_head_dim
305 + self.kv_lora_rank = kv_lora_rank
306 +
307 + # Compute weight matrices in the format expected by forward_impl
308 + # W_UK shape: (L, N, P) -> W_UK_T shape: (N, P, L)
309 + self.W_UK_T = W_UK.permute(1, 2, 0)
310 + # W_UV shape: (L, N, V) -> (N, L, V)
311 + self.W_UV = W_UV.transpose(0, 1)
312 +
313 + # Scale attributes needed by attention backends
314 + self._q_scale = torch.tensor(1.0, device=device)
315 + self._k_scale = torch.tensor(1.0, device=device)
316 + self._v_scale = torch.tensor(1.0, device=device)
317 + self._prob_scale = torch.tensor(1.0, device=device)
318 + self._q_scale_float = 1.0
319 + self._k_scale_float = 1.0
320 + self._v_scale_float = 1.0
321 +
322 + def forward_impl(
323 + self,
324 + q: torch.Tensor,
325 + kv_c: torch.Tensor,
326 + k_pe: torch.Tensor,
327 + kv_cache: torch.Tensor,
328 + attn_metadata,
329 + output: torch.Tensor,
330 + ) -> torch.Tensor:
331 + """Forward for sparse MLA - uses forward_mqa for all tokens."""
332 + # Write to KV cache
333 + kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto")
334 + if kv_cache.numel() > 0:
335 + ops.concat_and_cache_mla(
336 + kv_c,
337 + k_pe.squeeze(1),
338 + kv_cache,
339 + attn_metadata.slot_mapping.flatten(),
340 + kv_cache_dtype=kv_cache_dtype,
341 + scale=self._k_scale,
342 + )
343 +
344 + num_tokens = q.shape[0]
345 +
346 + # Sparse MLA uses forward_mqa for all tokens
347 + # Split q into nope and pe parts
348 + mqa_q_nope, mqa_q_pe = q.split(
349 + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
350 + )
351 +
352 + # Convert from (B, N, P) to (N, B, P)
353 + mqa_q_nope = mqa_q_nope.transpose(0, 1)
354 +
355 + # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
356 + mqa_ql_nope = torch.bmm(mqa_q_nope, self.W_UK_T)
357 +
358 + # Convert from (N, B, L) to (B, N, L)
359 + mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
360 +
361 + # Pass as tuple to forward_mqa
362 + mqa_q = (mqa_ql_nope, mqa_q_pe)
363 +
364 + attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
365 +
366 + # v_up projection: multiply by W_UV
367 + # attn_out shape: (B, N, L) where L = kv_lora_rank
368 + # W_UV shape: (N, L, V)
369 + # output shape: (B, N, V) -> flatten to (B, N*V)
370 + decode_output = torch.bmm(attn_out.transpose(0, 1), self.W_UV).transpose(0, 1)
371 + output[:num_tokens] = decode_output.reshape(
372 + num_tokens, self.num_heads * self.v_head_dim
373 + )
374 +
375 + return output
376 +
377 +
378 class MockMLAAttentionLayer(AttentionLayerBase):
278 - """A mock MLA attention layer for populating static_forward_context."""
379 + """A mock MLA attention layer for testing.
380 +
381 + This replicates the forward_impl logic from MLAAttention to allow
382 + testing MLA backends without the full layer infrastructure.
383 +
384 + The W_UK_T and W_UV weight matrices are created on the layer (like in
385 + MLAAttention.process_weights_after_loading), not on the impl.
386 + """
387
280 - def __init__(self, impl):
388 + def __init__(
389 + self,
390 + impl,
391 + num_heads: int,
392 + qk_nope_head_dim: int,
393 + qk_rope_head_dim: int,
394 + v_head_dim: int,
395 + kv_lora_rank: int,
396 + device: torch.device,
397 + kv_b_proj,
398 + ):
399 self.impl = impl
400 + self.num_heads = num_heads
401 + self.qk_nope_head_dim = qk_nope_head_dim
402 + self.qk_rope_head_dim = qk_rope_head_dim
403 + self.v_head_dim = v_head_dim
404 + self.kv_lora_rank = kv_lora_rank
405 +
406 + # Compute weight matrices from kv_b_proj (like MLAAttention does)
407 + # This replicates MLAAttention.process_weights_after_loading logic
408 + kv_b_proj_weight = kv_b_proj.weight.T
409 + kv_b_proj_weight = kv_b_proj_weight.view(
410 + kv_lora_rank,
411 + num_heads,
412 + qk_nope_head_dim + v_head_dim,
413 + )
414 + W_UK, W_UV = kv_b_proj_weight.split([qk_nope_head_dim, v_head_dim], dim=-1)
415 + # Convert from (L, N, V) to (N, L, V)
416 + self.W_UV = W_UV.transpose(0, 1)
417 + # Convert from (L, N, P) to (N, P, L)
418 + self.W_UK_T = W_UK.permute(1, 2, 0)
419 +
420 + # Scale attributes needed by attention backends
421 + self._q_scale = torch.tensor(1.0, device=device)
422 + self._k_scale = torch.tensor(1.0, device=device)
423 + self._v_scale = torch.tensor(1.0, device=device)
424 + self._prob_scale = torch.tensor(1.0, device=device)
425 + self._q_scale_float = 1.0
426 + self._k_scale_float = 1.0
427 + self._v_scale_float = 1.0
428
429 def get_attn_backend(self):
430 raise NotImplementedError
@@ -286,6 +432,83 @@ class MockMLAAttentionLayer(AttentionLayerBase):
432 def get_kv_cache_spec(self, vllm_config):
433 raise NotImplementedError
434
435 + def forward_impl(
436 + self,
437 + q: torch.Tensor,
438 + kv_c: torch.Tensor,
439 + k_pe: torch.Tensor,
440 + kv_cache: torch.Tensor,
441 + attn_metadata,
442 + output: torch.Tensor,
443 + ) -> torch.Tensor:
444 + """Replicates MLAAttention.forward_impl logic for testing."""
445 + # Write to KV cache
446 + if kv_cache.numel() > 0:
447 + ops.concat_and_cache_mla(
448 + kv_c,
449 + k_pe.squeeze(1),
450 + kv_cache,
451 + attn_metadata.slot_mapping.flatten(),
452 + kv_cache_dtype="auto",
453 + scale=self._k_scale,
454 + )
455 +
456 + # Determine decode vs prefill split
457 + num_decode_tokens = attn_metadata.num_decode_tokens or 0
458 + has_decode = (attn_metadata.num_decodes or 0) > 0
459 + has_prefill = (attn_metadata.num_prefills or 0) > 0
460 +
461 + # Run prefill with forward_mha
462 + if has_prefill:
463 + prefill_q = q[num_decode_tokens:]
464 + prefill_k_pe = k_pe[num_decode_tokens:]
465 + prefill_k_c = kv_c[num_decode_tokens:]
466 + self.impl.forward_mha(
467 + prefill_q,
468 + prefill_k_c,
469 + prefill_k_pe,
470 + kv_cache,
471 + attn_metadata,
472 + self._k_scale,
473 + output=output[num_decode_tokens:],
474 + )
475 +
476 + # Run decode with forward_mqa
477 + if has_decode:
478 + decode_q = q[:num_decode_tokens]
479 +
480 + # Split q into nope and pe parts
481 + mqa_q_nope, mqa_q_pe = decode_q.split(
482 + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
483 + )
484 +
485 + # Convert from (B, N, P) to (N, B, P)
486 + mqa_q_nope = mqa_q_nope.transpose(0, 1)
487 +
488 + # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
489 + mqa_ql_nope = torch.bmm(mqa_q_nope, self.W_UK_T)
490 +
491 + # Convert from (N, B, L) to (B, N, L)
492 + mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
493 +
494 + # Pass as tuple to forward_mqa
495 + mqa_q = (mqa_ql_nope, mqa_q_pe)
496 +
497 + attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
498 +
499 + # v_up projection: multiply by W_UV
500 + # attn_out shape: (B, N, L) where L = kv_lora_rank
501 + # W_UV shape: (N, L, V)
502 + # output shape: (B, N, V) -> flatten to (B, N*V)
503 + decode_output = torch.bmm(attn_out.transpose(0, 1), self.W_UV).transpose(
504 + 0, 1
505 + )
506 + output[:num_decode_tokens] = decode_output.reshape(
507 + num_decode_tokens, self.num_heads * self.v_head_dim
508 + )
509 +
510 + return output
511 +
512
513 def run_attention_backend(
514 backend: AttentionBackendEnum,
@@ -340,14 +563,31 @@ def run_attention_backend(
563 kv_b_proj=mock_kv_b_proj,
564 )
565
343 - # Process weights to create W_UK_T and W_UV attributes needed by MLA
566 + # Process weights on the impl
567 act_dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)
568 impl.process_weights_after_loading(act_dtype)
569
570 + # Initialize DCP attributes (normally set by MLAAttention.forward
571 + # before calling forward_mha, see mla_attention.py:511-512)
572 + if impl.dcp_world_size == -1:
573 + impl.dcp_world_size = 1
574 +
575 + # Create mock MLA layer
576 + mock_layer = MockMLAAttentionLayer(
577 + impl=impl,
578 + num_heads=num_heads,
579 + qk_nope_head_dim=qk_nope_head_dim,
580 + qk_rope_head_dim=qk_rope_head_dim,
581 + v_head_dim=v_head_dim,
582 + kv_lora_rank=kv_lora_rank,
583 + device=device,
584 + kv_b_proj=mock_kv_b_proj,
585 + )
586 +
587 # Populate static_forward_context with mock attention layers
588 for layer_name in layer_names:
589 vllm_config.compilation_config.static_forward_context[layer_name] = (
350 - MockMLAAttentionLayer(impl)
590 + mock_layer
591 )
592
593 # Build metadata
@@ -357,18 +597,15 @@ def run_attention_backend(
597 common_attn_metadata=common_attn_metadata,
598 )
599
360 - # Create mock layer and output buffer
361 - mock_layer = MockAttentionLayer(device)
600 + # Create output buffer
601 num_tokens = query.shape[0]
602 output = torch.empty(
603 num_tokens, num_heads * v_head_dim, dtype=query.dtype, device=query.device
604 )
605
606 # Run forward pass
368 - # NOTE: The query, key, and value are already shaped correctly
369 - # in the calling test function.
370 - output = impl.forward(
371 - mock_layer, query, kv_c, k_pe, kv_cache, attn_metadata, output=output
607 + output = mock_layer.forward_impl(
608 + query, kv_c, k_pe, kv_cache, attn_metadata, output
609 )
610
611 return output
tests/v1/attention/test_sparse_mla_backends.py
+16 -5
@@ -12,7 +12,7 @@ import torch
12 from tests.v1.attention.test_mla_backends import (
13 BATCH_SPECS,
14 BatchSpec,
15 - MockAttentionLayer,
15 + MockSparseMLAAttentionLayer,
16 create_and_prepopulate_kv_cache,
17 )
18 from tests.v1.attention.utils import (
@@ -408,20 +408,31 @@ def test_sparse_backend_decode_correctness(
408
409 impl.process_weights_after_loading(dtype)
410
411 - layer = MockAttentionLayer(device)
411 + # Create mock sparse MLA layer with weight matrices
412 + mock_layer = MockSparseMLAAttentionLayer(
413 + impl=impl,
414 + num_heads=num_heads,
415 + qk_nope_head_dim=qk_nope_head_dim,
416 + qk_rope_head_dim=qk_rope_head_dim,
417 + v_head_dim=v_head_dim,
418 + kv_lora_rank=kv_lora_rank,
419 + device=device,
420 + W_UK=W_UK,
421 + W_UV=W_UV,
422 + )
423 +
424 out_buffer = torch.empty(
425 metadata.num_actual_tokens, num_heads * v_head_dim, dtype=dtype, device=device
426 )
427
428 with torch.inference_mode():
417 - backend_output = impl.forward(
418 - layer,
429 + backend_output = mock_layer.forward_impl(
430 query_vllm,
431 kv_c_vllm,
432 k_pe_vllm,
433 kv_cache,
434 metadata,
424 - output=out_buffer,
435 + out_buffer,
436 )
437
438 assert backend_output.shape == sdpa_reference.shape
vllm/model_executor/layers/attention/attention.py
+2 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import TYPE_CHECKING
4 +from typing import TYPE_CHECKING, Any
5
6 import torch
7 import torch.nn as nn
@@ -562,7 +562,7 @@ direct_register_custom_op(
562
563 def get_attention_context(
564 layer_name: str,
565 -) -> tuple[dict | object | None, "Attention | MLAAttention", torch.Tensor]:
565 +) -> tuple[Any, "Attention | MLAAttention", torch.Tensor]:
566 """Extract attention context for a given layer.
567
568 This helper function extracts the attention metadata, attention layer
vllm/model_executor/layers/attention/mla_attention.py
+443 -440
@@ -63,7 +63,7 @@ W_UV project kv_c to v shape [Lkv, N, V]
63 W_O project v to h_t shape [N * V, H]
64
65
66 -## Compute Friendly Approach (i.e. "_forward_prefill"):
66 +## Compute Friendly Approach (i.e. "forward_mha"):
67
68 q_c = h_t @ W_DQ
69 q_nope = (q_c @ W_UQ).view(Sq, N, P)
@@ -91,7 +91,7 @@ NOTE: in the actual code,
91 `out_proj` is W_O
92
93
94 -## Data-Movement Friendly Approach (i.e. "_forward_decode"):
94 +## Data-Movement Friendly Approach (i.e. "forward_mqa"):
95
96 Runtime
97 q_c = h_t @ W_DQ
@@ -243,6 +243,7 @@ from vllm.v1.attention.backend import (
243 AttentionType,
244 CommonAttentionMetadata,
245 MLAAttentionImpl,
246 + SparseMLAAttentionImpl,
247 )
248 from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
249 from vllm.v1.attention.backends.utils import (
@@ -266,6 +267,9 @@ logger = init_logger(__name__)
267 class MLAAttention(nn.Module, AttentionLayerBase):
268 """Multi-Head Latent Attention layer.
269
270 + NOTE: Please read the comment at the top of the file before trying to
271 + understand this class
272 +
273 This class takes query, and compressed key/value tensors as input.
274 The class does the following:
275
@@ -289,6 +293,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
293 prefix: str = "",
294 use_sparse: bool = False,
295 indexer: object | None = None,
296 + q_pad_num_heads: int | None = None,
297 **extra_impl_args,
298 ):
299 super().__init__()
@@ -299,8 +304,14 @@ class MLAAttention(nn.Module, AttentionLayerBase):
304 self.v_head_dim = v_head_dim
305 self.q_lora_rank = q_lora_rank
306 self.kv_lora_rank = kv_lora_rank
307 + self.kv_b_proj = kv_b_proj
308 self.head_size = kv_lora_rank + qk_rope_head_dim
309 self.layer_name = prefix
310 + self.indexer = indexer
311 + self.q_pad_num_heads = q_pad_num_heads
312 +
313 + self.num_kv_heads = 1
314 + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
315
316 if cache_config is not None:
317 kv_cache_dtype = cache_config.cache_dtype
@@ -364,6 +375,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
375 v_head_dim=self.v_head_dim,
376 kv_b_proj=kv_b_proj,
377 indexer=indexer,
378 + q_pad_num_heads=q_pad_num_heads,
379 **extra_impl_args,
380 )
381
@@ -388,6 +400,26 @@ class MLAAttention(nn.Module, AttentionLayerBase):
400 self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
401 self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)
402
403 + self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
404 +
405 + # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported
406 + self.is_aiter_triton_fp4_bmm_enabled = (
407 + rocm_aiter_ops.is_fp4bmm_enabled()
408 + and self.kv_b_proj.weight.dtype == torch.bfloat16
409 + )
410 +
411 + # Attributes for forward_impl method
412 + self.chunked_prefill_workspace_size = (
413 + MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(
414 + get_current_vllm_config()
415 + )
416 + )
417 + self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(
418 + static=True,
419 + group_shape=GroupShape.PER_TENSOR,
420 + compile_native=True,
421 + )
422 +
423 def forward(
424 self,
425 q: torch.Tensor,
@@ -407,8 +439,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
439
440 if self.attn_backend.accept_output_buffer:
441 output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
410 - self.impl.forward(
411 - self,
442 + self.forward_impl(
443 q,
444 kv_c_normed,
445 k_pe,
@@ -418,8 +449,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
449 )
450 return output
451 else:
421 - return self.impl.forward(
422 - self, q, kv_c_normed, k_pe, self_kv_cache, attn_metadata
452 + return self.forward_impl(
453 + q, kv_c_normed, k_pe, self_kv_cache, attn_metadata
454 )
455 else:
456 if self.attn_backend.accept_output_buffer:
@@ -440,9 +471,282 @@ class MLAAttention(nn.Module, AttentionLayerBase):
471 self.layer_name,
472 )
473
474 + def forward_impl(
475 + self,
476 + q: torch.Tensor,
477 + k_c_normed: torch.Tensor, # key in unified attn
478 + k_pe: torch.Tensor, # value in unified attn
479 + kv_cache: torch.Tensor,
480 + attn_metadata: "MLACommonMetadata",
481 + output: torch.Tensor | None = None,
482 + output_scale: torch.Tensor | None = None,
483 + output_block_scale: torch.Tensor | None = None,
484 + ) -> torch.Tensor:
485 + assert output is not None, "Output tensor must be provided."
486 +
487 + if output_scale is not None or output_block_scale is not None:
488 + raise NotImplementedError(
489 + "fused output quantization is not yet supported for MLA"
490 + )
491 +
492 + if attn_metadata is None:
493 + # During the profile run try to simulate to worse case output size
494 + # for `self.kv_b_proj(kv_c_normed)` in `_compute_prefill_context`
495 + # since this can be large
496 + _ = torch.empty(
497 + (
498 + self.chunked_prefill_workspace_size,
499 + self.num_heads,
500 + self.qk_nope_head_dim + self.v_head_dim,
501 + ),
502 + device=k_c_normed.device,
503 + dtype=k_c_normed.dtype,
504 + )
505 +
506 + # The zero fill is required when used with DP + EP
507 + # to ensure all ranks within a DP group compute the
508 + # same expert outputs.
509 + return output.fill_(0)
510 +
511 + if self.impl.dcp_world_size == -1:
512 + self.impl.dcp_world_size = get_dcp_group().world_size
513 +
514 + fp8_attention = self.kv_cache_dtype.startswith("fp8")
515 +
516 + num_actual_toks = attn_metadata.num_actual_tokens
517 +
518 + # Inputs and outputs may be padded for CUDA graphs
519 + output_padded = output
520 + output = output[:num_actual_toks, ...]
521 + q = q[:num_actual_toks, ...]
522 + k_c_normed = k_c_normed[:num_actual_toks, ...]
523 + k_pe = k_pe[:num_actual_toks, ...]
524 +
525 + assert (
526 + attn_metadata.num_decodes is not None
527 + and attn_metadata.num_prefills is not None
528 + and attn_metadata.num_decode_tokens is not None
529 + )
530 +
531 + has_decode = attn_metadata.num_decodes > 0
532 + has_prefill = attn_metadata.num_prefills > 0
533 + num_decode_tokens = attn_metadata.num_decode_tokens
534 +
535 + decode_q = q[:num_decode_tokens]
536 +
537 + prefill_q = q[num_decode_tokens:]
538 + prefill_k_pe = k_pe[num_decode_tokens:]
539 + prefill_k_c_normed = k_c_normed[num_decode_tokens:]
540 +
541 + # write the latent and rope to kv cache
542 + if kv_cache.numel() > 0:
543 + ops.concat_and_cache_mla(
544 + k_c_normed,
545 + k_pe.squeeze(1),
546 + kv_cache,
547 + attn_metadata.slot_mapping.flatten(),
548 + kv_cache_dtype=self.kv_cache_dtype,
549 + scale=self._k_scale,
550 + )
551 +
552 + if fp8_attention:
553 + kv_cache = kv_cache.view(current_platform.fp8_dtype())
554 +
555 + # Sparse MLA impls only support forward_mqa (decode-style attention)
556 + is_sparse_impl = isinstance(self.impl, SparseMLAAttentionImpl)
557 +
558 + if has_prefill and not is_sparse_impl:
559 + self.impl.forward_mha(
560 + prefill_q,
561 + prefill_k_c_normed,
562 + prefill_k_pe,
563 + kv_cache,
564 + attn_metadata,
565 + self._k_scale,
566 + output=output[num_decode_tokens:],
567 + )
568 +
569 + if has_decode or (has_prefill and is_sparse_impl):
570 + # For sparse impl, we always use forward_mqa for all tokens
571 + # For non-sparse impl, we only use forward_mqa for decode tokens
572 + if is_sparse_impl:
573 + mqa_q = q
574 + mqa_output_slice = output
575 + else:
576 + assert attn_metadata.decode is not None
577 + mqa_q = decode_q
578 + mqa_output_slice = output[:num_decode_tokens]
579 +
580 + mqa_q_nope, mqa_q_pe = mqa_q.split(
581 + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
582 + )
583 +
584 + # Convert from (B, N, P) to (N, B, P)
585 + mqa_q_nope = mqa_q_nope.transpose(0, 1)
586 +
587 + if self.q_pad_num_heads is not None:
588 + B, N, L = mqa_q_pe.shape
589 + mqa_pe_padded = mqa_q_pe.new_empty((B, self.q_pad_num_heads, L))
590 + mqa_pe_padded.resize_((B, N, L))
591 + mqa_pe_padded.copy_(mqa_q_pe)
592 + mqa_q_pe = mqa_pe_padded
593 +
594 + if self.is_aiter_triton_fp4_bmm_enabled:
595 + from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
596 +
597 + mqa_ql_nope = batched_gemm_a16wfp4(
598 + mqa_q_nope,
599 + self.W_K,
600 + self.W_K_scale,
601 + transpose_bm=True,
602 + prequant=True,
603 + y_scale=self._q_scale if fp8_attention else None,
604 + )
605 + elif self.is_aiter_triton_fp8_bmm_enabled:
606 + # Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L)
607 + mqa_ql_nope = rocm_aiter_ops.triton_fp8_bmm(
608 + mqa_q_nope,
609 + self.W_K,
610 + self.W_K_scale,
611 + group_size=128,
612 + transpose_bm=True,
613 + )
614 + else:
615 + # Pads the head_dim if necessary (for the underlying kernel)
616 + N, B, P = mqa_q_nope.shape
617 + _, _, L = self.W_UK_T.shape
618 +
619 + if self.q_pad_num_heads is not None:
620 + mqa_ql_nope = mqa_q_nope.new_empty((self.q_pad_num_heads, B, L))
621 + mqa_ql_nope.resize_((N, B, L))
622 + else:
623 + mqa_ql_nope = mqa_q_nope.new_empty((N, B, L))
624 +
625 + # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
626 + torch.bmm(mqa_q_nope, self.W_UK_T, out=mqa_ql_nope)
627 +
628 + # Convert from (N, B, L) to (B, N, L)
629 + mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
630 +
631 + if fp8_attention:
632 + assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0]
633 + assert mqa_ql_nope.shape[1] == mqa_q_pe.shape[1]
634 + mqa_q = self._decode_concat_quant_fp8_op(
635 + mqa_ql_nope, mqa_q_pe, self._q_scale
636 + )
637 + else:
638 + mqa_q = (mqa_ql_nope, mqa_q_pe)
639 + if self.impl.dcp_world_size > 1:
640 + assert not fp8_attention, "DCP not support fp8 kvcache now."
641 + # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P)
642 + mqa_q = torch.cat(mqa_q, dim=-1)
643 + # mqa_q do allgather in head dim.
644 + mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)
645 +
646 + # call decode attn
647 + attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
648 +
649 + # correct dcp attn_out with lse.
650 + if self.impl.dcp_world_size > 1:
651 + attn_out = cp_lse_ag_out_rs(
652 + attn_out,
653 + lse,
654 + get_dcp_group(),
655 + is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),
656 + )
657 +
658 + # v_up projection
659 + self._v_up_proj(attn_out, out=mqa_output_slice)
660 + return output_padded
661 +
662 def process_weights_after_loading(self, act_dtype: torch.dtype):
444 - if hasattr(self.impl, "process_weights_after_loading"):
445 - self.impl.process_weights_after_loading(act_dtype)
663 + # we currently do not have quantized bmm's which are needed for
664 + # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform
665 + # the bmm's in 16-bit, the extra memory overhead of this is fairly low
666 + kv_b_proj_weight = get_and_maybe_dequant_weights(
667 + self.kv_b_proj, out_dtype=act_dtype
668 + ).T
669 +
670 + assert kv_b_proj_weight.shape == (
671 + self.kv_lora_rank,
672 + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
673 + ), (
674 + f"{kv_b_proj_weight.shape=}, "
675 + f"{self.kv_lora_rank=}, "
676 + f"{self.num_heads=}, "
677 + f"{self.qk_nope_head_dim=}, "
678 + f"{self.v_head_dim=}"
679 + )
680 + kv_b_proj_weight = kv_b_proj_weight.view(
681 + self.kv_lora_rank,
682 + self.num_heads,
683 + self.qk_nope_head_dim + self.v_head_dim,
684 + )
685 +
686 + W_UK, W_UV = kv_b_proj_weight.split(
687 + [self.qk_nope_head_dim, self.v_head_dim], dim=-1
688 + )
689 +
690 + # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported
691 + if self.is_aiter_triton_fp4_bmm_enabled:
692 + from vllm.model_executor.layers.quantization.quark.utils import (
693 + quark_quantize_weight_to_mxfp4,
694 + )
695 +
696 + self.W_K, self.W_K_scale = quark_quantize_weight_to_mxfp4(W_UK)
697 + # Convert from (L, N, P) to (N, L, P)
698 + self.W_K = self.W_K.transpose(0, 1)
699 + self.W_K_scale = self.W_K_scale.transpose(0, 1)
700 +
701 + self.W_V, self.W_V_scale = quark_quantize_weight_to_mxfp4(
702 + W_UV.permute(1, 2, 0)
703 + )
704 + elif self.is_aiter_triton_fp8_bmm_enabled:
705 + W_K = W_UK.transpose(0, 1) # 16 512 128
706 + W_V = W_UV.permute(1, 2, 0) # 16 128 512
707 + self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(
708 + W_K, dtype=current_platform.fp8_dtype()
709 + )
710 + self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(
711 + W_V, dtype=current_platform.fp8_dtype()
712 + )
713 +
714 + # The kernel operates on non-padded inputs. Hence, pre-compiling
715 + # triton kernel to avoid runtime compilation for unseen batch sizes
716 + # Pre-compile for batch sizes 1 to 1024 to cover most use-cases.
717 + # On DS-R1, this step adds roughly 50s to the model loading time.
718 + max_batch_size = 1024 # [ToDo] Find the optimal upper limit
719 + pre_compilation_list = list(range(1, max_batch_size + 1))
720 + if is_global_first_rank():
721 + pre_compilation_list = tqdm(
722 + pre_compilation_list,
723 + desc="[Aiter Triton] Pre-compiling fp8 BMM kernel",
724 + total=max_batch_size,
725 + )
726 +
727 + for m in pre_compilation_list:
728 + x = torch.empty(
729 + (self.W_K.shape[0], m, self.W_K.shape[2]),
730 + dtype=torch.bfloat16,
731 + device=self.W_K.device,
732 + )
733 + rocm_aiter_ops.triton_fp8_bmm(
734 + x, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
735 + )
736 +
737 + x = torch.empty(
738 + (self.W_V.shape[0], m, self.W_V.shape[2]),
739 + dtype=torch.bfloat16,
740 + device=self.W_V.device,
741 + )
742 + rocm_aiter_ops.triton_fp8_bmm(
743 + x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True
744 + )
745 + else:
746 + # Convert from (L, N, V) to (N, L, V)
747 + self.W_UV = W_UV.transpose(0, 1)
748 + # Convert from (L, N, P) to (N, P, L)
749 + self.W_UK_T = W_UK.permute(1, 2, 0)
750
751 # If we should not load quant weights, we initialize the scales to 1.0
752 # as the default value. See [Note: Register q/k/v/prob scales in state dict]
@@ -492,6 +796,41 @@ class MLAAttention(nn.Module, AttentionLayerBase):
796 cache_dtype_str=vllm_config.cache_config.cache_dtype,
797 )
798
799 + def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
800 + # Convert from (B, N, L) to (N, B, L)
801 + x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
802 + out = out.view(-1, self.num_heads, self.v_head_dim)
803 + if self.is_aiter_triton_fp4_bmm_enabled:
804 + out = rocm_aiter_ops.batched_gemm_a16wfp4(
805 + x,
806 + self.W_V,
807 + self.W_V_scale,
808 + out,
809 + transpose_bm=True,
810 + prequant=True,
811 + y_scale=None,
812 + )
813 + x = out.view(-1, self.num_heads * self.v_head_dim)
814 + elif self.is_aiter_triton_fp8_bmm_enabled:
815 + # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
816 + x = rocm_aiter_ops.triton_fp8_bmm(
817 + x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True, YQ=out
818 + )
819 + else:
820 + # Convert from (B, N * V) to (N, B, V)
821 + out = out.transpose(0, 1)
822 +
823 + # Multiply (N, B, L) x (N, L, V) -> (N, B, V)
824 + torch.bmm(x, self.W_UV, out=out) # Reuse "out" to make it "hot"
825 +
826 + # Convert from (N, B, V) to (B, N * V)
827 + out_new = out.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)
828 +
829 + # Adjust output buffer shape back to the original (B, N * V)
830 + N, B, V = out.shape
831 + out.resize_((B, N * V))
832 + out.copy_(out_new) # Copy result
833 +
834
835 @maybe_transfer_kv_layer
836 def unified_mla_attention(
@@ -500,8 +839,8 @@ def unified_mla_attention(
839 k_pe: torch.Tensor,
840 layer_name: str,
841 ) -> torch.Tensor:
503 - attn_metadata, self, kv_cache = get_attention_context(layer_name)
504 - output = self.impl.forward(self, q, kv_c_normed, k_pe, kv_cache, attn_metadata)
842 + attn_metadata, layer, kv_cache = get_attention_context(layer_name)
843 + output = layer.forward_impl(q, kv_c_normed, k_pe, kv_cache, attn_metadata)
844
845 return output
846
@@ -534,9 +873,8 @@ def unified_mla_attention_with_output(
873 output_scale: torch.Tensor | None = None,
874 output_block_scale: torch.Tensor | None = None,
875 ) -> None:
537 - attn_metadata, self, kv_cache = get_attention_context(layer_name)
538 - self.impl.forward(
539 - self,
876 + attn_metadata, layer, kv_cache = get_attention_context(layer_name)
877 + layer.forward_impl(
878 q,
879 kv_c_normed,
880 k_pe,
@@ -1460,247 +1798,104 @@ def reorg_kvcache(
1798 padded_local_chunk_seq_lens_lst: local chunk context lengths
1799 under current CP rank.
1800 local_context_lens_allranks: local context lengths on each CP rank.
1463 - sum_seq_len: the sum of cp_chunk_seq_lens_lst.
1464 - max_seq_len: the max value of cp_chunk_seq_lens_lst.
1465 - chunk_size: the local padded max context chunk from
1466 - chunked_context_metadata building.
1467 - chunk_idx: chunk idx of chunked_prefill.
1468 - toks: the number of tokens for local gather cache.
1469 - """
1470 - kv_c_segments = []
1471 - k_pe_segments = []
1472 - src_token_idx = 0
1473 - max_seq_len_check = 0
1474 - for padded_local_chunk_seq_len, local_context_lens in zip(
1475 - padded_local_chunk_seq_lens_lst, local_context_lens_allranks
1476 - ):
1477 - cur_seq_len = 0
1478 - for rank, local_context_len in enumerate(local_context_lens):
1479 - # Note(qcs): We split the context into multiple chunks,
1480 - # depending on the size of the workspace.
1481 - # local_context in dcp0: |-----------------|
1482 - # local_context in dcp1: |--------------|
1483 - # n*padded_local_chunk: |-----|-----|-----|
1484 - # local_chunk_len in dcp1: |-----|-----|--|
1485 - # so we need update the last chunk length in dcp1.
1486 - local_chunk_len = min(
1487 - max(0, local_context_len - chunk_idx * chunk_size),
1488 - padded_local_chunk_seq_len,
1489 - )
1490 - if local_chunk_len != 0:
1491 - kv_c_segment = allgatered_kv_c_normed[
1492 - rank * toks + src_token_idx : rank * toks
1493 - + src_token_idx
1494 - + local_chunk_len
1495 - ]
1496 - k_pe_segment = allgatered_k_pe[
1497 - rank * toks + src_token_idx : rank * toks
1498 - + src_token_idx
1499 - + local_chunk_len
1500 - ]
1501 - kv_c_segments.append(kv_c_segment)
1502 - k_pe_segments.append(k_pe_segment)
1503 - cur_seq_len += local_chunk_len
1504 - max_seq_len_check = max(max_seq_len_check, cur_seq_len)
1505 - src_token_idx += padded_local_chunk_seq_len
1506 - reorganized_kv_c_normed = torch.cat(kv_c_segments, dim=0)
1507 - reorganized_k_pe = torch.cat(k_pe_segments, dim=0)
1508 - assert reorganized_kv_c_normed.shape[0] == sum_seq_len
1509 - assert reorganized_k_pe.shape[0] == sum_seq_len
1510 - assert max_seq_len_check == max_seq_len
1511 - return reorganized_kv_c_normed, reorganized_k_pe
1512 -
1513 -
1514 -# TODO(Lucas): rename MLACommonBaseImpl -> MLACommonImpl,
1515 -# and MLACommonImpl -> MLACommonDenseImpl or somthing like that
1516 -class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
1517 - """
1518 - NOTE: Please read the comment at the top of the file before trying to
1519 - understand this class
1520 - """
1521 -
1522 - def __init__(
1523 - self,
1524 - num_heads: int,
1525 - head_size: int,
1526 - scale: float,
1527 - num_kv_heads: int,
1528 - alibi_slopes: list[float] | None,
1529 - sliding_window: int | None,
1530 - kv_cache_dtype: str,
1531 - logits_soft_cap: float | None,
1532 - attn_type: str,
1533 - kv_sharing_target_layer_name: str | None,
1534 - # MLA Specific Arguments
1535 - q_lora_rank: int | None,
1536 - kv_lora_rank: int,
1537 - qk_nope_head_dim: int,
1538 - qk_rope_head_dim: int,
1539 - qk_head_dim: int,
1540 - v_head_dim: int,
1541 - kv_b_proj: ColumnParallelLinear,
1542 - indexer=None,
1543 - q_pad_num_heads: int | None = None,
1544 - ) -> None:
1545 - if kv_sharing_target_layer_name is not None:
1546 - raise NotImplementedError("KV sharing is not supported for MLA")
1547 -
1548 - self.num_heads = num_heads
1549 - self.head_size = head_size
1550 - self.scale = float(scale)
1551 - self.num_kv_heads = num_kv_heads
1552 - self.kv_cache_dtype = kv_cache_dtype
1553 -
1554 - self.q_lora_rank = q_lora_rank
1555 - self.kv_lora_rank = kv_lora_rank
1556 - self.qk_nope_head_dim = qk_nope_head_dim
1557 - self.qk_rope_head_dim = qk_rope_head_dim
1558 - self.qk_head_dim = qk_head_dim
1559 - self.v_head_dim = v_head_dim
1560 - self.kv_b_proj = kv_b_proj
1561 - self.indexer = indexer
1562 - self.q_pad_num_heads = q_pad_num_heads
1563 - self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
1564 -
1565 - # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported
1566 - self.is_aiter_triton_fp4_bmm_enabled = (
1567 - rocm_aiter_ops.is_fp4bmm_enabled()
1568 - and self.kv_b_proj.weight.dtype == torch.bfloat16
1569 - )
1570 -
1571 - def process_weights_after_loading(self, act_dtype: torch.dtype):
1572 - # we currently do not have quantized bmm's which are needed for
1573 - # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform
1574 - # the bmm's in 16-bit, the extra memory overhead of this is fairly low
1575 - kv_b_proj_weight = get_and_maybe_dequant_weights(
1576 - self.kv_b_proj, out_dtype=act_dtype
1577 - ).T
1578 -
1579 - assert kv_b_proj_weight.shape == (
1580 - self.kv_lora_rank,
1581 - self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
1582 - ), (
1583 - f"{kv_b_proj_weight.shape=}, "
1584 - f"{self.kv_lora_rank=}, "
1585 - f"{self.num_heads=}, "
1586 - f"{self.qk_nope_head_dim=}, "
1587 - f"{self.v_head_dim=}"
1588 - )
1589 - kv_b_proj_weight = kv_b_proj_weight.view(
1590 - self.kv_lora_rank,
1591 - self.num_heads,
1592 - self.qk_nope_head_dim + self.v_head_dim,
1593 - )
1594 -
1595 - W_UK, W_UV = kv_b_proj_weight.split(
1596 - [self.qk_nope_head_dim, self.v_head_dim], dim=-1
1597 - )
1598 -
1599 - # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported
1600 - if self.is_aiter_triton_fp4_bmm_enabled:
1601 - from vllm.model_executor.layers.quantization.quark.utils import (
1602 - quark_quantize_weight_to_mxfp4,
1603 - )
1604 -
1605 - self.W_K, self.W_K_scale = quark_quantize_weight_to_mxfp4(W_UK)
1606 - # Convert from (L, N, P) to (N, L, P)
1607 - self.W_K = self.W_K.transpose(0, 1)
1608 - self.W_K_scale = self.W_K_scale.transpose(0, 1)
1609 -
1610 - self.W_V, self.W_V_scale = quark_quantize_weight_to_mxfp4(
1611 - W_UV.permute(1, 2, 0)
1612 - )
1613 - elif self.is_aiter_triton_fp8_bmm_enabled:
1614 - W_K = W_UK.transpose(0, 1) # 16 512 128
1615 - W_V = W_UV.permute(1, 2, 0) # 16 128 512
1616 - self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(
1617 - W_K, dtype=current_platform.fp8_dtype()
1618 - )
1619 - self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(
1620 - W_V, dtype=current_platform.fp8_dtype()
1621 - )
1622 -
1623 - # The kernel operates on non-padded inputs. Hence, pre-compiling
1624 - # triton kernel to avoid runtime compilation for unseen batch sizes
1625 - # Pre-compile for batch sizes 1 to 1024 to cover most use-cases.
1626 - # On DS-R1, this step adds roughly 50s to the model loading time.
1627 - max_batch_size = 1024 # [ToDo] Find the optimal upper limit
1628 - pre_compilation_list = list(range(1, max_batch_size + 1))
1629 - if is_global_first_rank():
1630 - pre_compilation_list = tqdm(
1631 - pre_compilation_list,
1632 - desc="[Aiter Triton] Pre-compiling fp8 BMM kernel",
1633 - total=max_batch_size,
1634 - )
1635 -
1636 - for m in pre_compilation_list:
1637 - x = torch.empty(
1638 - (self.W_K.shape[0], m, self.W_K.shape[2]),
1639 - dtype=torch.bfloat16,
1640 - device=self.W_K.device,
1641 - )
1642 - rocm_aiter_ops.triton_fp8_bmm(
1643 - x, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
1644 - )
1645 -
1646 - x = torch.empty(
1647 - (self.W_V.shape[0], m, self.W_V.shape[2]),
1648 - dtype=torch.bfloat16,
1649 - device=self.W_V.device,
1650 - )
1651 - rocm_aiter_ops.triton_fp8_bmm(
1652 - x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True
1653 - )
1654 - else:
1655 - # Convert from (L, N, V) to (N, L, V)
1656 - self.W_UV = W_UV.transpose(0, 1)
1657 - # Convert from (L, N, P) to (N, P, L)
1658 - self.W_UK_T = W_UK.permute(1, 2, 0)
1659 -
1660 - def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
1661 - # Convert from (B, N, L) to (N, B, L)
1662 - x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
1663 - out = out.view(-1, self.num_heads, self.v_head_dim)
1664 - if self.is_aiter_triton_fp4_bmm_enabled:
1665 - out = rocm_aiter_ops.batched_gemm_a16wfp4(
1666 - x,
1667 - self.W_V,
1668 - self.W_V_scale,
1669 - out,
1670 - transpose_bm=True,
1671 - prequant=True,
1672 - y_scale=None,
1673 - )
1674 - x = out.view(-1, self.num_heads * self.v_head_dim)
1675 - elif self.is_aiter_triton_fp8_bmm_enabled:
1676 - # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
1677 - x = rocm_aiter_ops.triton_fp8_bmm(
1678 - x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True, YQ=out
1801 + sum_seq_len: the sum of cp_chunk_seq_lens_lst.
1802 + max_seq_len: the max value of cp_chunk_seq_lens_lst.
1803 + chunk_size: the local padded max context chunk from
1804 + chunked_context_metadata building.
1805 + chunk_idx: chunk idx of chunked_prefill.
1806 + toks: the number of tokens for local gather cache.
1807 + """
1808 + kv_c_segments = []
1809 + k_pe_segments = []
1810 + src_token_idx = 0
1811 + max_seq_len_check = 0
1812 + for padded_local_chunk_seq_len, local_context_lens in zip(
1813 + padded_local_chunk_seq_lens_lst, local_context_lens_allranks
1814 + ):
1815 + cur_seq_len = 0
1816 + for rank, local_context_len in enumerate(local_context_lens):
1817 + # Note(qcs): We split the context into multiple chunks,
1818 + # depending on the size of the workspace.
1819 + # local_context in dcp0: |-----------------|
1820 + # local_context in dcp1: |--------------|
1821 + # n*padded_local_chunk: |-----|-----|-----|
1822 + # local_chunk_len in dcp1: |-----|-----|--|
1823 + # so we need update the last chunk length in dcp1.
1824 + local_chunk_len = min(
1825 + max(0, local_context_len - chunk_idx * chunk_size),
1826 + padded_local_chunk_seq_len,
1827 )
1680 - else:
1681 - # Convert from (B, N * V) to (N, B, V)
1682 - out = out.transpose(0, 1)
1683 -
1684 - # Multiply (N, B, L) x (N, L, V) -> (N, B, V)
1685 - torch.bmm(x, self.W_UV, out=out) # Reuse "out" to make it "hot"
1686 -
1687 - # Convert from (N, B, V) to (B, N * V)
1688 - out_new = out.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)
1689 -
1690 - # Adjust output buffer shape back to the original (B, N * V)
1691 - N, B, V = out.shape
1692 - out.resize_((B, N * V))
1693 - out.copy_(out_new) # Copy result
1828 + if local_chunk_len != 0:
1829 + kv_c_segment = allgatered_kv_c_normed[
1830 + rank * toks + src_token_idx : rank * toks
1831 + + src_token_idx
1832 + + local_chunk_len
1833 + ]
1834 + k_pe_segment = allgatered_k_pe[
1835 + rank * toks + src_token_idx : rank * toks
1836 + + src_token_idx
1837 + + local_chunk_len
1838 + ]
1839 + kv_c_segments.append(kv_c_segment)
1840 + k_pe_segments.append(k_pe_segment)
1841 + cur_seq_len += local_chunk_len
1842 + max_seq_len_check = max(max_seq_len_check, cur_seq_len)
1843 + src_token_idx += padded_local_chunk_seq_len
1844 + reorganized_kv_c_normed = torch.cat(kv_c_segments, dim=0)
1845 + reorganized_k_pe = torch.cat(k_pe_segments, dim=0)
1846 + assert reorganized_kv_c_normed.shape[0] == sum_seq_len
1847 + assert reorganized_k_pe.shape[0] == sum_seq_len
1848 + assert max_seq_len_check == max_seq_len
1849 + return reorganized_kv_c_normed, reorganized_k_pe
1850
1851
1696 -class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1852 +class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1853 """
1854 NOTE: Please read the comment at the top of the file before trying to
1855 understand this class
1856 """
1857
1702 - def __init__(self, *args, **kwargs) -> None:
1703 - super().__init__(*args, **kwargs)
1858 + def __init__(
1859 + self,
1860 + num_heads: int,
1861 + head_size: int,
1862 + scale: float,
1863 + num_kv_heads: int,
1864 + alibi_slopes: list[float] | None,
1865 + sliding_window: int | None,
1866 + kv_cache_dtype: str,
1867 + logits_soft_cap: float | None,
1868 + attn_type: str,
1869 + kv_sharing_target_layer_name: str | None,
1870 + # MLA Specific Arguments
1871 + q_lora_rank: int | None,
1872 + kv_lora_rank: int,
1873 + qk_nope_head_dim: int,
1874 + qk_rope_head_dim: int,
1875 + qk_head_dim: int,
1876 + v_head_dim: int,
1877 + kv_b_proj: ColumnParallelLinear,
1878 + indexer: object | None = None,
1879 + q_pad_num_heads: int | None = None,
1880 + ) -> None:
1881 + if kv_sharing_target_layer_name is not None:
1882 + raise NotImplementedError("KV sharing is not supported for MLA")
1883 +
1884 + self.num_heads = num_heads
1885 + self.head_size = head_size
1886 + self.scale = float(scale)
1887 + self.num_kv_heads = num_kv_heads
1888 + self.kv_cache_dtype = kv_cache_dtype
1889 +
1890 + self.q_lora_rank = q_lora_rank
1891 + self.kv_lora_rank = kv_lora_rank
1892 + self.qk_nope_head_dim = qk_nope_head_dim
1893 + self.qk_rope_head_dim = qk_rope_head_dim
1894 + self.qk_head_dim = qk_head_dim
1895 + self.v_head_dim = v_head_dim
1896 + self.kv_b_proj = kv_b_proj
1897 + self.indexer = indexer
1898 + self.q_pad_num_heads = q_pad_num_heads
1899
1900 if use_trtllm_ragged_deepseek_prefill():
1901 logger.info_once(
@@ -1750,19 +1945,9 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1945
1946 self.dcp_world_size: int = -1
1947
1753 - self.chunked_prefill_workspace_size = (
1754 - MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(
1755 - get_current_vllm_config()
1756 - )
1757 - )
1948 self.cp_kv_cache_interleave_size: int = (
1949 get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size
1950 )
1761 - self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(
1762 - static=True,
1763 - group_shape=GroupShape.PER_TENSOR,
1764 - compile_native=True,
1765 - )
1951
1952 def _flash_attn_varlen_diff_headdims(
1953 self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs
@@ -2193,7 +2378,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
2378
2379 return output, output_lse
2380
2196 - def _forward_prefill(
2381 + def forward_mha(
2382 self,
2383 q: torch.Tensor,
2384 kv_c_normed: torch.Tensor,
@@ -2258,7 +2443,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
2443 output.copy_(output_prefill)
2444
2445 @abstractmethod
2261 - def _forward_decode(
2446 + def forward_mqa(
2447 self,
2448 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
2449 kv_c_and_k_pe_cache: torch.Tensor,
@@ -2266,185 +2451,3 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
2451 layer: AttentionLayer,
2452 ) -> tuple[torch.Tensor, torch.Tensor | None]:
2453 raise NotImplementedError
2269 -
2270 - def forward(
2271 - self,
2272 - layer: AttentionLayer,
2273 - q: torch.Tensor,
2274 - k_c_normed: torch.Tensor, # key in unified attn
2275 - k_pe: torch.Tensor, # value in unified attn
2276 - kv_cache: torch.Tensor,
2277 - attn_metadata: M,
2278 - output: torch.Tensor | None = None,
2279 - output_scale: torch.Tensor | None = None,
2280 - output_block_scale: torch.Tensor | None = None,
2281 - ) -> torch.Tensor:
2282 - assert output is not None, "Output tensor must be provided."
2283 -
2284 - if output_scale is not None or output_block_scale is not None:
2285 - raise NotImplementedError(
2286 - "fused output quantization is not yet supported for MLACommonImpl"
2287 - )
2288 -
2289 - if attn_metadata is None:
2290 - # During the profile run try to simulate to worse case output size
2291 - # for `self.kv_b_proj(kv_c_normed)` in `_compute_prefill_context`
2292 - # since this can be large
2293 - _ = torch.empty(
2294 - (
2295 - self.chunked_prefill_workspace_size,
2296 - self.num_heads,
2297 - self.qk_nope_head_dim + self.v_head_dim,
2298 - ),
2299 - device=k_c_normed.device,
2300 - dtype=k_c_normed.dtype,
2301 - )
2302 -
2303 - # The zero fill is required when used with DP + EP
2304 - # to ensure all ranks within a DP group compute the
2305 - # same expert outputs.
2306 - return output.fill_(0)
2307 -
2308 - if self.dcp_world_size == -1:
2309 - self.dcp_world_size = get_dcp_group().world_size
2310 -
2311 - fp8_attention = self.kv_cache_dtype.startswith("fp8")
2312 -
2313 - num_actual_toks = attn_metadata.num_actual_tokens
2314 -
2315 - # Inputs and outputs may be padded for CUDA graphs
2316 - output_padded = output
2317 - output = output[:num_actual_toks, ...]
2318 - q = q[:num_actual_toks, ...]
2319 - k_c_normed = k_c_normed[:num_actual_toks, ...]
2320 - k_pe = k_pe[:num_actual_toks, ...]
2321 -
2322 - assert (
2323 - attn_metadata.num_decodes is not None
2324 - and attn_metadata.num_prefills is not None
2325 - and attn_metadata.num_decode_tokens is not None
2326 - )
2327 -
2328 - has_decode = attn_metadata.num_decodes > 0
2329 - has_prefill = attn_metadata.num_prefills > 0
2330 - num_decode_tokens = attn_metadata.num_decode_tokens
2331 -
2332 - decode_q = q[:num_decode_tokens]
2333 -
2334 - prefill_q = q[num_decode_tokens:]
2335 - prefill_k_pe = k_pe[num_decode_tokens:]
2336 - prefill_k_c_normed = k_c_normed[num_decode_tokens:]
2337 -
2338 - # write the latent and rope to kv cache
2339 - if kv_cache.numel() > 0:
2340 - ops.concat_and_cache_mla(
2341 - k_c_normed,
2342 - k_pe.squeeze(1),
2343 - kv_cache,
2344 - attn_metadata.slot_mapping.flatten(),
2345 - kv_cache_dtype=self.kv_cache_dtype,
2346 - scale=layer._k_scale,
2347 - )
2348 -
2349 - if fp8_attention:
2350 - kv_cache = kv_cache.view(current_platform.fp8_dtype())
2351 -
2352 - if has_prefill:
2353 - self._forward_prefill(
2354 - prefill_q,
2355 - prefill_k_c_normed,
2356 - prefill_k_pe,
2357 - kv_cache,
2358 - attn_metadata,
2359 - layer._k_scale,
2360 - output=output[num_decode_tokens:],
2361 - )
2362 -
2363 - if has_decode:
2364 - assert attn_metadata.decode is not None
2365 -
2366 - decode_q_nope, decode_q_pe = decode_q.split(
2367 - [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
2368 - )
2369 -
2370 - # Convert from (B, N, P) to (N, B, P)
2371 - decode_q_nope = decode_q_nope.transpose(0, 1)
2372 -
2373 - if self.q_pad_num_heads is not None:
2374 - B, N, L = decode_q_pe.shape
2375 - decode_pe_padded = decode_q_pe.new_empty((B, self.q_pad_num_heads, L))
2376 - decode_pe_padded.resize_((B, N, L))
2377 - decode_pe_padded.copy_(decode_q_pe)
2378 - decode_q_pe = decode_pe_padded
2379 -
2380 - if self.is_aiter_triton_fp4_bmm_enabled:
2381 - from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
2382 -
2383 - decode_ql_nope = batched_gemm_a16wfp4(
2384 - decode_q_nope,
2385 - self.W_K,
2386 - self.W_K_scale,
2387 - transpose_bm=True,
2388 - prequant=True,
2389 - y_scale=layer._q_scale if fp8_attention else None,
2390 - )
2391 - elif self.is_aiter_triton_fp8_bmm_enabled:
2392 - # Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L)
2393 - decode_ql_nope = rocm_aiter_ops.triton_fp8_bmm(
2394 - decode_q_nope,
2395 - self.W_K,
2396 - self.W_K_scale,
2397 - group_size=128,
2398 - transpose_bm=True,
2399 - )
2400 - else:
2401 - # Pads the head_dim if necessary (for the underlying kernel)
2402 - N, B, P = decode_q_nope.shape
2403 - _, _, L = self.W_UK_T.shape
2404 -
2405 - if self.q_pad_num_heads is not None:
2406 - decode_ql_nope = decode_q_nope.new_empty(
2407 - (self.q_pad_num_heads, B, L)
2408 - )
2409 - decode_ql_nope.resize_((N, B, L))
2410 - else:
2411 - decode_ql_nope = decode_q_nope.new_empty((N, B, L))
2412 -
2413 - # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
2414 - torch.bmm(decode_q_nope, self.W_UK_T, out=decode_ql_nope)
2415 -
2416 - # Convert from (N, B, L) to (B, N, L)
2417 - decode_ql_nope = decode_ql_nope.transpose(0, 1)
2418 -
2419 - if fp8_attention:
2420 - assert decode_ql_nope.shape[0] == decode_q_pe.shape[0]
2421 - assert decode_ql_nope.shape[1] == decode_q_pe.shape[1]
2422 - decode_q = self._decode_concat_quant_fp8_op(
2423 - decode_ql_nope, decode_q_pe, layer._q_scale
2424 - )
2425 - else:
2426 - decode_q = (decode_ql_nope, decode_q_pe)
2427 - if self.dcp_world_size > 1:
2428 - assert not fp8_attention, "DCP not support fp8 kvcache now."
2429 - # concatenate decode_ql_nope and decode_q_pe -> (B, N, L + P)
2430 - decode_q = torch.cat(decode_q, dim=-1)
2431 - # decode_q do allgather in head dim.
2432 - decode_q = get_dcp_group().all_gather(decode_q, dim=1)
2433 -
2434 - # call decode attn
2435 - attn_out, lse = self._forward_decode(
2436 - decode_q, kv_cache, attn_metadata, layer
2437 - )
2438 -
2439 - # correct dcp attn_out with lse.
2440 - if self.dcp_world_size > 1:
2441 - attn_out = cp_lse_ag_out_rs(
2442 - attn_out,
2443 - lse,
2444 - get_dcp_group(),
2445 - is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),
2446 - )
2447 -
2448 - # v_up projection
2449 - self._v_up_proj(attn_out, out=output[:num_decode_tokens])
2450 - return output_padded
vllm/v1/attention/backend.py
+82 -13
@@ -67,7 +67,7 @@ class AttentionBackend(ABC):
67
68 @staticmethod
69 @abstractmethod
70 - def get_impl_cls() -> type["AttentionImpl"]:
70 + def get_impl_cls() -> type["AttentionImplBase"]:
71 raise NotImplementedError
72
73 @staticmethod
@@ -594,7 +594,14 @@ class AttentionLayer(Protocol):
594 ) -> torch.Tensor: ...
595
596
597 -class AttentionImpl(ABC, Generic[T]):
597 +class AttentionImplBase(ABC, Generic[T]):
598 + """Base class for attention implementations.
599 +
600 + Contains common attributes and initialization logic shared by both
601 + standard AttentionImpl and MLAAttentionImpl. Does not define a forward
602 + method - subclasses define their own forward interfaces.
603 + """
604 +
605 # Required attributes that all impls should have
606 num_heads: int
607 head_size: int
@@ -662,6 +669,13 @@ class AttentionImpl(ABC, Generic[T]):
669 )
670 return self
671
672 + def process_weights_after_loading(self, act_dtype: torch.dtype):
673 + pass
674 +
675 +
676 +class AttentionImpl(AttentionImplBase[T], Generic[T]):
677 + """Standard attention implementation with forward method."""
678 +
679 @abstractmethod
680 def __init__(
681 self,
@@ -704,11 +718,10 @@ class AttentionImpl(ABC, Generic[T]):
718 """
719 return False
720
707 - def process_weights_after_loading(self, act_dtype: torch.dtype):
708 - pass
721
722 +class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
723 + """MLA attention implementation with forward_mqa and forward_mha methods."""
724
711 -class MLAAttentionImpl(AttentionImpl[T], Generic[T]):
725 @abstractmethod
726 def __init__(
727 self,
@@ -731,22 +744,78 @@ class MLAAttentionImpl(AttentionImpl[T], Generic[T]):
744 v_head_dim: int,
745 kv_b_proj: "ColumnParallelLinear",
746 indexer: object | None = None,
747 + q_pad_num_heads: int | None = None,
748 ) -> None:
749 raise NotImplementedError
750
751 @abstractmethod
738 - def forward(
752 + def forward_mha(
753 self,
740 - layer: AttentionLayer,
741 - hidden_states_or_cq: torch.Tensor,
754 + q: torch.Tensor,
755 kv_c_normed: torch.Tensor,
756 k_pe: torch.Tensor,
744 - kv_cache: torch.Tensor,
757 + kv_c_and_k_pe_cache: torch.Tensor,
758 attn_metadata: T,
746 - output: torch.Tensor | None = None,
747 - output_scale: torch.Tensor | None = None,
748 - output_block_scale: torch.Tensor | None = None,
749 - ) -> torch.Tensor:
759 + k_scale: torch.Tensor,
760 + output: torch.Tensor,
761 + ) -> None:
762 + """MHA-style prefill forward pass."""
763 + raise NotImplementedError
764 +
765 + @abstractmethod
766 + def forward_mqa(
767 + self,
768 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
769 + kv_c_and_k_pe_cache: torch.Tensor,
770 + attn_metadata: T,
771 + layer: AttentionLayer,
772 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
773 + """MQA-style decode forward pass."""
774 + raise NotImplementedError
775 +
776 +
777 +class SparseMLAAttentionImpl(AttentionImplBase[T], Generic[T]):
778 + """Sparse MLA attention implementation with only forward_mqa method.
779 +
780 + Sparse MLA implementations only support decode (MQA-style) attention.
781 + They do not support prefill (MHA-style) attention.
782 + """
783 +
784 + @abstractmethod
785 + def __init__(
786 + self,
787 + num_heads: int,
788 + head_size: int,
789 + scale: float,
790 + num_kv_heads: int,
791 + alibi_slopes: list[float] | None,
792 + sliding_window: int | None,
793 + kv_cache_dtype: str,
794 + logits_soft_cap: float | None,
795 + attn_type: str,
796 + kv_sharing_target_layer_name: str | None,
797 + # MLA Specific Arguments
798 + q_lora_rank: int | None,
799 + kv_lora_rank: int,
800 + qk_nope_head_dim: int,
801 + qk_rope_head_dim: int,
802 + qk_head_dim: int,
803 + v_head_dim: int,
804 + kv_b_proj: "ColumnParallelLinear",
805 + indexer: object | None = None,
806 + q_pad_num_heads: int | None = None,
807 + ) -> None:
808 + raise NotImplementedError
809 +
810 + @abstractmethod
811 + def forward_mqa(
812 + self,
813 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
814 + kv_c_and_k_pe_cache: torch.Tensor,
815 + attn_metadata: T,
816 + layer: AttentionLayer,
817 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
818 + """MQA-style decode forward pass."""
819 raise NotImplementedError
820
821
vllm/v1/attention/backends/mla/cutlass_mla.py
+1 -1
@@ -244,7 +244,7 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
244
245 return out, lse
246
247 - def _forward_decode(
247 + def forward_mqa(
248 self,
249 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
250 kv_c_and_k_pe_cache: torch.Tensor,
vllm/v1/attention/backends/mla/flashattn_mla.py
+1 -1
@@ -293,7 +293,7 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
293 "FlashAttnMLA V1 with FP8 KV cache not yet supported"
294 )
295
296 - def _forward_decode(
296 + def forward_mqa(
297 self,
298 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
299 kv_c_and_k_pe_cache: torch.Tensor,
vllm/v1/attention/backends/mla/flashinfer_mla.py
+1 -1
@@ -150,7 +150,7 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
150 self.bmm1_scale: float | None = None
151 self.bmm2_scale: float | None = None
152
153 - def _forward_decode(
153 + def forward_mqa(
154 self,
155 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
156 kv_c_and_k_pe_cache: torch.Tensor,
vllm/v1/attention/backends/mla/flashmla.py
+1 -1
@@ -234,7 +234,7 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
234 "FlashMLAImpl"
235 )
236
237 - def _forward_decode(
237 + def forward_mqa(
238 self,
239 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
240 kv_c_and_k_pe_cache: torch.Tensor,
vllm/v1/attention/backends/mla/flashmla_sparse.py
+24 -70
@@ -11,7 +11,6 @@ from vllm.config import VllmConfig, get_current_vllm_config
11 from vllm.config.cache import CacheDType
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.attention.mla_attention import (
14 - MLACommonBaseImpl,
14 get_mla_dims,
15 )
16 from vllm.platforms import current_platform
@@ -25,6 +24,7 @@ from vllm.v1.attention.backend import (
24 AttentionMetadataBuilder,
25 CommonAttentionMetadata,
26 MultipleOf,
27 + SparseMLAAttentionImpl,
28 )
29 from vllm.v1.attention.backends.utils import (
30 reshape_attn_output_for_spec_decode,
@@ -686,7 +686,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
686 return metadata
687
688
689 -class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):
689 +class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
690 @staticmethod
691 def _compute_fp8_decode_padded_heads(num_heads: int) -> int:
692 # FP8 decode kernel only supports h_q = 64 or 128
@@ -710,19 +710,12 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):
710 indexer: "Indexer | None" = None,
711 **mla_args,
712 ) -> None:
713 - super().__init__(
714 - num_heads,
715 - head_size,
716 - scale,
717 - num_kv_heads,
718 - alibi_slopes,
719 - sliding_window,
720 - kv_cache_dtype,
721 - logits_soft_cap,
722 - attn_type,
723 - kv_sharing_target_layer_name,
724 - **mla_args,
725 - )
713 + self.num_heads = num_heads
714 + self.head_size = head_size
715 + self.scale = float(scale)
716 + self.num_kv_heads = num_kv_heads
717 + self.kv_cache_dtype = kv_cache_dtype
718 + self.kv_lora_rank: int = mla_args["kv_lora_rank"]
719 self.softmax_scale = scale
720 assert indexer is not None
721 self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
@@ -974,78 +967,39 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):
967 output = output[:, : self.num_heads, :]
968 return output
969
977 - def forward(
970 + def forward_mqa(
971 self,
972 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
973 + kv_c_and_k_pe_cache: torch.Tensor,
974 + attn_metadata: FlashMLASparseMetadata,
975 layer: AttentionLayer,
980 - q: torch.Tensor,
981 - k_c_normed: torch.Tensor, # key in unified attn
982 - k_pe: torch.Tensor, # value in unified attn
983 - kv_cache: torch.Tensor,
984 - attn_metadata: FlashMLASparseMetadata | None,
985 - output: torch.Tensor | None = None,
986 - output_scale: torch.Tensor | None = None,
987 - output_block_scale: torch.Tensor | None = None,
988 - ) -> torch.Tensor:
976 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
977 # NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use
978 # MQA 576/512 approach for both prefill and decode
979
992 - assert output is not None, "Output tensor must be provided."
993 -
994 - if output_scale is not None or output_block_scale is not None:
995 - raise NotImplementedError(
996 - "fused output quantization is not yet supported for MLACommonImpl"
997 - )
998 -
999 - if attn_metadata is None:
1000 - # Dummy run - no need to allocate buffers
1001 - # The zero fill is required when used with DP + EP
1002 - # to ensure all ranks within a DP group compute the
1003 - # same expert outputs.
1004 - return output.fill_(0)
1005 -
1006 - num_actual_toks = attn_metadata.num_actual_tokens
980 + # Concatenate q if it's a tuple (ql_nope, q_pe)
981 + if isinstance(q, tuple):
982 + q = torch.cat(q, dim=-1)
983
1008 - # Inputs and outputs may be padded for CUDA graphs
984 + num_actual_toks = q.shape[0]
985
1010 - q = q[:num_actual_toks, ...]
1011 - k_c_normed = k_c_normed[:num_actual_toks, ...]
1012 - k_pe = k_pe[:num_actual_toks, ...]
986 + # Get topk indices
987 assert self.topk_indices_buffer is not None
988 topk_indices = self.topk_indices_buffer[:num_actual_toks]
989
1016 - q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
1017 - # Convert from (B, N, P) to (N, B, P)
1018 - q_nope = q_nope.transpose(0, 1)
1019 - # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
1020 - ql_nope = torch.bmm(q_nope, self.W_UK_T)
1021 - # Convert from (N, B, L) to (B, N, L)
1022 - ql_nope = ql_nope.transpose(0, 1)
1023 -
990 use_fp8_cache = self.kv_cache_dtype == "fp8_ds_mla"
991
1026 - q = torch.cat([ql_nope, q_pe], dim=-1)
1027 -
1028 - # write the latent and rope to kv cache
1029 - if kv_cache.numel() > 0:
1030 - ops.concat_and_cache_mla(
1031 - k_c_normed,
1032 - k_pe.squeeze(1),
1033 - kv_cache,
1034 - attn_metadata.slot_mapping.flatten(),
1035 - kv_cache_dtype=self.kv_cache_dtype,
1036 - scale=layer._k_scale,
1037 - )
1038 -
992 if not use_fp8_cache:
1040 - attn_out = self._forward_bf16_kv(q, kv_cache, topk_indices, attn_metadata)
993 + attn_out = self._forward_bf16_kv(
994 + q, kv_c_and_k_pe_cache, topk_indices, attn_metadata
995 + )
996 elif attn_metadata.fp8_use_mixed_batch:
997 attn_out = self._forward_fp8_kv_mixed_batch(
1043 - q, kv_cache, topk_indices, attn_metadata
998 + q, kv_c_and_k_pe_cache, topk_indices, attn_metadata
999 )
1000 else:
1001 attn_out = self._forward_fp8_kv_separate_prefill_decode(
1047 - q, kv_cache, topk_indices, attn_metadata
1002 + q, kv_c_and_k_pe_cache, topk_indices, attn_metadata
1003 )
1004
1050 - self._v_up_proj(attn_out, out=output[:num_actual_toks])
1051 - return output
1005 + return attn_out, None
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+1 -1
@@ -241,7 +241,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
241
242 return output
243
244 - def _forward_decode(
244 + def forward_mqa(
245 self,
246 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
247 kv_c_and_k_pe_cache: torch.Tensor,
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+21 -77
@@ -7,12 +7,10 @@ from typing import TYPE_CHECKING, ClassVar
7 import numpy as np
8 import torch
9
10 -from vllm import _custom_ops as ops
10 from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.config import VllmConfig
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.attention.mla_attention import (
15 - MLACommonBaseImpl,
14 get_mla_dims,
15 )
16 from vllm.triton_utils import tl, triton
@@ -23,6 +21,7 @@ from vllm.v1.attention.backend import (
21 AttentionMetadata,
22 AttentionMetadataBuilder,
23 CommonAttentionMetadata,
24 + SparseMLAAttentionImpl,
25 )
26 from vllm.v1.attention.backends.mla.flashmla_sparse import (
27 triton_convert_req_index_to_global_index,
@@ -269,7 +268,7 @@ def reference_mla_sparse_prefill(
268 return (result, lse)
269
270
272 -class ROCMAiterMLASparseImpl(MLACommonBaseImpl[ROCMAiterMLASparseMetadata]):
271 +class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata]):
272 def __init__(
273 self,
274 num_heads: int,
@@ -287,23 +286,15 @@ class ROCMAiterMLASparseImpl(MLACommonBaseImpl[ROCMAiterMLASparseMetadata]):
286 indexer: "Indexer | None" = None,
287 **mla_args,
288 ) -> None:
290 - super().__init__(
291 - num_heads,
292 - head_size,
293 - scale,
294 - num_kv_heads,
295 - alibi_slopes,
296 - sliding_window,
297 - kv_cache_dtype,
298 - logits_soft_cap,
299 - attn_type,
300 - kv_sharing_target_layer_name,
301 - **mla_args,
302 - )
289 + self.num_heads = num_heads
290 + self.head_size = head_size
291 + self.scale = float(scale)
292 + self.num_kv_heads = num_kv_heads
293 + self.kv_cache_dtype = kv_cache_dtype
294 + self.kv_lora_rank: int = mla_args["kv_lora_rank"]
295 self.softmax_scale = scale
296 assert indexer is not None
297 self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
306 - self.is_fp8bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
298
299 def _forward_bf16_kv(
300 self,
@@ -342,56 +333,23 @@ class ROCMAiterMLASparseImpl(MLACommonBaseImpl[ROCMAiterMLASparseMetadata]):
333
334 return output[:, : self.num_heads, :]
335
345 - def forward(
336 + def forward_mqa(
337 self,
347 - layer: AttentionLayer,
348 - q: torch.Tensor,
349 - k_c_normed: torch.Tensor, # key in unified attn
350 - k_pe: torch.Tensor, # value in unified attn
351 - kv_cache: torch.Tensor,
338 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
339 + kv_c_and_k_pe_cache: torch.Tensor,
340 attn_metadata: ROCMAiterMLASparseMetadata,
353 - output: torch.Tensor | None = None,
354 - output_scale: torch.Tensor | None = None,
355 - output_block_scale: torch.Tensor | None = None,
356 - ) -> torch.Tensor:
341 + layer: AttentionLayer,
342 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
343 # NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use
344 # MQA 576/512 approach for both prefill and decode
345
360 - assert output is not None, "Output tensor must be provided."
361 -
362 - if output_scale is not None or output_block_scale is not None:
363 - raise NotImplementedError(
364 - "fused output quantization is not yet supported for ROCMAiterMLASparse"
365 - )
366 -
367 - if attn_metadata is None:
368 - # The zero fill is required when used with DP + EP
369 - # to ensure all ranks within a DP group compute the
370 - # same expert outputs.
371 - return output.fill_(0)
372 -
373 - num_actual_toks = attn_metadata.num_actual_tokens
374 -
375 - # Inputs and outputs may be padded for CUDA graphs
376 -
377 - q = q[:num_actual_toks, ...]
378 - k_c_normed = k_c_normed[:num_actual_toks, ...]
379 - k_pe = k_pe[:num_actual_toks, ...]
380 -
381 - q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
382 - # Convert from (B, N, P) to (N, B, P)
383 - q_nope = q_nope.transpose(0, 1)
384 - if self.is_fp8bmm_enabled:
385 - # Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L)
386 - ql_nope = rocm_aiter_ops.triton_fp8_bmm(
387 - q_nope, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
388 - )
389 - else:
390 - # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
391 - ql_nope = torch.bmm(q_nope, self.W_UK_T)
392 - # Convert from (N, B, L) to (B, N, L)
393 - ql_nope = ql_nope.transpose(0, 1)
346 + # Concatenate q if it's a tuple (ql_nope, q_pe)
347 + if isinstance(q, tuple):
348 + q = torch.cat(q, dim=-1)
349
350 + num_actual_toks = q.shape[0]
351 +
352 + # Get topk indices
353 assert self.topk_indices_buffer is not None
354 topk_indices = self.topk_indices_buffer[:num_actual_toks]
355
@@ -403,22 +361,8 @@ class ROCMAiterMLASparseImpl(MLACommonBaseImpl[ROCMAiterMLASparseMetadata]):
361 NUM_TOPK_TOKENS=attn_metadata.topk_tokens,
362 )
363
406 - q = torch.cat([ql_nope, q_pe], dim=-1)
407 -
408 - # write the latent and rope to kv cache
409 - if kv_cache.numel() > 0:
410 - ops.concat_and_cache_mla(
411 - k_c_normed,
412 - k_pe.squeeze(1),
413 - kv_cache,
414 - attn_metadata.slot_mapping.flatten(),
415 - kv_cache_dtype=self.kv_cache_dtype,
416 - scale=layer._k_scale,
417 - )
418 -
364 attn_out = self._forward_bf16_kv(
420 - q, kv_cache, topk_indices_global, attn_metadata
365 + q, kv_c_and_k_pe_cache, topk_indices_global, attn_metadata
366 )
367
423 - self._v_up_proj(attn_out, out=output[:num_actual_toks])
424 - return output
368 + return attn_out, None
vllm/v1/attention/backends/mla/triton_mla.py
+1 -1
@@ -110,7 +110,7 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
110 **kwargs,
111 )
112
113 - def _forward_decode(
113 + def forward_mqa(
114 self,
115 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
116 kv_c_and_k_pe_cache: torch.Tensor,