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 10, 14:00 UTC
[Chore] Migrate V0 attention utils (#31891)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Jan 7, 2026 at 21:44 UTC
b665bbc2d4277feebf12a6fa38f24d367b5ed25e
10 files changed
+30
-47
tests/kernels/mamba/test_causal_conv1d.py
+1
-1
@@ -7,12 +7,12 @@ import torch
7
import torch.nn.functional as F
8
from einops import rearrange
9
10
-from vllm.attention.backends.utils import PAD_SLOT_ID
10
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
11
causal_conv1d_fn,
12
causal_conv1d_update,
13
)
14
from vllm.utils.torch_utils import set_random_seed
15
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
16
17
18
def causal_conv1d_ref(
tests/kernels/mamba/test_mamba_ssm.py
+1
-1
@@ -8,12 +8,12 @@ from einops import rearrange, repeat
8
9
from tests.kernels.utils import opcheck
10
from vllm import _custom_ops as ops # noqa: F401
11
-from vllm.attention.backends.utils import PAD_SLOT_ID
11
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
12
selective_scan_fn,
13
selective_state_update,
14
)
15
from vllm.utils.torch_utils import set_random_seed
16
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
17
18
19
def selective_state_update_ref(
vllm/attention/backends/utils.py
deleted
-33
@@ -1,33 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Attention backend utils"""
4
-
5
-from dataclasses import dataclass
6
-
7
-from vllm.config import ModelConfig
8
-from vllm.logger import init_logger
9
-
10
-logger = init_logger(__name__)
11
-
12
-PAD_SLOT_ID = -1
13
-
14
-
15
-@dataclass
16
-class MLADims:
17
- q_lora_rank: int | None
18
- kv_lora_rank: int
19
- qk_nope_head_dim: int
20
- qk_rope_head_dim: int
21
- v_head_dim: int
22
-
23
-
24
-def get_mla_dims(model_config: ModelConfig) -> MLADims:
25
- hf_text_config = model_config.hf_text_config
26
-
27
- return MLADims(
28
- q_lora_rank=getattr(hf_text_config, "q_lora_rank", None),
29
- kv_lora_rank=hf_text_config.kv_lora_rank,
30
- qk_nope_head_dim=hf_text_config.qk_nope_head_dim,
31
- qk_rope_head_dim=hf_text_config.qk_rope_head_dim,
32
- v_head_dim=hf_text_config.v_head_dim,
33
- )
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+1
-1
@@ -8,8 +8,8 @@
8
import numpy as np
9
import torch
10
11
-from vllm.attention.backends.utils import PAD_SLOT_ID
11
from vllm.triton_utils import tl, triton
12
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
13
14
15
@triton.jit()
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+1
-1
@@ -8,8 +8,8 @@ import torch
8
from packaging import version
9
10
from vllm import _custom_ops as ops
11
-from vllm.attention.backends.utils import PAD_SLOT_ID
11
from vllm.triton_utils import HAS_TRITON, tl, triton
12
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
13
14
TRITON3 = HAS_TRITON and (version.parse(triton.__version__) >= version.parse("3.0.0"))
15
vllm/v1/attention/backends/gdn_attn.py
+1
-1
@@ -7,9 +7,9 @@ from dataclasses import dataclass
7
import torch
8
9
from vllm.attention.backends.abstract import AttentionBackend
10
-from vllm.attention.backends.utils import PAD_SLOT_ID
10
from vllm.config import VllmConfig
11
from vllm.v1.attention.backends.utils import (
12
+ PAD_SLOT_ID,
13
AttentionCGSupport,
14
AttentionMetadataBuilder,
15
CommonAttentionMetadata,
vllm/v1/attention/backends/mla/common.py
+22
-2
@@ -205,11 +205,10 @@ from vllm.attention.backends.abstract import (
205
AttentionMetadata,
206
MLAAttentionImpl,
207
)
208
-from vllm.attention.backends.utils import get_mla_dims
208
from vllm.attention.ops.common import cp_lse_ag_out_rs
209
from vllm.attention.ops.merge_attn_states import merge_attn_states
210
from vllm.attention.utils.fa_utils import get_flash_attn_version
212
-from vllm.config import VllmConfig, get_current_vllm_config
211
+from vllm.config import ModelConfig, VllmConfig, get_current_vllm_config
212
from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank
213
from vllm.logger import init_logger
214
from vllm.model_executor.layers.batch_invariant import (
@@ -479,6 +478,27 @@ def use_trtllm_ragged_deepseek_prefill() -> bool:
478
)
479
480
481
+@dataclass
482
+class MLADims:
483
+ q_lora_rank: int | None
484
+ kv_lora_rank: int
485
+ qk_nope_head_dim: int
486
+ qk_rope_head_dim: int
487
+ v_head_dim: int
488
+
489
+
490
+def get_mla_dims(model_config: ModelConfig) -> MLADims:
491
+ hf_text_config = model_config.hf_text_config
492
+
493
+ return MLADims(
494
+ q_lora_rank=getattr(hf_text_config, "q_lora_rank", None),
495
+ kv_lora_rank=hf_text_config.kv_lora_rank,
496
+ qk_nope_head_dim=hf_text_config.qk_nope_head_dim,
497
+ qk_rope_head_dim=hf_text_config.qk_rope_head_dim,
498
+ v_head_dim=hf_text_config.v_head_dim,
499
+ )
500
+
501
+
502
class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
503
"""
504
NOTE: Please read the comment at the top of the file before trying to
vllm/v1/attention/backends/mla/flashmla_sparse.py
+1
-2
@@ -13,7 +13,6 @@ from vllm.attention.backends.abstract import (
13
AttentionMetadata,
14
MultipleOf,
15
)
16
-from vllm.attention.backends.utils import get_mla_dims
16
from vllm.attention.ops.flashmla import (
17
flash_mla_sparse_prefill,
18
flash_mla_with_kvcache,
@@ -26,7 +25,7 @@ from vllm.platforms import current_platform
25
from vllm.platforms.interface import DeviceCapability
26
from vllm.triton_utils import tl, triton
27
from vllm.utils.math_utils import cdiv
29
-from vllm.v1.attention.backends.mla.common import MLACommonBaseImpl
28
+from vllm.v1.attention.backends.mla.common import MLACommonBaseImpl, get_mla_dims
29
from vllm.v1.attention.backends.utils import (
30
AttentionCGSupport,
31
AttentionMetadataBuilder,
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+1
-4
@@ -14,12 +14,9 @@ from vllm.attention.backends.abstract import (
14
AttentionLayer,
15
AttentionMetadata,
16
)
17
-from vllm.attention.backends.utils import get_mla_dims
17
from vllm.config import VllmConfig
18
from vllm.logger import init_logger
20
-from vllm.v1.attention.backends.mla.common import (
21
- MLACommonBaseImpl,
22
-)
19
+from vllm.v1.attention.backends.mla.common import MLACommonBaseImpl, get_mla_dims
20
from vllm.v1.attention.backends.mla.flashmla_sparse import (
21
triton_convert_req_index_to_global_index,
22
)
vllm/v1/worker/gpu/block_table.py
+1
-1
@@ -4,9 +4,9 @@ from collections.abc import Iterable
4
5
import torch
6
7
-from vllm.attention.backends.utils import PAD_SLOT_ID
7
from vllm.triton_utils import tl, triton
8
from vllm.utils.math_utils import cdiv
9
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
10
from vllm.v1.utils import CpuGpuBuffer
11
12