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 15, 01:00 UTC
[V0 Deprecation] Deprecate BlockSparse Attention & Phi3-Small (#21217)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jul 19, 2025 at 13:53 UTC
752c6ade2e0f38a26cdaaed6ffae8f72781e2d61
38 files changed
+65
-2435
.buildkite/scripts/hardware_ci/run-amd-test.sh
-1
@@ -108,7 +108,6 @@ fi
108
if [[ $commands == *" kernels/attention"* ]]; then
109
commands="${commands} \
110
--ignore=kernels/attention/test_attention_selector.py \
111
- --ignore=kernels/attention/test_blocksparse_attention.py \
111
--ignore=kernels/attention/test_encoder_decoder_attn.py \
112
--ignore=kernels/attention/test_flash_attn.py \
113
--ignore=kernels/attention/test_flashinfer.py \
docs/models/supported_models.md
-1
@@ -376,7 +376,6 @@ Specified using `--task generate`.
376
| `OrionForCausalLM` | Orion | `OrionStarAI/Orion-14B-Base`, `OrionStarAI/Orion-14B-Chat`, etc. | | ✅︎ | ✅︎ |
377
| `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ | ✅︎ |
378
| `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
379
-| `Phi3SmallForCausalLM` | Phi-3-Small | `microsoft/Phi-3-small-8k-instruct`, `microsoft/Phi-3-small-128k-instruct`, etc. | | ✅︎ | ✅︎ |
379
| `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
380
| `Phi4FlashForCausalLM` | Phi-4-mini-flash-reasoning | `microsoft/microsoft/Phi-4-mini-instruct`, etc. | | | |
381
| `PersimmonForCausalLM` | Persimmon | `adept/persimmon-8b-base`, `adept/persimmon-8b-chat`, etc. | | ✅︎ | ✅︎ |
tests/kernels/attention/test_blocksparse_attention.py
deleted
-441
@@ -1,441 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import random
5
-from typing import Optional
6
-
7
-import pytest
8
-import torch
9
-
10
-from tests.kernels.allclose_default import get_default_atol, get_default_rtol
11
-from vllm import _custom_ops as ops
12
-from vllm.attention.ops.blocksparse_attention.interface import (
13
- LocalStridedBlockSparseAttn)
14
-from vllm.platforms import current_platform
15
-from vllm.utils import get_max_shared_memory_bytes
16
-
17
-FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
18
-# This will change depending on the compute capability.
19
-# - 512 as a buffer
20
-MAX_SEQ_LEN = get_max_shared_memory_bytes() // FLOAT32_BYTES - 512
21
-# MAX_SEQ_LEN = 2771
22
-
23
-# There may not be enough gpu memory due to large NUM_BLOCKS.
24
-# Reduce NUM_BLOCKS when it happens.
25
-NUM_BLOCKS = 4321 # Arbitrary values for testing
26
-PARTITION_SIZE = 512
27
-DTYPES = [torch.half, torch.bfloat16]
28
-NUM_GEN_SEQS = [3] # Arbitrary values for testing
29
-NUM_PREFILL_SEQS = [3] # Arbitrary values for testing
30
-NUM_HEADS = [(40, 40)] # Arbitrary values for testing
31
-
32
-HEAD_SIZES = [64, 112]
33
-BLOCK_SIZES = [16]
34
-USE_ALIBI = [False, True]
35
-KV_CACHE_DTYPE = ["auto", "fp8"]
36
-SEEDS = [0]
37
-CUDA_DEVICES = ['cuda:0']
38
-BLOCKSPARSE_LOCAL_BLOCKS = [16]
39
-BLOCKSPARSE_VERT_STRIDES = [8]
40
-
41
-BLOCKSPARSE_BLOCK_SIZES = [64]
42
-BLOCKSPARSE_HEADS_SLIDINGS = [2, -1]
43
-BLOCKSPARSE_HOMO_HEADS = [True, False]
44
-
45
-
46
-def ref_masked_attention(
47
- query: torch.Tensor,
48
- key: torch.Tensor,
49
- value: torch.Tensor,
50
- scale: float,
51
- attn_mask: Optional[torch.Tensor] = None,
52
-) -> torch.Tensor:
53
- attn_weights = scale * torch.einsum("qhd,khd->hqk", query, key).float()
54
- if attn_mask is not None:
55
- attn_weights = attn_weights + attn_mask.float()
56
- attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)
57
- out = torch.einsum("hqk,khd->qhd", attn_weights, value)
58
- return out
59
-
60
-
61
-def ref_single_query_cached_kv_attention(
62
- output: torch.Tensor,
63
- query: torch.Tensor,
64
- num_queries_per_kv: int,
65
- key_cache: torch.Tensor,
66
- value_cache: torch.Tensor,
67
- block_tables: torch.Tensor,
68
- seq_lens: torch.Tensor,
69
- scale: float,
70
- alibi_slopes: Optional[torch.Tensor],
71
- tp_rank: int = 0,
72
- blocksparse_local_blocks: int = 0,
73
- blocksparse_vert_stride: int = 1,
74
- blocksparse_block_size: int = 64,
75
- blocksparse_head_sliding_step: int = 0,
76
-) -> None:
77
- num_query_heads = query.shape[1]
78
- num_kv_heads = value_cache.shape[1]
79
- head_size = value_cache.shape[2]
80
- block_size = value_cache.shape[3]
81
- num_seqs = query.shape[0]
82
-
83
- block_tables_lst = block_tables.cpu().tolist()
84
- seq_lens_lst = seq_lens.cpu().tolist()
85
- for i in range(num_seqs):
86
- q = query[i].unsqueeze(0)
87
- block_table = block_tables_lst[i]
88
- seq_len = int(seq_lens_lst[i])
89
-
90
- keys_lst: list[torch.Tensor] = []
91
- values_lst: list[torch.Tensor] = []
92
- for j in range(seq_len):
93
- block_number = int(block_table[j // block_size])
94
- block_offset = j % block_size
95
-
96
- k = key_cache[block_number, :, :, block_offset, :]
97
- k = k.reshape(num_kv_heads, head_size)
98
- keys_lst.append(k)
99
-
100
- v = value_cache[block_number, :, :, block_offset]
101
- values_lst.append(v)
102
- keys = torch.stack(keys_lst, dim=0)
103
- values = torch.stack(values_lst, dim=0)
104
- if num_queries_per_kv > 1:
105
- # Handle MQA and GQA
106
- keys = torch.repeat_interleave(keys, num_queries_per_kv, dim=1)
107
- values = torch.repeat_interleave(values, num_queries_per_kv, dim=1)
108
-
109
- alibi_bias = None
110
- if alibi_slopes is not None:
111
- # Create the ALiBi bias used in the paged attention kernel.
112
- position_ids = torch.arange(seq_len).int()
113
- alibi_bias = (position_ids - seq_len + 1).float()
114
- alibi_bias = alibi_slopes.view(-1, 1, 1) * alibi_bias.view(
115
- 1, 1, -1)
116
-
117
- if blocksparse_vert_stride >= 1:
118
- bsize = blocksparse_block_size
119
- hsliding = blocksparse_head_sliding_step
120
- vert = blocksparse_vert_stride
121
- locals = blocksparse_local_blocks
122
- qb = (seq_len - 1) // bsize
123
- attn_mask = q.new_zeros(
124
- (num_query_heads, 1, seq_len)).float() - torch.inf
125
- for h in range(num_query_heads):
126
- if hsliding >= 0: # slide with q heads
127
- bs_offset = (tp_rank * num_query_heads + h) * hsliding + 1
128
- else: # slide with kv heads
129
- bs_offset = (tp_rank * num_kv_heads +
130
- h // num_queries_per_kv) * (-hsliding) + 1
131
- for kb in range(qb + 1):
132
- kj = kb * bsize
133
- if (qb - kb) < locals or \
134
- (kb + bs_offset) % vert == 0:
135
- attn_mask[h, 0, kj:min(kj + bsize, seq_len)] = 0
136
- if alibi_bias is not None:
137
- attn_mask += alibi_bias
138
- else:
139
- attn_mask = alibi_bias
140
-
141
- out = ref_masked_attention(q, keys, values, scale, attn_mask=attn_mask)
142
- out = out.view(num_query_heads, head_size)
143
- output[i].copy_(out, non_blocking=True)
144
-
145
-
146
-@pytest.mark.parametrize("version", ["v1", "v2"])
147
-@pytest.mark.parametrize("num_seqs", NUM_GEN_SEQS)
148
-@pytest.mark.parametrize("num_heads", NUM_HEADS)
149
-@pytest.mark.parametrize("head_size", HEAD_SIZES)
150
-@pytest.mark.parametrize("use_alibi", USE_ALIBI)
151
-@pytest.mark.parametrize("block_size", BLOCK_SIZES)
152
-@pytest.mark.parametrize("dtype", DTYPES)
153
-@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
154
-@pytest.mark.parametrize("seed", SEEDS)
155
-@pytest.mark.parametrize("device", CUDA_DEVICES)
156
-@pytest.mark.parametrize("blocksparse_local_blocks", BLOCKSPARSE_LOCAL_BLOCKS)
157
-@pytest.mark.parametrize("blocksparse_vert_stride", BLOCKSPARSE_VERT_STRIDES)
158
-@pytest.mark.parametrize("blocksparse_block_size", BLOCKSPARSE_BLOCK_SIZES)
159
-@pytest.mark.parametrize("blocksparse_head_sliding_step",
160
- BLOCKSPARSE_HEADS_SLIDINGS)
161
-def test_paged_attention(
162
- kv_cache_factory,
163
- version: str,
164
- num_seqs: int,
165
- num_heads: tuple[int, int],
166
- head_size: int,
167
- use_alibi: bool,
168
- block_size: int,
169
- dtype: torch.dtype,
170
- kv_cache_dtype: str,
171
- seed: int,
172
- device: str,
173
- blocksparse_local_blocks: int,
174
- blocksparse_vert_stride: int,
175
- blocksparse_block_size: int,
176
- blocksparse_head_sliding_step: int,
177
-) -> None:
178
- current_platform.seed_everything(seed)
179
- torch.set_default_device(device)
180
- scale = float(1.0 / (head_size**0.5))
181
- num_query_heads, num_kv_heads = num_heads
182
- query = torch.empty(num_seqs, num_query_heads, head_size, dtype=dtype)
183
- query.uniform_(-scale, scale)
184
-
185
- assert num_query_heads % num_kv_heads == 0
186
- num_queries_per_kv = num_query_heads // num_kv_heads
187
- alibi_slopes = None
188
- if use_alibi:
189
- alibi_slopes = torch.rand(num_query_heads, dtype=torch.float)
190
-
191
- seq_lens = [random.randint(1, MAX_SEQ_LEN) for _ in range(num_seqs)]
192
- seq_lens[-1] = MAX_SEQ_LEN
193
- max_seq_len = max(seq_lens)
194
- seq_lens = torch.tensor(seq_lens, dtype=torch.int)
195
-
196
- # Create the block tables.
197
- max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
198
- block_tables = []
199
- for _ in range(num_seqs):
200
- block_table = [
201
- random.randint(0, NUM_BLOCKS - 1)
202
- for _ in range(max_num_blocks_per_seq)
203
- ]
204
- block_tables.append(block_table)
205
- block_tables = torch.tensor(block_tables, dtype=torch.int)
206
-
207
- # Create the KV caches.
208
- key_caches, value_caches = kv_cache_factory(NUM_BLOCKS, block_size, 1,
209
- num_kv_heads, head_size,
210
- kv_cache_dtype, dtype, seed,
211
- device)
212
- key_cache, value_cache = key_caches[0], value_caches[0]
213
-
214
- # Using default kv_scale
215
- k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
216
- tp_rank = 0
217
-
218
- # Call the paged attention kernel.
219
- output = torch.empty_like(query)
220
- if version == "v1":
221
- ops.paged_attention_v1(
222
- output,
223
- query,
224
- key_cache,
225
- value_cache,
226
- num_kv_heads,
227
- scale,
228
- block_tables,
229
- seq_lens,
230
- block_size,
231
- max_seq_len,
232
- alibi_slopes,
233
- kv_cache_dtype,
234
- k_scale,
235
- v_scale,
236
- tp_rank=tp_rank,
237
- blocksparse_local_blocks=blocksparse_local_blocks,
238
- blocksparse_vert_stride=blocksparse_vert_stride,
239
- blocksparse_block_size=blocksparse_block_size,
240
- blocksparse_head_sliding_step=blocksparse_head_sliding_step,
241
- )
242
- elif version == "v2":
243
- num_partitions = ((max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE)
244
- assert PARTITION_SIZE % block_size == 0
245
- num_seqs, num_heads, head_size = output.shape
246
- tmp_output = torch.empty(
247
- size=(num_seqs, num_heads, num_partitions, head_size),
248
- dtype=output.dtype,
249
- )
250
- exp_sums = torch.empty(
251
- size=(num_seqs, num_heads, num_partitions),
252
- dtype=torch.float32,
253
- )
254
- max_logits = torch.empty_like(exp_sums)
255
- ops.paged_attention_v2(
256
- output,
257
- exp_sums,
258
- max_logits,
259
- tmp_output,
260
- query,
261
- key_cache,
262
- value_cache,
263
- num_kv_heads,
264
- scale,
265
- block_tables,
266
- seq_lens,
267
- block_size,
268
- max_seq_len,
269
- alibi_slopes,
270
- kv_cache_dtype,
271
- k_scale,
272
- v_scale,
273
- tp_rank=tp_rank,
274
- blocksparse_local_blocks=blocksparse_local_blocks,
275
- blocksparse_vert_stride=blocksparse_vert_stride,
276
- blocksparse_block_size=blocksparse_block_size,
277
- blocksparse_head_sliding_step=blocksparse_head_sliding_step,
278
- )
279
- else:
280
- raise AssertionError(f"Unknown version: {version}")
281
-
282
- # Run the reference implementation.
283
- if kv_cache_dtype == "fp8":
284
- # Convert cache data back to dtype.
285
- x = 16 // torch.tensor([], dtype=dtype).element_size()
286
- key_cache_shape = (NUM_BLOCKS, num_kv_heads, head_size // x,
287
- block_size, x)
288
- dequantized_key_cache = torch.empty(size=key_cache_shape,
289
- dtype=dtype,
290
- device=device)
291
- ops.convert_fp8(dequantized_key_cache, key_cache)
292
- key_cache = dequantized_key_cache
293
-
294
- value_cache_shape = value_cache.shape
295
- dequantized_value_cache = torch.empty(size=value_cache_shape,
296
- dtype=dtype,
297
- device=device)
298
- ops.convert_fp8(dequantized_value_cache, value_cache)
299
- value_cache = dequantized_value_cache
300
-
301
- ref_output = torch.empty_like(query)
302
- ref_single_query_cached_kv_attention(
303
- ref_output,
304
- query,
305
- num_queries_per_kv,
306
- key_cache,
307
- value_cache,
308
- block_tables,
309
- seq_lens,
310
- scale,
311
- alibi_slopes,
312
- tp_rank,
313
- blocksparse_local_blocks,
314
- blocksparse_vert_stride,
315
- blocksparse_block_size,
316
- blocksparse_head_sliding_step,
317
- )
318
-
319
- # NOTE(woosuk): Due to the kernel-level differences in the two
320
- # implementations, there is a small numerical difference in the two
321
- # outputs. Thus, we use a relaxed tolerance for the test.
322
- atol = get_default_atol(output) if current_platform.is_rocm() else 1e-3
323
- rtol = get_default_rtol(output) if current_platform.is_rocm() else 1e-5
324
-
325
- # NOTE(zhaoyang): FP8 KV Cache will introduce quantization error,
326
- # so we use a relaxed tolerance for the test.
327
- atol, rtol = 1e-3, 1e-5
328
- if kv_cache_dtype == "fp8":
329
- atol, rtol = 1e-2, 1e-5
330
- torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
331
-
332
-
333
-def ref_multi_query_kv_attention(
334
- cu_seq_lens: list[int],
335
- query: torch.Tensor,
336
- key: torch.Tensor,
337
- value: torch.Tensor,
338
- scale: float,
339
- dtype: torch.dtype,
340
-) -> torch.Tensor:
341
- num_seqs = len(cu_seq_lens) - 1
342
- ref_outputs = []
343
- for i in range(num_seqs):
344
- start_idx = cu_seq_lens[i]
345
- end_idx = cu_seq_lens[i + 1]
346
- seq_len = end_idx - start_idx
347
-
348
- # Create attention mask.
349
- attn_mask = torch.triu(torch.ones(seq_len, seq_len, dtype=dtype),
350
- diagonal=1)
351
- attn_mask = attn_mask * torch.finfo(dtype).min
352
- attn_mask = attn_mask.to(dtype=dtype)
353
-
354
- ref_output = ref_masked_attention(
355
- query[start_idx:end_idx],
356
- key[start_idx:end_idx],
357
- value[start_idx:end_idx],
358
- scale,
359
- attn_mask=attn_mask,
360
- )
361
- ref_outputs.append(ref_output)
362
- ref_output = torch.cat(ref_outputs, dim=0)
363
- return ref_output
364
-
365
-
366
-@pytest.mark.parametrize("num_seqs", NUM_PREFILL_SEQS)
367
-@pytest.mark.parametrize("num_heads", NUM_HEADS)
368
-@pytest.mark.parametrize("head_size", HEAD_SIZES)
369
-@pytest.mark.parametrize("blocksparse_local_blocks", BLOCKSPARSE_LOCAL_BLOCKS)
370
-@pytest.mark.parametrize("blocksparse_vert_stride", BLOCKSPARSE_VERT_STRIDES)
371
-@pytest.mark.parametrize("blocksparse_block_size", BLOCKSPARSE_BLOCK_SIZES)
372
-@pytest.mark.parametrize("blocksparse_homo_heads", BLOCKSPARSE_HOMO_HEADS)
373
-@pytest.mark.parametrize("dtype", DTYPES)
374
-@pytest.mark.parametrize("seed", SEEDS)
375
-@pytest.mark.parametrize("device", CUDA_DEVICES)
376
-@torch.inference_mode()
377
-def test_varlen_blocksparse_attention_prefill(
378
- num_seqs: int,
379
- num_heads: tuple[int, int],
380
- head_size: int,
381
- blocksparse_local_blocks: int,
382
- blocksparse_vert_stride: int,
383
- blocksparse_block_size: int,
384
- blocksparse_homo_heads: bool,
385
- dtype: torch.dtype,
386
- seed: int,
387
- device: str,
388
-) -> None:
389
- current_platform.seed_everything(seed)
390
- torch.set_default_device(device)
391
- # MAX_SEQ_LEN sometimes causes OOM in the reference implementation.
392
- # As the xformers library is already tested with its own tests, we can use
393
- # a smaller MAX_SEQ_LEN here.
394
- max_len = min(MAX_SEQ_LEN, 4096)
395
- seq_lens = random.sample(range(1, max_len), num_seqs)
396
- cu_seq_lens = torch.cumsum(torch.tensor([0] + seq_lens), dim=0)
397
- num_tokens = sum(seq_lens)
398
-
399
- scale = float(1.0 / (head_size**0.5))
400
- num_query_heads, num_kv_heads = num_heads
401
- assert num_query_heads % num_kv_heads == 0
402
- num_queries_per_kv = num_query_heads // num_kv_heads
403
-
404
- qkv = torch.empty(num_tokens,
405
- num_query_heads + 2 * num_kv_heads,
406
- head_size,
407
- dtype=dtype)
408
- qkv.uniform_(-scale, scale)
409
- query, key, value = qkv.split(
410
- [num_query_heads, num_kv_heads, num_kv_heads], dim=1)
411
-
412
- bs_attn_op = LocalStridedBlockSparseAttn(
413
- num_query_heads,
414
- max_len,
415
- local_blocks=blocksparse_local_blocks,
416
- vert_stride=blocksparse_vert_stride,
417
- block_size=blocksparse_block_size,
418
- device=device,
419
- dtype=dtype,
420
- homo_head=blocksparse_homo_heads)
421
-
422
- output = bs_attn_op(query,
423
- key,
424
- value,
425
- cu_seq_lens.to(device),
426
- sm_scale=scale)
427
-
428
- if num_queries_per_kv > 1:
429
- # Handle MQA and GQA
430
- key = torch.repeat_interleave(key, num_queries_per_kv, dim=1)
431
- value = torch.repeat_interleave(value, num_queries_per_kv, dim=1)
432
-
433
- ref_output = ref_multi_query_kv_attention(
434
- cu_seq_lens.tolist(),
435
- query,
436
- key,
437
- value,
438
- scale,
439
- dtype,
440
- )
441
- torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
tests/kernels/attention/test_rocm_attention_selector.py
+24
-8
@@ -33,8 +33,12 @@ def test_selector(monkeypatch: pytest.MonkeyPatch):
33
34
# change the attention backend to triton MLA
35
m.setenv(STR_BACKEND_ENV_VAR, "TRITON_MLA")
36
- backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False,
37
- False, True)
36
+ backend = get_attn_backend(576,
37
+ torch.bfloat16,
38
+ "auto",
39
+ 16,
40
+ False,
41
+ use_mla=True)
42
assert (backend.get_name() == "TRITON_MLA"
43
or backend.get_name() == "TRITON_MLA_VLLM_V1")
44
@@ -42,15 +46,23 @@ def test_selector(monkeypatch: pytest.MonkeyPatch):
46
# If use_mla is true
47
# The selected backend is triton MLA
48
m.setenv(STR_BACKEND_ENV_VAR, None)
45
- backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False,
46
- False, True)
49
+ backend = get_attn_backend(576,
50
+ torch.bfloat16,
51
+ "auto",
52
+ 16,
53
+ False,
54
+ use_mla=True)
55
assert (backend.get_name() == "TRITON_MLA"
56
or backend.get_name() == "TRITON_MLA_VLLM_V1")
57
58
# change the attention backend to AITER MLA
59
m.setenv(STR_BACKEND_ENV_VAR, "ROCM_AITER_MLA")
52
- backend = get_attn_backend(576, torch.bfloat16, "auto", 1, False,
53
- False, True)
60
+ backend = get_attn_backend(576,
61
+ torch.bfloat16,
62
+ "auto",
63
+ 1,
64
+ False,
65
+ use_mla=True)
66
assert (backend.get_name() == "ROCM_AITER_MLA"
67
or backend.get_name() == "ROCM_AITER_MLA_VLLM_V1")
68
@@ -60,7 +72,11 @@ def test_selector(monkeypatch: pytest.MonkeyPatch):
72
# The selected backend is ROCM_AITER_MLA
73
m.setenv(STR_BACKEND_ENV_VAR, None)
74
m.setenv("VLLM_ROCM_USE_AITER", "1")
63
- backend = get_attn_backend(576, torch.bfloat16, "auto", 1, False,
64
- False, True)
75
+ backend = get_attn_backend(576,
76
+ torch.bfloat16,
77
+ "auto",
78
+ 1,
79
+ False,
80
+ use_mla=True)
81
assert (backend.get_name() == "ROCM_AITER_MLA"
82
or backend.get_name() == "ROCM_AITER_MLA_VLLM_V1")
tests/models/registry.py
-4
@@ -247,10 +247,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
247
"PersimmonForCausalLM": _HfExamplesInfo("adept/persimmon-8b-chat"),
248
"PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2"),
249
"Phi3ForCausalLM": _HfExamplesInfo("microsoft/Phi-3-mini-4k-instruct"),
250
- # Blocksparse attention not supported in V1 yet
251
- "Phi3SmallForCausalLM": _HfExamplesInfo("microsoft/Phi-3-small-8k-instruct",
252
- trust_remote_code=True,
253
- v0_only=True),
250
"Phi4FlashForCausalLM": _HfExamplesInfo("microsoft/Phi-4-mini-flash-reasoning", # noqa: E501
251
trust_remote_code=True,
252
v0_only=True,
vllm/attention/backends/abstract.py
-1
@@ -269,7 +269,6 @@ class AttentionImpl(ABC, Generic[T]):
269
alibi_slopes: Optional[List[float]] = None,
270
sliding_window: Optional[int] = None,
271
kv_cache_dtype: str = "auto",
272
- blocksparse_params: Optional[Dict[str, Any]] = None,
272
logits_soft_cap: Optional[float] = None,
273
attn_type: str = AttentionType.DECODER,
274
kv_sharing_target_layer_name: Optional[str] = None,
vllm/attention/backends/blocksparse_attn.py
deleted
-466
@@ -1,466 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from dataclasses import dataclass, field
5
-from typing import Any, Dict, List, Optional, Tuple, Type
6
-
7
-import torch
8
-
9
-from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
10
- AttentionLayer,
11
- AttentionMetadata, AttentionType)
12
-from vllm.attention.backends.utils import (CommonAttentionState,
13
- CommonMetadataBuilder)
14
-from vllm.attention.ops.blocksparse_attention.interface import (
15
- LocalStridedBlockSparseAttn, get_head_sliding_step)
16
-from vllm.attention.ops.paged_attn import PagedAttention
17
-from vllm.distributed import (get_tensor_model_parallel_rank,
18
- get_tensor_model_parallel_world_size)
19
-
20
-
21
-@dataclass
22
-class BlocksparseParams:
23
- max_seqlen: int
24
-
25
- # Num q heads per tensor-parallel rank/partition
26
- num_heads: int # per TP partition
27
- # Num kv heads per tensor-parallel rank/partition
28
- num_kv_heads: int
29
-
30
- # block size used for blocksparse attention.
31
- # This is the block_size used in `local_blocks`, `vert_stride`.
32
- block_size: int
33
-
34
- # Number of blocks for local attention, i.e., number of
35
- # local attended tokens / `sparse_block_size`
36
- local_blocks: int
37
-
38
- # Attend to one block per every `vert_stride` blocks.
39
- # Controlling the sparsity
40
- vert_stride: int
41
- """
42
- If to use the same vertical stride offset for all heads,
43
- i.e., attend to the same block of tokens on all heads.
44
- By default, it is False, i.e., attention on the non-local
45
- blocks depends on the `head_idx`, that is on
46
- blocks satisfying
47
- `(block_idx + head_idx * head_sliding_step + 1) % vert_stride == 0`
48
- where `head_sliding_step=max(1, int(vert_stride / num_total_heads))`,
49
- `block_idx = position_id // sparse_block_size`.
50
- See `..ops.blocksparse_attention.utils:get_sparse_attn_mask`
51
- for more detail.
52
- """
53
- homo_head: bool = False
54
-
55
- # If within a group, the kv offsets that each q attends is the same or no.
56
- homo_head_group: bool = False
57
-
58
- # Decided by homo_head and homo_head group
59
- head_sliding_step: int = field(init=False)
60
-
61
- # range of q heads to for a TP rank
62
- active_head_range: Tuple = field(init=False)
63
-
64
- def __post_init__(self):
65
- assert self.block_size > 0
66
- assert self.local_blocks >= 0
67
- assert self.vert_stride >= 1
68
-
69
- tp_size = get_tensor_model_parallel_world_size()
70
- tp_rank = get_tensor_model_parallel_rank()
71
- total_heads = tp_size * self.num_heads
72
- total_kv_heads = tp_size * self.num_kv_heads
73
-
74
- if self.homo_head:
75
- self.head_sliding_step = 0
76
- elif self.homo_head_group:
77
- head_sliding_step = get_head_sliding_step(total_kv_heads,
78
- self.vert_stride)
79
- # negative indicates sliding along kv heads, i.e., homo q group
80
- self.head_sliding_step = -head_sliding_step
81
- else:
82
- self.head_sliding_step = get_head_sliding_step(
83
- total_heads, self.vert_stride)
84
-
85
- self.active_head_range = (
86
- tp_rank * self.num_heads,
87
- (tp_rank + 1) * self.num_heads,
88
- )
89
-
90
-
91
-class BlocksparseFlashAttentionBackend(AttentionBackend):
92
-
93
- @staticmethod
94
- def get_name() -> str:
95
- return "BLOCK_SPARSE_FLASH_ATTN"
96
-
97
- @staticmethod
98
- def get_impl_cls() -> Type["BlocksparseFlashAttentionImpl"]:
99
- return BlocksparseFlashAttentionImpl
100
-
101
- @staticmethod
102
- def get_metadata_cls() -> Type["AttentionMetadata"]:
103
- return BlocksparseFlashAttentionMetadata
104
-
105
- @staticmethod
106
- def get_builder_cls() -> Type["BlocksparseFlashAttentionMetadataBuilder"]:
107
- return BlocksparseFlashAttentionMetadataBuilder
108
-
109
- @staticmethod
110
- def get_state_cls() -> Type["CommonAttentionState"]:
111
- return CommonAttentionState
112
-
113
- @staticmethod
114
- def get_kv_cache_shape(
115
- num_blocks: int,
116
- block_size: int,
117
- num_kv_heads: int,
118
- head_size: int,
119
- ) -> Tuple[int, ...]:
120
- return PagedAttention.get_kv_cache_shape(num_blocks, block_size,
121
- num_kv_heads, head_size)
122
-
123
- @staticmethod
124
- def swap_blocks(
125
- src_kv_cache: torch.Tensor,
126
- dst_kv_cache: torch.Tensor,
127
- src_to_dst: Dict[int, int],
128
- ) -> None:
129
- PagedAttention.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
130
-
131
- @staticmethod
132
- def copy_blocks(
133
- kv_caches: List[torch.Tensor],
134
- src_to_dists: Dict[int, List[int]],
135
- ) -> None:
136
- PagedAttention.copy_blocks(kv_caches, src_to_dists)
137
-
138
-
139
-@dataclass
140
-class BlocksparseFlashAttentionMetadata(AttentionMetadata):
141
- """A copy of Metadata for FlashAttentionBackend,
142
- to avoid having to install flash_attn.
143
-
144
- NOTE: Any python object stored here is not updated when it is
145
- cuda-graph replayed. If you have values that need to be changed
146
- dynamically, it should be stored in tensor. The tensor has to be
147
- updated from `CUDAGraphRunner.forward` API.
148
- """
149
- # (batch_size,). The sequence length per sequence. Sequence length means
150
- # the computed tokens + new tokens None if it is a decoding.
151
- seq_lens: Optional[List[int]]
152
- # seq_lens stored as a tensor.
153
- seq_lens_tensor: Optional[torch.Tensor]
154
-
155
- # NOTE(sang): Definition of context_len, query_len, and seq_len.
156
- # |---------- N-1 iteration --------|
157
- # |---------------- N iteration ---------------------|
158
- # |- tokenA -|......................|-- newTokens ---|
159
- # |---------- context_len ----------|
160
- # |-------------------- seq_len ----------------------|
161
- # |-- query_len ---|
162
-
163
- # Maximum query length in the batch. None for decoding.
164
- max_query_len: Optional[int]
165
- # Maximum sequence length among prefill batch. 0 if there are decoding
166
- # requests only.
167
- max_prefill_seq_len: int
168
- # Maximum sequence length among decode batch. 0 if there are prefill
169
- # requests only.
170
- max_decode_seq_len: int
171
- # (batch_size + 1,). The cumulative subquery lengths of the sequences in
172
- # the batch, used to index into subquery. E.g., if the subquery length
173
- # is [4, 6], it is [0, 4, 10].
174
- query_start_loc: Optional[torch.Tensor]
175
- # (batch_size + 1,). The cumulative sequence lengths of the sequences in
176
- # the batch, used to index into sequence. E.g., if the sequence length is
177
- # [4, 6], it is [0, 4, 10].
178
- seq_start_loc: Optional[torch.Tensor]
179
- # (batch_size,) A tensor of context lengths (tokens that are computed
180
- # so far).
181
- context_lens_tensor: Optional[torch.Tensor]
182
-
183
- # (batch_size, max_blocks_per_seq).
184
- # Block addresses per sequence. (Seq id -> list of physical block)
185
- # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks
186
- # in the kv cache. Each block can contain up to block_size tokens.
187
- # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph
188
- # captured.
189
- block_tables: Optional[torch.Tensor]
190
-
191
- # Whether or not if cuda graph is enabled.
192
- # Cuda-graph is currently enabled for decoding only.
193
- # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.
194
- use_cuda_graph: bool
195
-
196
- # Max number of query tokens for among request in the batch.
197
- max_decode_query_len: Optional[int] = None
198
-
199
- _cached_prefill_metadata: Optional[
200
- "BlocksparseFlashAttentionMetadata"] = None
201
- _cached_decode_metadata: Optional[
202
- "BlocksparseFlashAttentionMetadata"] = None
203
-
204
- @property
205
- def prefill_metadata(
206
- self) -> Optional["BlocksparseFlashAttentionMetadata"]:
207
- if self.num_prefills == 0:
208
- return None
209
-
210
- if self._cached_prefill_metadata is not None:
211
- return self._cached_prefill_metadata
212
-
213
- assert self.seq_lens is not None
214
- assert self.seq_lens_tensor is not None
215
- assert self.query_start_loc is not None
216
- assert self.context_lens_tensor is not None
217
- assert self.block_tables is not None
218
- assert self.seq_start_loc is not None
219
-
220
- self._cached_prefill_metadata = BlocksparseFlashAttentionMetadata(
221
- num_prefills=self.num_prefills,
222
- num_prefill_tokens=self.num_prefill_tokens,
223
- num_decode_tokens=0,
224
- slot_mapping=self.slot_mapping[:self.num_prefill_tokens],
225
- multi_modal_placeholder_index_maps=self.
226
- multi_modal_placeholder_index_maps,
227
- enable_kv_scales_calculation=self.enable_kv_scales_calculation,
228
- seq_lens=self.seq_lens[:self.num_prefills],
229
- seq_lens_tensor=self.seq_lens_tensor[:self.num_prefills],
230
- max_query_len=self.max_query_len,
231
- max_prefill_seq_len=self.max_prefill_seq_len,
232
- max_decode_seq_len=0,
233
- query_start_loc=self.query_start_loc[:self.num_prefills + 1],
234
- seq_start_loc=self.seq_start_loc[:self.num_prefills + 1],
235
- context_lens_tensor=self.context_lens_tensor[:self.num_prefills],
236
- block_tables=self.block_tables[:self.num_prefills],
237
- use_cuda_graph=False,
238
- )
239
- return self._cached_prefill_metadata
240
-
241
- @property
242
- def decode_metadata(self) -> Optional["BlocksparseFlashAttentionMetadata"]:
243
- if self.num_decode_tokens == 0:
244
- return None
245
-
246
- if self._cached_decode_metadata is not None:
247
- return self._cached_decode_metadata
248
- assert self.block_tables is not None
249
- assert self.seq_lens_tensor is not None
250
-
251
- self._cached_decode_metadata = BlocksparseFlashAttentionMetadata(
252
- num_prefills=0,
253
- num_prefill_tokens=0,
254
- num_decode_tokens=self.num_decode_tokens,
255
- slot_mapping=self.slot_mapping[self.num_prefill_tokens:],
256
- multi_modal_placeholder_index_maps=None,
257
- enable_kv_scales_calculation=False,
258
- seq_lens=None,
259
- seq_lens_tensor=self.seq_lens_tensor[self.num_prefills:],
260
- max_query_len=None,
261
- max_prefill_seq_len=0,
262
- max_decode_seq_len=self.max_decode_seq_len,
263
- query_start_loc=None,
264
- seq_start_loc=None,
265
- context_lens_tensor=None,
266
- block_tables=self.block_tables[self.num_prefills:],
267
- use_cuda_graph=self.use_cuda_graph,
268
- )
269
- return self._cached_decode_metadata
270
-
271
-
272
-class BlocksparseFlashAttentionMetadataBuilder(
273
- CommonMetadataBuilder[BlocksparseFlashAttentionMetadata]):
274
-
275
- _metadata_cls = BlocksparseFlashAttentionMetadata
276
-
277
-
278
-class BlocksparseFlashAttentionImpl(AttentionImpl):
279
- """
280
- If the input tensors contain prompt tokens, the layout is as follows:
281
- |<--------------- num_prompt_tokens -------------->|
282
- |<--prompt_0-->|<--prompt_1-->|...|<--prompt_N-1-->|
283
-
284
- Otherwise, the layout is as follows:
285
- |<------------------ num_generation_tokens (M) ----------------->|
286
- |<--generation_0-->|..........|<--generation_M-1-->|<--padding-->|
287
-
288
- Generation tokens can contain padding when cuda-graph is used.
289
- Currently, prompt tokens don't contain any padding.
290
-
291
- The prompts might have different lengths, while the generation tokens
292
- always have length 1.
293
-
294
- """
295
-
296
- def __init__(
297
- self,
298
- num_heads: int,
299
- head_size: int,
300
- scale: float,
301
- num_kv_heads: int,
302
- alibi_slopes: Optional[List[float]],
303
- sliding_window: Optional[int],
304
- kv_cache_dtype: str,
305
- blocksparse_params: Optional[Dict[str, Any]] = None,
306
- logits_soft_cap: Optional[float] = None,
307
- attn_type: str = AttentionType.DECODER,
308
- kv_sharing_target_layer_name: Optional[str] = None,
309
- ) -> None:
310
- if kv_sharing_target_layer_name is not None:
311
- raise NotImplementedError("KV sharing is not supported in V0 "
312
- "BLOCK_SPARSE_FLASH_ATTN Backend.")
313
- assert blocksparse_params is not None
314
- assert alibi_slopes is None, ValueError(
315
- "Alibi not support for blocksparse flash attention.")
316
- assert sliding_window is None, ValueError(
317
- "sliding_window is invalid for blocksparse attention.")
318
- assert logits_soft_cap is None, ValueError(
319
- "logits_soft_cap is invalid for blocksparse attention.")
320
-
321
- if "num_heads" not in blocksparse_params:
322
- blocksparse_params["num_heads"] = num_heads
323
- if "num_kv_heads" not in blocksparse_params:
324
- blocksparse_params["num_kv_heads"] = num_kv_heads or num_heads
325
- self.blocksparse_params = BlocksparseParams(**blocksparse_params)
326
- self.kv_cache_dtype = kv_cache_dtype
327
-
328
- self.num_heads = num_heads
329
- self.head_size = head_size
330
- self.scale = float(scale)
331
- self.alibi_slopes = alibi_slopes
332
- self.num_kv_heads = num_kv_heads
333
-
334
- self.num_queries_per_kv = self.num_heads // self.num_kv_heads
335
-
336
- self.local_blocks = self.blocksparse_params.local_blocks
337
- self.vert_stride = self.blocksparse_params.vert_stride
338
- self.sparse_block_size = self.blocksparse_params.block_size
339
- self.head_sliding_step = self.blocksparse_params.head_sliding_step
340
-
341
- supported_head_sizes = PagedAttention.get_supported_head_sizes()
342
- if head_size not in supported_head_sizes:
343
- raise ValueError(
344
- f"Head size {head_size} is not supported by PagedAttention. "
345
- f"Supported head sizes are: {supported_head_sizes}.")
346
-
347
- self.tp_size = get_tensor_model_parallel_world_size()
348
- self.tp_rank = get_tensor_model_parallel_rank()
349
-
350
- total_num_heads = num_heads * self.tp_size
351
- self.bs_attn = LocalStridedBlockSparseAttn(
352
- total_num_heads,
353
- self.blocksparse_params.max_seqlen,
354
- self.blocksparse_params.local_blocks,
355
- self.blocksparse_params.vert_stride,
356
- self.blocksparse_params.block_size,
357
- homo_head=self.blocksparse_params.homo_head,
358
- active_head_range=self.blocksparse_params.active_head_range,
359
- )
360
-
361
- if attn_type != AttentionType.DECODER:
362
- raise NotImplementedError("Encoder self-attention and "
363
- "encoder/decoder cross-attention "
364
- "are not implemented for "
365
- "BlocksparseFlashAttentionImpl")
366
-
367
- def forward(
368
- self,
369
- layer: AttentionLayer,
370
- query: torch.Tensor,
371
- key: torch.Tensor,
372
- value: torch.Tensor,
373
- kv_cache: torch.Tensor,
374
- attn_metadata: BlocksparseFlashAttentionMetadata,
375
- output: Optional[torch.Tensor] = None,
376
- output_scale: Optional[torch.Tensor] = None,
377
- ) -> torch.Tensor:
378
- """Forward pass with FlashAttention and PagedAttention.
379
-
380
- Args:
381
- query: shape = [num_tokens, num_heads * head_size]
382
- key: shape = [num_tokens, num_kv_heads * head_size]
383
- value: shape = [num_tokens, num_kv_heads * head_size]
384
- kv_cache = [2, num_blocks, block_size * num_kv_heads * head_size]
385
- NOTE: kv_cache will be an empty tensor with shape [0]
386
- for profiling run.
387
- attn_metadata: Metadata for attention.
388
- Returns:
389
- shape = [num_tokens, num_heads * head_size]
390
- """
391
- if output_scale is not None:
392
- raise NotImplementedError(
393
- "fused output quantization is not yet supported"
394
- " for BlocksparseFlashAttentionImpl")
395
-
396
- num_tokens, hidden_size = query.shape
397
- # Reshape the query, key, and value tensors.
398
- query = query.view(-1, self.num_heads, self.head_size)
399
- key = key.view(-1, self.num_kv_heads, self.head_size)
400
- value = value.view(-1, self.num_kv_heads, self.head_size)
401
-
402
- if kv_cache.numel() > 0:
403
- key_cache, value_cache = PagedAttention.split_kv_cache(
404
- kv_cache, self.num_kv_heads, self.head_size)
405
-
406
- # Reshape the input keys and values and store them in the cache.
407
- # If kv_cache is not provided, the new key and value tensors are
408
- # not cached. This happens during the initial memory profiling run.
409
-
410
- PagedAttention.write_to_paged_cache(
411
- key,
412
- value,
413
- key_cache,
414
- value_cache,
415
- attn_metadata.slot_mapping,
416
- self.kv_cache_dtype,
417
- layer._k_scale,
418
- layer._v_scale,
419
- )
420
-
421
- if prefill_meta := attn_metadata.prefill_metadata:
422
-
423
- # Prompt run.
424
- # normal attention
425
- # When block_tables are not filled, it means q and k are the
426
- # prompt, and they have the same length.
427
-
428
- assert kv_cache.numel() == 0 \
429
- or prefill_meta.block_tables is None \
430
- or prefill_meta.block_tables.numel() == 0, \
431
- "Does not support prefix-enabled attention."
432
-
433
- output = self.bs_attn(
434
- q=query,
435
- k=key,
436
- v=value,
437
- cu_seqlens_q=prefill_meta.seq_start_loc,
438
- cu_seqlens_k=prefill_meta.seq_start_loc,
439
- sm_scale=self.scale,
440
- )
441
-
442
- if decode_meta := attn_metadata.decode_metadata:
443
- # Decoding run.
444
- output = PagedAttention.forward_decode(
445
- query,
446
- key_cache,
447
- value_cache,
448
- decode_meta.block_tables,
449
- decode_meta.seq_lens_tensor,
450
- self.blocksparse_params.max_seqlen,
451
- self.kv_cache_dtype,
452
- self.num_kv_heads,
453
- self.scale,
454
- self.alibi_slopes,
455
- layer._k_scale,
456
- layer._v_scale,
457
- tp_rank=self.tp_rank,
458
- blocksparse_local_blocks=self.local_blocks,
459
- blocksparse_vert_stride=self.vert_stride,
460
- blocksparse_block_size=self.sparse_block_size,
461
- blocksparse_head_sliding_step=self.head_sliding_step,
462
- )
463
-
464
- assert output is not None
465
- # Reshape the output tensor.
466
- return output.view(num_tokens, hidden_size)
vllm/attention/backends/differential_flash_attn.py
-4
@@ -667,7 +667,6 @@ class DifferentialFlashAttentionImpl(AttentionImpl):
667
alibi_slopes: Optional[List[float]],
668
sliding_window: Optional[int],
669
kv_cache_dtype: str,
670
- blocksparse_params: Optional[Dict[str, Any]] = None,
670
logits_soft_cap: Optional[float] = None,
671
attn_type: str = AttentionType.DECODER,
672
kv_sharing_target_layer_name: Optional[str] = None,
@@ -680,9 +679,6 @@ class DifferentialFlashAttentionImpl(AttentionImpl):
679
differential_flash_attention_config
680
self.used_shared_kv_cache = kv_sharing_target_layer_name is not None
681
self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
683
- if blocksparse_params is not None:
684
- raise ValueError(
685
- "FlashAttention does not support block-sparse attention.")
682
if use_irope:
683
logger.warning(
684
"Using irope in V0 is not supported yet, it will fall back "
vllm/attention/backends/dual_chunk_flash_attn.py
-1
@@ -287,7 +287,6 @@ class DualChunkFlashAttentionImpl(FlashAttentionImpl):
287
alibi_slopes: Optional[List[float]],
288
sliding_window: Optional[int],
289
kv_cache_dtype: str,
290
- blocksparse_params: Optional[Dict[str, Any]] = None,
290
logits_soft_cap: Optional[float] = None,
291
attn_type: str = AttentionType.DECODER,
292
kv_sharing_target_layer_name: Optional[str] = None,
vllm/attention/backends/flash_attn.py
+1
-5
@@ -4,7 +4,7 @@
4
from collections import defaultdict
5
from dataclasses import dataclass
6
from itertools import accumulate
7
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
7
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type
8
9
import torch
10
@@ -615,7 +615,6 @@ class FlashAttentionImpl(AttentionImpl):
615
alibi_slopes: Optional[List[float]],
616
sliding_window: Optional[int],
617
kv_cache_dtype: str,
618
- blocksparse_params: Optional[Dict[str, Any]] = None,
618
logits_soft_cap: Optional[float] = None,
619
attn_type: str = AttentionType.DECODER,
620
kv_sharing_target_layer_name: Optional[str] = None,
@@ -624,9 +623,6 @@ class FlashAttentionImpl(AttentionImpl):
623
if kv_sharing_target_layer_name is not None:
624
raise NotImplementedError("KV sharing is not supported in V0 "
625
"FLASH_ATTN backend.")
627
- if blocksparse_params is not None:
628
- raise ValueError(
629
- "FlashAttention does not support block-sparse attention.")
626
if use_irope:
627
logger.warning(
628
"Using irope in V0 is not supported yet, it will fall back "
vllm/attention/backends/flashinfer.py
-1
@@ -999,7 +999,6 @@ class FlashInferImpl(AttentionImpl):
999
alibi_slopes: Optional[List[float]],
1000
sliding_window: Optional[int],
1001
kv_cache_dtype: str,
1002
- blocksparse_params: Optional[Dict[str, Any]] = None,
1002
logits_soft_cap: Optional[float] = None,
1003
attn_type: str = AttentionType.DECODER,
1004
kv_sharing_target_layer_name: Optional[str] = None,
vllm/attention/backends/flashmla.py
+4
-8
@@ -3,7 +3,7 @@
3
4
from contextlib import contextmanager
5
from dataclasses import dataclass
6
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
6
+from typing import TYPE_CHECKING, List, Optional, Tuple, Type
7
8
import torch
9
@@ -181,7 +181,6 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
181
alibi_slopes: Optional[List[float]],
182
sliding_window: Optional[int],
183
kv_cache_dtype: str,
184
- blocksparse_params: Optional[Dict[str, Any]],
184
logits_soft_cap: Optional[float],
185
attn_type: str,
186
kv_sharing_target_layer_name: Optional[str] = None,
@@ -189,20 +188,17 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
188
**mla_args) -> None:
189
super().__init__(num_heads, head_size, scale, num_kv_heads,
190
alibi_slopes, sliding_window, kv_cache_dtype,
192
- blocksparse_params, logits_soft_cap, attn_type,
191
+ logits_soft_cap, attn_type,
192
kv_sharing_target_layer_name, **mla_args)
193
194
assert is_flashmla_supported(), \
195
"FlashMLA is not supported on this device"
196
198
- unsupported_features = [
199
- alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
200
- ]
197
+ unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
198
if any(unsupported_features):
199
raise NotImplementedError(
200
"FlashMLAImpl does not support one of the following: "
204
- "alibi_slopes, sliding_window, blocksparse_params, "
205
- "logits_soft_cap")
201
+ "alibi_slopes, sliding_window, logits_soft_cap")
202
203
if attn_type != AttentionType.DECODER:
204
raise NotImplementedError("Encoder self-attention and "
vllm/attention/backends/mla/common.py
-1
@@ -997,7 +997,6 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
997
alibi_slopes: Optional[List[float]],
998
sliding_window: Optional[int],
999
kv_cache_dtype: str,
1000
- blocksparse_params: Optional[Dict[str, Any]],
1000
logits_soft_cap: Optional[float],
1001
attn_type: str,
1002
kv_sharing_target_layer_name: Optional[str],
vllm/attention/backends/rocm_aiter_mla.py
+4
-8
@@ -3,7 +3,7 @@
3
4
from contextlib import contextmanager
5
from dataclasses import dataclass
6
-from typing import TYPE_CHECKING, Any, Optional, Type, Union
6
+from typing import TYPE_CHECKING, Optional, Type, Union
7
8
import torch
9
@@ -367,7 +367,6 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
367
alibi_slopes: Optional[list[float]],
368
sliding_window: Optional[int],
369
kv_cache_dtype: str,
370
- blocksparse_params: Optional[dict[str, Any]],
370
logits_soft_cap: Optional[float],
371
attn_type: str,
372
kv_sharing_target_layer_name: Optional[str],
@@ -375,17 +374,14 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
374
**mla_args) -> None:
375
super().__init__(num_heads, head_size, scale, num_kv_heads,
376
alibi_slopes, sliding_window, kv_cache_dtype,
378
- blocksparse_params, logits_soft_cap, attn_type,
377
+ logits_soft_cap, attn_type,
378
kv_sharing_target_layer_name, **mla_args)
379
381
- unsupported_features = [
382
- alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
383
- ]
380
+ unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
381
if any(unsupported_features):
382
raise NotImplementedError(
383
"Aiter MLA does not support one of the following: "
387
- "alibi_slopes, sliding_window, blocksparse_params, "
388
- "logits_soft_cap")
384
+ "alibi_slopes, sliding_window, logits_soft_cap")
385
386
from aiter import flash_attn_varlen_func
387
self.flash_attn_varlen_func = flash_attn_varlen_func
vllm/attention/backends/rocm_flash_attn.py
+1
-5
@@ -4,7 +4,7 @@
4
import itertools
5
from dataclasses import dataclass
6
from functools import cache
7
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
7
+from typing import TYPE_CHECKING, List, Optional, Tuple, Type
8
9
import torch
10
@@ -494,7 +494,6 @@ class ROCmFlashAttentionImpl(AttentionImpl):
494
alibi_slopes: Optional[List[float]],
495
sliding_window: Optional[int],
496
kv_cache_dtype: str,
497
- blocksparse_params: Optional[Dict[str, Any]] = None,
497
logits_soft_cap: Optional[float] = None,
498
attn_type: str = AttentionType.DECODER,
499
kv_sharing_target_layer_name: Optional[str] = None,
@@ -507,9 +506,6 @@ class ROCmFlashAttentionImpl(AttentionImpl):
506
logger.warning_once(
507
"Using irope in ROCm Flash Attention is not supported yet, it "
508
"will fail back to global attention for long context.")
510
- if blocksparse_params is not None:
511
- raise ValueError(
512
- "ROCmFlashAttention does not support blocksparse attention.")
509
if use_irope:
510
logger.warning(
511
"Using irope in V0 is not supported yet, it will fall back "
vllm/attention/backends/triton_mla.py
+4
-8
@@ -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 Any, Dict, List, Optional, Type
4
+from typing import List, Optional, Type
5
6
import torch
7
@@ -35,7 +35,6 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
35
alibi_slopes: Optional[List[float]],
36
sliding_window: Optional[int],
37
kv_cache_dtype: str,
38
- blocksparse_params: Optional[Dict[str, Any]],
38
logits_soft_cap: Optional[float],
39
attn_type: str,
40
kv_sharing_target_layer_name: Optional[str],
@@ -43,17 +42,14 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
42
**mla_args) -> None:
43
super().__init__(num_heads, head_size, scale, num_kv_heads,
44
alibi_slopes, sliding_window, kv_cache_dtype,
46
- blocksparse_params, logits_soft_cap, attn_type,
45
+ logits_soft_cap, attn_type,
46
kv_sharing_target_layer_name, **mla_args)
47
49
- unsupported_features = [
50
- alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
51
- ]
48
+ unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
49
if any(unsupported_features):
50
raise NotImplementedError(
51
"TritonMLAImpl does not support one of the following: "
55
- "alibi_slopes, sliding_window, blocksparse_params, "
56
- "logits_soft_cap")
52
+ "alibi_slopes, sliding_window, logits_soft_cap")
53
54
if attn_type != AttentionType.DECODER:
55
raise NotImplementedError("Encoder self-attention and "
vllm/attention/backends/xformers.py
+1
-5
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer with xFormers and PagedAttention."""
4
from dataclasses import dataclass
5
-from typing import Any, Dict, List, Optional, Tuple, Type
5
+from typing import Dict, List, Optional, Tuple, Type
6
7
import torch
8
from xformers import ops as xops
@@ -387,7 +387,6 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
387
alibi_slopes: Optional[List[float]],
388
sliding_window: Optional[int],
389
kv_cache_dtype: str,
390
- blocksparse_params: Optional[Dict[str, Any]] = None,
390
logits_soft_cap: Optional[float] = None,
391
attn_type: str = AttentionType.DECODER,
392
kv_sharing_target_layer_name: Optional[str] = None,
@@ -396,9 +395,6 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
395
if kv_sharing_target_layer_name is not None:
396
raise NotImplementedError("KV sharing is not supported in V0 "
397
"XFORMERS backend.")
399
- if blocksparse_params is not None:
400
- raise ValueError(
401
- "XFormers does not support block-sparse attention.")
398
if logits_soft_cap is not None:
399
logger.warning_once("XFormers does not support logits soft cap. "
400
"Outputs may be slightly off.")
vllm/attention/layer.py
+2
-4
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer."""
4
-from typing import Any, Dict, List, Optional
4
+from typing import List, Optional
5
6
import torch
7
import torch.nn as nn
@@ -74,7 +74,6 @@ class Attention(nn.Module):
74
alibi_slopes: Optional[List[float]] = None,
75
cache_config: Optional[CacheConfig] = None,
76
quant_config: Optional[QuantizationConfig] = None,
77
- blocksparse_params: Optional[Dict[str, Any]] = None,
77
logits_soft_cap: Optional[float] = None,
78
per_layer_sliding_window: Optional[int] = None,
79
use_mla: bool = False,
@@ -163,12 +162,11 @@ class Attention(nn.Module):
162
kv_cache_dtype,
163
block_size,
164
is_attention_free,
166
- blocksparse_params is not None,
165
use_mla=use_mla)
166
impl_cls = attn_backend.get_impl_cls()
167
self.impl = impl_cls(num_heads, head_size, scale, num_kv_heads,
168
alibi_slopes, sliding_window, kv_cache_dtype,
171
- blocksparse_params, logits_soft_cap, attn_type,
169
+ logits_soft_cap, attn_type,
170
kv_sharing_target_layer_name, **extra_impl_args)
171
self.backend = backend_name_to_enum(attn_backend.get_name())
172
self.dtype = dtype
vllm/attention/ops/blocksparse_attention/__init__.py
vllm/attention/ops/blocksparse_attention/blocksparse_attention_kernel.py
deleted
-433
@@ -1,433 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import torch
5
-
6
-from vllm.triton_utils import tl, triton
7
-
8
-
9
-def blocksparse_flash_attn_varlen_fwd(
10
- q,
11
- k,
12
- v, # (#tokens, n_heads, head_size)
13
- cu_seqlens_k,
14
- cu_seqlens_q,
15
- sm_scale,
16
- sparse_layout,
17
- *,
18
- block_size=64,
19
- q_block_size=None,
20
- max_seqlen=None):
21
- # split q to blocks
22
-
23
- assert isinstance(sparse_layout, (list, tuple))
24
-
25
- _, n_heads, head_size = q.shape
26
- batch_size = cu_seqlens_k.size(0) - 1
27
- q_block_size = q_block_size or block_size
28
-
29
- assert q.dim() == k.dim() == v.dim() == 3
30
- assert q.size(1) % k.size(1) == 0
31
- assert q.size(2) == k.size(2)
32
- # TODO(linxihui): allow k, v to have different head_size
33
- assert k.shape == v.shape
34
- assert cu_seqlens_k.dim() == 1
35
-
36
- q_k_ratio = q.size(1) // k.size(1)
37
-
38
- if cu_seqlens_q is None:
39
- if q.size(0) == batch_size: # decoding only
40
- cu_seqlens_q = torch.arange(
41
- 0,
42
- batch_size + 1,
43
- dtype=cu_seqlens_k.dtype,
44
- device=cu_seqlens_k.device,
45
- )
46
- elif q.size(0) == k.size(0):
47
- cu_seqlens_q = cu_seqlens_k
48
- else:
49
- raise ValueError("cu_seqlens_q must be specified\
50
- if it mix of prefilling and decoding.")
51
- else:
52
- assert cu_seqlens_k.size(0) == cu_seqlens_q.size(0)
53
-
54
- # switch to use cpu to avoid too many kernel launches when iterated over
55
- q_lens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).cpu()
56
- k_lens = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]).cpu()
57
-
58
- assert torch.logical_or(q_lens == 1, k_lens == q_lens).all(), (
59
- "length of q should either be 1 (decoding) or same as k (prefilling).")
60
-
61
- if max_seqlen:
62
- assert k_lens.max() <= max_seqlen
63
-
64
- n_blocks = (q_lens + q_block_size - 1) // q_block_size
65
-
66
- q_batch_ids = torch.tensor(
67
- [i for i, n in enumerate(n_blocks) for _ in range(n)],
68
- dtype=cu_seqlens_q.dtype,
69
- device=cu_seqlens_q.device,
70
- )
71
- q_start_sids = torch.tensor(
72
- [i * q_block_size for n in n_blocks for i in range(n)],
73
- dtype=cu_seqlens_q.dtype,
74
- device=cu_seqlens_q.device,
75
- )
76
-
77
- out = q.new_empty(q.shape)
78
- cu_seqlens_q = cu_seqlens_q.contiguous()
79
- cu_seqlens_k = cu_seqlens_k.contiguous()
80
-
81
- layout_crow_indices, layout_col_indices = sparse_layout
82
- block_d = triton.next_power_of_2(head_size)
83
-
84
- decoding_only = (q_lens == 1).all().item()
85
- grid = (len(q_start_sids), n_heads, 1)
86
-
87
- _fwd_kernel_batch_inference[grid](
88
- q,
89
- k,
90
- v,
91
- out,
92
- sm_scale,
93
- cu_seqlens_q[:-1],
94
- cu_seqlens_q[1:],
95
- cu_seqlens_k[:-1],
96
- cu_seqlens_k[1:],
97
- q_batch_ids,
98
- q_start_sids,
99
- 0,
100
- *q.stride(),
101
- 0,
102
- *k.stride(),
103
- 0,
104
- *v.stride(),
105
- 0,
106
- *out.stride(),
107
- layout_crow_indices,
108
- layout_col_indices,
109
- *layout_crow_indices.stride(),
110
- *layout_col_indices.stride(),
111
- q_k_ratio,
112
- HAS_BATCH_DIM=False,
113
- D_HEAD=head_size,
114
- BLOCK_M=q_block_size,
115
- BLOCK_N=block_size,
116
- BLOCK_D=block_d,
117
- BLOCK_M_LOADING=(16 if decoding_only else
118
- q_block_size), # smaller for decoding
119
- EVEN_D=block_d == head_size,
120
- num_warps=1 if decoding_only else 4,
121
- num_stages=3)
122
-
123
- return out
124
-
125
-
126
-@triton.jit
127
-def _fwd_kernel_inner(
128
- acc,
129
- l_i,
130
- m_i,
131
- q,
132
- Q,
133
- k_block_col_idx,
134
- layout_col_ptr,
135
- layout_col_stride_h,
136
- layout_col_stride_m,
137
- k_ptrs,
138
- v_ptrs,
139
- off_h,
140
- offs_m,
141
- offs_n,
142
- offs_d,
143
- stride_kt,
144
- stride_vt,
145
- sm_scale,
146
- k_seqlen,
147
- past_len,
148
- LAST_K_BLOCK: tl.constexpr,
149
- BLOCK_M_LOADING: tl.constexpr,
150
- BLOCK_N: tl.constexpr,
151
- D_HEAD: tl.constexpr,
152
- EVEN_D: tl.constexpr,
153
- M_LT_N: tl.constexpr,
154
-):
155
- k_block_id = tl.load(layout_col_ptr + off_h * layout_col_stride_h +
156
- k_block_col_idx * layout_col_stride_m).to(tl.int32)
157
- start_n = k_block_id * BLOCK_N
158
- if LAST_K_BLOCK:
159
- if EVEN_D:
160
- k = tl.load(
161
- k_ptrs + start_n * stride_kt,
162
- mask=offs_n[None, :] + start_n < k_seqlen,
163
- other=0.0,
164
- )
165
- else:
166
- k = tl.load(
167
- k_ptrs + start_n * stride_kt,
168
- mask=(offs_n[None, :] + start_n < k_seqlen) &
169
- (offs_d[:, None] < D_HEAD),
170
- other=0.0,
171
- )
172
- else:
173
- if EVEN_D:
174
- k = tl.load(k_ptrs + start_n * stride_kt)
175
- else:
176
- k = tl.load(k_ptrs + start_n * stride_kt,
177
- mask=offs_d[:, None] < D_HEAD,
178
- other=0.0)
179
-
180
- qk = tl.zeros([BLOCK_M_LOADING, BLOCK_N], dtype=tl.float32)
181
- qk += tl.dot(q, k)
182
- qk *= sm_scale
183
-
184
- # the following is needed only when LAST_K_BLOCK or BLOCK_M < BLOCK_N
185
- if LAST_K_BLOCK | M_LT_N:
186
- qk += tl.where(
187
- offs_m[:, None] + past_len >= (start_n + offs_n[None, :]),
188
- 0,
189
- float("-inf"),
190
- )
191
-
192
- # flash-attn2
193
- m_ij = tl.maximum(m_i, tl.max(qk, 1))
194
- p = tl.math.exp2(qk - m_ij[:, None])
195
- l_ij = tl.sum(p, 1)
196
- alpha = tl.math.exp2(m_i - m_ij)
197
- acc = acc * alpha[:, None]
198
- # update m_i
199
- m_i = m_ij
200
- l_i = l_i * alpha + l_ij
201
-
202
- p = p.to(Q.dtype.element_ty)
203
- # update acc
204
- if LAST_K_BLOCK:
205
- if EVEN_D:
206
- v = tl.load(
207
- v_ptrs + start_n * stride_vt,
208
- mask=offs_n[:, None] + start_n < k_seqlen,
209
- other=0.0,
210
- )
211
- else:
212
- v = tl.load(
213
- v_ptrs + start_n * stride_vt,
214
- mask=(offs_n[:, None] + start_n < k_seqlen) &
215
- (offs_d[None, :] < D_HEAD),
216
- other=0.0,
217
- )
218
- else:
219
- if EVEN_D:
220
- v = tl.load(v_ptrs + start_n * stride_vt)
221
- else:
222
- v = tl.load(v_ptrs + start_n * stride_vt,
223
- mask=offs_d[None, :] < D_HEAD,
224
- other=0.0)
225
-
226
- acc += tl.dot(p, v)
227
-
228
- return acc, l_i, m_i
229
-
230
-
231
-@triton.heuristics({
232
- "M_LT_N":
233
- lambda kwargs: kwargs["BLOCK_M"] < kwargs["BLOCK_N"],
234
-})
235
-@triton.jit
236
-def _fwd_kernel_batch_inference(
237
- Q,
238
- K,
239
- V,
240
- Out,
241
- sm_scale,
242
- q_batch_starts,
243
- q_batch_ends,
244
- k_batch_starts,
245
- k_batch_ends,
246
- q_batch_ids,
247
- q_start_sids,
248
- stride_qb,
249
- stride_qt,
250
- stride_qh,
251
- stride_qd,
252
- stride_kb,
253
- stride_kt,
254
- stride_kh,
255
- stride_kd,
256
- stride_vb,
257
- stride_vt,
258
- stride_vh,
259
- stride_vd,
260
- stride_ob,
261
- stride_ot,
262
- stride_oh,
263
- stride_od,
264
- layout_crow_ptr,
265
- layout_col_ptr,
266
- layout_crow_stride_h,
267
- layout_crow_stride_m,
268
- layout_col_stride_h,
269
- layout_col_stride_m,
270
- q_k_ratio,
271
- HAS_BATCH_DIM: tl.constexpr,
272
- D_HEAD: tl.constexpr,
273
- BLOCK_M: tl.constexpr,
274
- BLOCK_N: tl.constexpr,
275
- BLOCK_D: tl.constexpr,
276
- BLOCK_M_LOADING: tl.constexpr,
277
- EVEN_D: tl.constexpr,
278
- M_LT_N: tl.constexpr,
279
-):
280
- """
281
- NOTATION:
282
- pid: position id
283
- sid: storage id
284
- sbid: storage block id
285
- pbid: position block id
286
- offs_m, offs_n: storage offsets of m-dim(q, row) and n-dim(k, col)
287
-
288
- TODO(linxihui):
289
- Optimize grouped-attn
290
- """
291
- off_zm = tl.program_id(0)
292
- off_h = tl.program_id(1)
293
-
294
- off_h_for_kv = off_h // q_k_ratio
295
-
296
- if HAS_BATCH_DIM:
297
- off_z = tl.program_id(2)
298
- Q += off_z * stride_qb
299
- K += off_z * stride_kb
300
- V += off_z * stride_vb
301
- Out += off_z * stride_ob
302
- start_m = off_zm
303
- q_start_sid = start_m * BLOCK_M # always 0 for decoding
304
- else:
305
- off_z = tl.load(q_batch_ids + off_zm).to(tl.int32) # [0, 0, 0, 1]
306
- q_start_sid = tl.load(q_start_sids + off_zm)
307
- start_m = q_start_sid // BLOCK_M # q_sbid
308
-
309
- offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M_LOADING)
310
- offs_n = tl.arange(0, BLOCK_N)
311
- offs_d = tl.arange(0, BLOCK_D)
312
-
313
- q_cu_start = tl.load(q_batch_starts + off_z).to(tl.int32)
314
- q_seqlen = tl.load(q_batch_ends + off_z).to(tl.int32) - q_cu_start
315
- k_cu_start = tl.load(k_batch_starts + off_z).to(tl.int32)
316
- k_seqlen = tl.load(k_batch_ends + off_z).to(tl.int32) - k_cu_start
317
- past_len = k_seqlen - q_seqlen
318
-
319
- Q += q_cu_start * stride_qt + off_h * stride_qh
320
- K += k_cu_start * stride_kt + off_h_for_kv * stride_kh
321
- V += k_cu_start * stride_vt + off_h_for_kv * stride_vh
322
- Out += q_cu_start * stride_ot + off_h * stride_oh
323
-
324
- q_pbid = (past_len + q_start_sid) // BLOCK_M
325
-
326
- if EVEN_D:
327
- q = tl.load(
328
- Q + offs_m[:, None] * stride_qt + offs_d[None, :] * stride_qd,
329
- mask=offs_m[:, None] < q_seqlen,
330
- other=0.0,
331
- )
332
- else:
333
- q = tl.load(
334
- Q + offs_m[:, None] * stride_qt + offs_d[None, :] * stride_qd,
335
- mask=(offs_m[:, None] < q_seqlen) & (offs_d[None, :] < D_HEAD),
336
- other=0.0,
337
- )
338
-
339
- sparse_crow_ptr = (layout_crow_ptr + off_h * layout_crow_stride_h +
340
- q_pbid * layout_crow_stride_m)
341
-
342
- # TODO(linxihui): load at once, with any Triton version
343
- # that supports `tl.split`, e.g., Triton 3.0
344
- k_block_start = tl.load(sparse_crow_ptr).to(tl.int32)
345
- k_block_end = tl.load(sparse_crow_ptr + 1).to(tl.int32)
346
-
347
- m_i = tl.zeros([BLOCK_M_LOADING], dtype=tl.float32) - float("inf")
348
- l_i = tl.zeros([BLOCK_M_LOADING], dtype=tl.float32)
349
- acc = tl.zeros([BLOCK_M_LOADING, BLOCK_D], dtype=tl.float32)
350
-
351
- k_ptrs = K + offs_n[None, :] * stride_kt + offs_d[:, None] * stride_kd
352
- v_ptrs = V + offs_n[:, None] * stride_vt + offs_d[None, :] * stride_vd
353
-
354
- sm_scale *= (
355
- 1.44269504 # 1/log2 as we use base2 for exponential and logarithm
356
- )
357
-
358
- for k_block_col_idx in range(k_block_start, k_block_end - 1):
359
- acc, l_i, m_i = _fwd_kernel_inner(
360
- acc,
361
- l_i,
362
- m_i,
363
- q,
364
- Q,
365
- k_block_col_idx,
366
- layout_col_ptr,
367
- layout_col_stride_h,
368
- layout_col_stride_m,
369
- k_ptrs,
370
- v_ptrs,
371
- off_h,
372
- offs_m,
373
- offs_n,
374
- offs_d,
375
- stride_kt,
376
- stride_vt,
377
- sm_scale,
378
- k_seqlen,
379
- past_len,
380
- False,
381
- BLOCK_M_LOADING,
382
- BLOCK_N,
383
- D_HEAD,
384
- EVEN_D,
385
- M_LT_N,
386
- )
387
-
388
- acc, l_i, m_i = _fwd_kernel_inner(
389
- acc,
390
- l_i,
391
- m_i,
392
- q,
393
- Q,
394
- k_block_end - 1,
395
- layout_col_ptr,
396
- layout_col_stride_h,
397
- layout_col_stride_m,
398
- k_ptrs,
399
- v_ptrs,
400
- off_h,
401
- offs_m,
402
- offs_n,
403
- offs_d,
404
- stride_kt,
405
- stride_vt,
406
- sm_scale,
407
- k_seqlen,
408
- past_len,
409
- True,
410
- BLOCK_M_LOADING,
411
- BLOCK_N,
412
- D_HEAD,
413
- EVEN_D,
414
- M_LT_N,
415
- )
416
-
417
- # flash-attn 2
418
- m_i += tl.math.log2(l_i)
419
- acc = acc / l_i[:, None]
420
-
421
- # write output
422
- if EVEN_D:
423
- tl.store(
424
- Out + offs_m[:, None] * stride_ot + offs_d[None, :] * stride_od,
425
- acc,
426
- mask=offs_m[:, None] < q_seqlen,
427
- )
428
- else:
429
- tl.store(
430
- Out + offs_m[:, None] * stride_ot + offs_d[None, :] * stride_od,
431
- acc,
432
- mask=(offs_m[:, None] < q_seqlen) & (offs_d[None, :] < D_HEAD),
433
- )
vllm/attention/ops/blocksparse_attention/interface.py
deleted
-239
@@ -1,239 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import math
5
-
6
-import torch
7
-
8
-from vllm.platforms import current_platform
9
-
10
-from .utils import (dense_to_crow_col, get_head_sliding_step,
11
- get_sparse_attn_mask)
12
-
13
-IS_COMPUTE_8_OR_ABOVE = current_platform.has_device_capability(80)
14
-
15
-if IS_COMPUTE_8_OR_ABOVE:
16
- from .blocksparse_attention_kernel import blocksparse_flash_attn_varlen_fwd
17
-
18
-
19
-class LocalStridedBlockSparseAttn(torch.nn.Module):
20
-
21
- def __init__(
22
- self,
23
- n_heads,
24
- max_seqlen,
25
- local_blocks,
26
- vert_stride,
27
- block_size,
28
- device=None,
29
- dtype=None,
30
- homo_head=False,
31
- active_head_range=None,
32
- q_block_size=None,
33
- use_spda=None,
34
- ):
35
- super().__init__()
36
- if use_spda is None:
37
- use_spda = current_platform.is_rocm() or \
38
- current_platform.is_cpu() or not \
39
- IS_COMPUTE_8_OR_ABOVE
40
- device = device or (torch.cuda.current_device()
41
- if current_platform.is_cuda_alike() else "cpu")
42
- device = torch.device(device)
43
- # NOTE: vllm CPU backend support BF16 instead of FP16.
44
- dtype = dtype or (torch.bfloat16 if IS_COMPUTE_8_OR_ABOVE
45
- or device.type == "cpu" else torch.half)
46
-
47
- self.n_heads = n_heads
48
- self.max_seqlen = max_seqlen
49
- self.local_blocks = local_blocks
50
- self.vert_stride = vert_stride
51
- self.use_spda = use_spda
52
- self.dtype = dtype
53
- self.device = device
54
- self.block_size = block_size
55
- self.q_block_size = q_block_size
56
- self.homo_head = homo_head
57
- self.active_head_range = active_head_range
58
- self.head_sliding_step = get_head_sliding_step(n_heads, vert_stride,
59
- homo_head)
60
-
61
- sparse_layout, sparse_pattern, self.dense_attn_mask = (
62
- self.get_attn_pattern(dtype, device))
63
-
64
- if q_block_size is not None and q_block_size != block_size:
65
- if q_block_size > block_size:
66
- assert q_block_size % block_size == 0
67
- blocks_to_merge = q_block_size // block_size
68
- shape = sparse_pattern.shape
69
- sparse_pattern = sparse_pattern.view(shape[0], -1,
70
- blocks_to_merge,
71
- shape[-1])
72
- sparse_pattern = sparse_pattern.sum(2)
73
- sparse_layout = dense_to_crow_col(sparse_pattern)
74
- else:
75
- raise ValueError(
76
- "Does not support smaller q_block_size. It will be slower."
77
- )
78
-
79
- self.sparse_layout = sparse_layout
80
-
81
- def get_attn_pattern(self, dtype, device):
82
- sparse_layout, sparse_pattern, dense_attn_mask = get_sparse_attn_mask(
83
- self.n_heads,
84
- self.max_seqlen,
85
- self.max_seqlen,
86
- dtype,
87
- device,
88
- block_size=self.block_size,
89
- local_blocks=self.local_blocks,
90
- vert_stride=self.vert_stride,
91
- homo_head=self.homo_head,
92
- return_dense=self.use_spda,
93
- dense_mask_type="bias",
94
- )
95
- if (not self.homo_head) and (self.active_head_range is not None):
96
- assert isinstance(self.active_head_range, tuple)
97
- assert (len(self.active_head_range) == 2)
98
- h_start, h_end = self.active_head_range
99
- sparse_layout = tuple(x[h_start:h_end] for x in sparse_layout)
100
- if self.use_spda:
101
- dense_attn_mask = dense_attn_mask[h_start:h_end]
102
- return sparse_layout, sparse_pattern, dense_attn_mask
103
-
104
- def varlen_attn(self,
105
- q,
106
- k,
107
- v,
108
- cu_seqlens_k,
109
- cu_seqlens_q=None,
110
- sm_scale=None):
111
- """
112
- q, k, v: shape = (num_tokens, num_heads_q/kv, head_size).
113
- Support grouped attention, with `q[:, i*r:(i*r + r)]`
114
- is correspondent to `k[:, i]`, where `r` is the q/k ratio.
115
- cu_seqlens_k: shape=(batch_size + 1,),
116
- indicating segment of samples,
117
- e.g., `k[cu_seqlen[i]:cu_seqlne[i+1]]` is q of sample i
118
- cu_seqlens_q: shape=(batch_size + 1, ).
119
- Default None: same as cu_seqlens_k for prefilling or
120
- [0, 1, .., batch_size] for decoding.
121
- The only case you need to specify is when q is a mix of
122
- prefilling and decoding.
123
- sm_scale: softmax scale, default to 1/sqrt(head_size).
124
-
125
- return: tensor of shape as q.
126
- """
127
- assert (
128
- IS_COMPUTE_8_OR_ABOVE
129
- ), "Requires compute capability of 8 or above (Ampere or newer) to use \
130
- Triton kernel."
131
-
132
- sm_scale = sm_scale or 1.0 / math.sqrt(q.size(-1))
133
-
134
- return blocksparse_flash_attn_varlen_fwd(
135
- q,
136
- k,
137
- v,
138
- cu_seqlens_k,
139
- cu_seqlens_q,
140
- sm_scale,
141
- self.sparse_layout,
142
- block_size=self.block_size,
143
- q_block_size=self.q_block_size,
144
- max_seqlen=self.max_seqlen,
145
- )
146
-
147
- @staticmethod
148
- def transpose_and_pad(x, cu_seqlens, maxlen, head_repeats=1):
149
- """
150
- :param x: (total_tokens, n_heads, head_size)
151
- :return: (batch, n_heads, length, head_size)
152
- """
153
- x_padded = x.new_empty(
154
- len(cu_seqlens) - 1, x.size(1), head_repeats, maxlen, x.size(2))
155
- cu_seqlens = cu_seqlens.cpu()
156
- for i, (s, e) in enumerate(zip(cu_seqlens[:-1], cu_seqlens[1:])):
157
- x_padded[i, :, :, :e - s].copy_(x[s:e].transpose(0,
158
- 1).unsqueeze(1))
159
- return x_padded.flatten(1, 2)
160
-
161
- @staticmethod
162
- def transpose_and_unpad(x_padded, cu_seqlens):
163
- """
164
- :param x_padded: (batch, n_heads, length, head_size)
165
- :return: (total_tokens, n_heads, head_size)
166
- """
167
- cu_seqlens = cu_seqlens.cpu()
168
- total_n_tokens = cu_seqlens[-1]
169
- x = x_padded.new_empty(total_n_tokens, x_padded.size(1),
170
- x_padded.size(3))
171
- for i, (s, e) in enumerate(zip(cu_seqlens[:-1], cu_seqlens[1:])):
172
- x[s:e].copy_(x_padded[i, :, :e - s].transpose(0, 1))
173
- return x
174
-
175
- def spda(self, q, k, v, cu_seqlens_k, cu_seqlens_q=None, sm_scale=None):
176
- """For CPU, V100 or other older GPUs.
177
- NOTE: torch SPDA supports nested tensor,
178
- but seems extremely slow. Choose to pad instead.
179
- """
180
- assert (cu_seqlens_q is None or
181
- (cu_seqlens_q
182
- == cu_seqlens_k).all()), "Can only handle prompt with SPDA."
183
- assert q.size(0) == k.size(0), "can only handle prompt with SPDA."
184
-
185
- assert q.size(1) % k.size(1) == 0
186
- q_k_ratio = q.size(1) // k.size(1)
187
- sm_scale = sm_scale or 1.0 / math.sqrt(q.size(-1))
188
- cu_seqlens = cu_seqlens_k.cpu()
189
- maxlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
190
-
191
- if (self.dense_attn_mask.dtype != q.dtype
192
- or self.dense_attn_mask.device != q.device):
193
- _, _, self.dense_attn_mask = self.get_attn_pattern(
194
- q.dtype, q.device)
195
- attn_mask = self.dense_attn_mask[None, :, :maxlen, :maxlen]
196
-
197
- q2 = self.transpose_and_pad(q, cu_seqlens, maxlen, 1)
198
- k2, v2 = (self.transpose_and_pad(x, cu_seqlens, maxlen, q_k_ratio)
199
- for x in [k, v])
200
- spda_output = torch.nn.functional.scaled_dot_product_attention(
201
- q2, k2, v2, attn_mask=attn_mask, scale=sm_scale)
202
- return self.transpose_and_unpad(spda_output, cu_seqlens)
203
-
204
- def forward(self, q, k, v, cu_seqlens_k, cu_seqlens_q=None, sm_scale=None):
205
- """Dispatch to `varlen_attn` (Ampere or newer) or
206
- `self.spda`(cpu, Volta, Turing or older)based on
207
- the type of device used and cuda compute capability.
208
-
209
- q, k, v: shape = (num_tokens, num_heads_q/kv, head_size).
210
- Support grouped attention, with `q[:, i*r:(i*r + r)]`
211
- is correspondent to `k[:, i]`, where `r` is the q/k ratio.
212
- cu_seqlens_k: shape=(batch_size + 1,), indicating segment of samples,
213
- e.g., `k[cu_seqlen[i]:cu_seqlne[i+1]]` is q of sample i
214
- cu_seqlens_q: shape=(batch_size + 1, ).
215
- Default None: same as cu_seqlens_k for prefilling or
216
- [0, 1, .., batch_size] for decoding.
217
- The only case you need to specify
218
- is when q is a mix of prefilling
219
- and decoding.
220
- sm_scale: softmax scale, default to 1/sqrt(head_size).
221
-
222
- return: tensor of shape as q.
223
- """
224
- assert k.dim() == 3
225
- if self.use_spda:
226
- return self.spda(
227
- q,
228
- k,
229
- v,
230
- cu_seqlens_k,
231
- cu_seqlens_q=cu_seqlens_q,
232
- sm_scale=sm_scale,
233
- )
234
- return self.varlen_attn(q,
235
- k,
236
- v,
237
- cu_seqlens_k,
238
- cu_seqlens_q=cu_seqlens_q,
239
- sm_scale=sm_scale)
vllm/attention/ops/blocksparse_attention/utils.py
deleted
-246
@@ -1,246 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-# Helper functions for 3D sparse pattern
5
-# These function are not optimized and very inefficient.
6
-# Avoid calling them too frequent or use a cache mechanism.
7
-
8
-from functools import lru_cache
9
-
10
-import numpy as np
11
-import torch
12
-
13
-from vllm.triton_utils import triton
14
-
15
-
16
-class csr_matrix:
17
- """Simple implementation of CSR matrix conversion without scipy.
18
- This replaced scipy.sparse.csr_matrix() previously used."""
19
-
20
- def __init__(self, input_array):
21
- if not isinstance(input_array, np.ndarray):
22
- raise ValueError("Input must be a NumPy array")
23
-
24
- self.shape = input_array.shape
25
- rows, cols = self.shape
26
- data = []
27
- indices = []
28
- indptr = [0]
29
-
30
- for i in range(rows):
31
- for j in range(cols):
32
- if input_array[i, j]:
33
- data.append(input_array[i, j])
34
- indices.append(j)
35
- indptr.append(len(indices))
36
-
37
- self.data = np.array(data)
38
- self.indices = np.array(indices)
39
- self.indptr = np.array(indptr)
40
-
41
-
42
-def dense_to_crow_col(x: torch.Tensor):
43
- """Turning a 2D/3D torch tensor (x) to CSR rows/cols indexing.
44
- NOTE: col_indices padded -1
45
- """
46
- device = x.device
47
- pad = -1
48
- dim = x.dim()
49
- assert x.dim() in (2, 3)
50
- if x.dim() == 2:
51
- x = x[None]
52
- x = [csr_matrix(xi.bool().cpu().numpy()) for xi in x]
53
- crows = torch.vstack([torch.from_numpy(xi.indptr) for xi in x])
54
- cols = [torch.from_numpy(xi.indices) for xi in x]
55
- max_cols = max(len(xi) for xi in cols)
56
- cols = [
57
- torch.cat([xi, pad + xi.new_zeros(max_cols - xi.shape[0])])
58
- for xi in cols
59
- ]
60
- cols = torch.vstack(cols)
61
- if dim == 2:
62
- crows = crows[0]
63
- cols = cols[0]
64
- return crows.to(device), cols.to(device)
65
-
66
-
67
-def crow_col_to_dense(crows: torch.Tensor,
68
- cols: torch.Tensor,
69
- dtype: torch.dtype = torch.float16):
70
- dim = crows.dim()
71
- if dim == 1:
72
- crows = crows[None]
73
- cols = cols[None]
74
- device = crows.device
75
- crows, cols = crows.cpu(), cols.cpu() # faster in cpu
76
- shape = (crows.shape[0], crows.shape[1] - 1, cols.max() + 1)
77
- x = torch.zeros(shape, dtype=dtype)
78
- for i in range(shape[0]):
79
- for j in range(shape[1]):
80
- x[i, j, cols[i, crows[i, j]:crows[i, j + 1]]] = 1
81
- if dim == 1:
82
- x = x[0]
83
- return x.to(device)
84
-
85
-
86
-def dense_to_ccol_row(x: torch.Tensor):
87
- """Similar, but to CSC format"""
88
- x = x.transpose(-2, -1)
89
- return dense_to_crow_col(x)
90
-
91
-
92
-def ccol_row_to_dense(ccol: torch.Tensor,
93
- rows: torch.Tensor,
94
- dtype: torch.dtype = torch.float16):
95
- return crow_col_to_dense(ccol, rows, dtype).permute(0, 2, 1).contiguous()
96
-
97
-
98
-def _get_sparse_attn_mask_homo_head(
99
- q_len: int,
100
- max_seqlen: int,
101
- dtype: torch.dtype,
102
- device: torch.device,
103
- block_size: int = 128,
104
- local_blocks: int = 4,
105
- vert_stride: int = 4,
106
- return_dense: bool = False,
107
-):
108
- """
109
- :return: a tuple of 3:
110
- - tuple of crow_indices, col_indices representation
111
- of CSR format.
112
- - block dense mask
113
- - all token dense mask (be aware that it can be
114
- OOM if it is too big) if `return_dense==True`,
115
- otherwise, None
116
- """
117
- with torch.no_grad():
118
- num_blocks = triton.cdiv(max_seqlen, block_size)
119
- q_pos = torch.arange(num_blocks)[:, None]
120
- k_pos = torch.arange(num_blocks)[None]
121
- mask_vert_strided = (torch.arange(num_blocks) + 1) % vert_stride == 0
122
- block_mask_dense = (((q_pos >= k_pos)
123
- & ((q_pos - k_pos < local_blocks)
124
- | mask_vert_strided)).to(device).to(dtype))
125
- num_blocks_q = triton.cdiv(q_len, block_size)
126
- block_mask_dense_output = (dense_to_crow_col(
127
- block_mask_dense[-num_blocks_q:].contiguous()))
128
- if return_dense:
129
- mask_dense = torch.kron(
130
- block_mask_dense,
131
- block_mask_dense.new_ones((block_size, block_size)),
132
- )
133
- causal_mask = torch.tril(torch.ones(
134
- max_seqlen, max_seqlen)).type_as(mask_dense)[-q_len:]
135
- mask_dense = mask_dense[-q_len:, :max_seqlen] * causal_mask
136
- return (
137
- block_mask_dense_output,
138
- block_mask_dense,
139
- mask_dense,
140
- )
141
- else:
142
- return (
143
- block_mask_dense_output,
144
- block_mask_dense,
145
- None,
146
- )
147
-
148
-
149
-def binary_mask_to_bias(mask_dense: torch.Tensor):
150
- mask_dense = 1 - mask_dense
151
- mask_dense.masked_fill_(mask_dense.bool(), -torch.inf)
152
- return mask_dense
153
-
154
-
155
-def get_head_sliding_step(n_heads: int,
156
- vert_stride: int,
157
- homo_head: bool = False):
158
- if homo_head:
159
- return 0
160
- return max(1, int(vert_stride / n_heads))
161
-
162
-
163
-@lru_cache
164
-def get_sparse_attn_mask(
165
- n_heads: int,
166
- q_len: int,
167
- max_seqlen: int,
168
- dtype: torch.dtype,
169
- device: torch.device,
170
- block_size: int = 64,
171
- local_blocks: int = 4,
172
- vert_stride: int = 4,
173
- homo_head: bool = True,
174
- return_dense: bool = False,
175
- dense_mask_type: str = "binary",
176
-):
177
- """
178
- :param dense_mask_type: "binary" (0 for skip token, 1 for others)
179
- or "bias" (-inf for skip token, 0 or others)
180
- :return: a tuple of 3:
181
- - tuple of crow_indices, col_indices representation
182
- of CSR format.
183
- - block dense mask
184
- - all token dense mask (be aware that it can be OOM if it
185
- is too big) if `return_dense==True`, otherwise, None
186
- """
187
- assert dense_mask_type in ("binary", "bias")
188
- if homo_head:
189
- with torch.no_grad():
190
- (crow, col), block_mask_dense, mask_dense = (
191
- _get_sparse_attn_mask_homo_head(
192
- q_len,
193
- max_seqlen,
194
- dtype,
195
- device,
196
- block_size,
197
- local_blocks,
198
- vert_stride,
199
- return_dense,
200
- ))
201
- crow = crow[None].expand(n_heads, crow.shape[0])
202
- col = col[None].expand(n_heads, col.shape[0])
203
- if return_dense:
204
- mask_dense = mask_dense[None].expand(n_heads,
205
- *mask_dense.shape)
206
- if dense_mask_type == "bias":
207
- mask_dense = binary_mask_to_bias(mask_dense)
208
- return (crow, col), block_mask_dense, mask_dense
209
-
210
- with torch.no_grad():
211
- num_blocks = triton.cdiv(max_seqlen, block_size)
212
- q_pos = torch.arange(num_blocks)[None, :, None]
213
- k_pos = torch.arange(num_blocks)[None, None]
214
- head_sliding_step = get_head_sliding_step(n_heads, vert_stride)
215
- mask_vert_strided = [
216
- (torch.arange(num_blocks) + h * head_sliding_step + 1) %
217
- vert_stride == 0 for h in range(n_heads)
218
- ]
219
- mask_vert_strided = torch.vstack(mask_vert_strided).unsqueeze(1)
220
- block_mask_dense = (((q_pos >= k_pos)
221
- & ((q_pos - k_pos < local_blocks)
222
- | mask_vert_strided)).to(device).to(dtype))
223
- num_blocks_q = triton.cdiv(q_len, block_size)
224
- block_mask_dense_output = block_mask_dense[:, -num_blocks_q:]
225
- if return_dense:
226
- mask_dense = torch.kron(
227
- block_mask_dense,
228
- block_mask_dense.new_ones((block_size, block_size)),
229
- )
230
- causal_mask = torch.tril(torch.ones(
231
- max_seqlen, max_seqlen)).type_as(mask_dense)[-q_len:]
232
- mask_dense = mask_dense[..., -q_len:, :max_seqlen] * causal_mask[None]
233
- if dense_mask_type == "bias":
234
- mask_dense = binary_mask_to_bias(mask_dense)
235
-
236
- return (
237
- dense_to_crow_col(block_mask_dense_output),
238
- block_mask_dense,
239
- mask_dense,
240
- )
241
- else:
242
- return (
243
- dense_to_crow_col(block_mask_dense_output),
244
- block_mask_dense,
245
- None,
246
- )
vllm/attention/selector.py
-9
@@ -143,7 +143,6 @@ def get_attn_backend(
143
kv_cache_dtype: Optional[str],
144
block_size: int,
145
is_attention_free: bool,
146
- is_blocksparse: bool = False,
146
use_mla: bool = False,
147
) -> type[AttentionBackend]:
148
"""Selects which attention backend to use and lazily imports it."""
@@ -157,7 +156,6 @@ def get_attn_backend(
156
kv_cache_dtype=kv_cache_dtype,
157
block_size=block_size,
158
is_attention_free=is_attention_free,
160
- is_blocksparse=is_blocksparse,
159
use_v1=envs.VLLM_USE_V1,
160
use_mla=use_mla,
161
)
@@ -170,16 +168,9 @@ def _cached_get_attn_backend(
168
kv_cache_dtype: Optional[str],
169
block_size: int,
170
is_attention_free: bool,
173
- is_blocksparse: bool = False,
171
use_v1: bool = False,
172
use_mla: bool = False,
173
) -> type[AttentionBackend]:
177
- if is_blocksparse:
178
- logger.info("Using BlocksparseFlashAttention backend.")
179
- from vllm.attention.backends.blocksparse_attn import (
180
- BlocksparseFlashAttentionBackend)
181
- return BlocksparseFlashAttentionBackend
182
-
174
# If there are no attention layers (e.g. we are running Mamba),
175
# use the placeholder NO_ATTENTION
176
if is_attention_free:
vllm/model_executor/models/phi3_small.py
deleted
-465
@@ -1,465 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import math
5
-from collections.abc import Iterable
6
-from typing import Optional, Union
7
-
8
-import torch
9
-from torch import nn
10
-from transformers.configuration_utils import PretrainedConfig
11
-
12
-from vllm.attention import Attention
13
-from vllm.config import CacheConfig, VllmConfig
14
-from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
15
- get_tensor_model_parallel_world_size)
16
-from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
17
- QKVParallelLinear,
18
- RowParallelLinear)
19
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
20
-from vllm.model_executor.layers.quantization import QuantizationConfig
21
-from vllm.model_executor.layers.rotary_embedding import get_rope
22
-from vllm.model_executor.layers.vocab_parallel_embedding import (
23
- DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
24
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
25
-from vllm.model_executor.sampling_metadata import SamplingMetadata
26
-from vllm.platforms import current_platform
27
-from vllm.sequence import IntermediateTensors
28
-
29
-from .interfaces import SupportsPP
30
-from .utils import (AutoWeightsLoader, WeightsMapper, is_pp_missing_parameter,
31
- make_empty_intermediate_tensors_factory, make_layers,
32
- maybe_prefix)
33
-
34
-
35
-def load_column_parallel_weight(param: torch.nn.Parameter,
36
- loaded_weight: torch.Tensor):
37
- tp = get_tensor_model_parallel_world_size()
38
- rk = get_tensor_model_parallel_rank()
39
- assert param.size(0) * tp == loaded_weight.size(0)
40
- s = rk * param.size(0)
41
- e = (rk + 1) * param.size(0)
42
- loaded_weight = loaded_weight[s:e]
43
- assert param.shape == loaded_weight.shape
44
- param.data.copy_(loaded_weight)
45
-
46
-
47
-class HeadMajorQKVParallelLinear(QKVParallelLinear):
48
-
49
- def weight_loader(self, param: torch.nn.Parameter,
50
- loaded_weight: torch.Tensor):
51
- return load_column_parallel_weight(param, loaded_weight)
52
-
53
-
54
-class HeadMajorColumnParallelLinear(MergedColumnParallelLinear):
55
-
56
- def weight_loader(self, param: torch.nn.Parameter,
57
- loaded_weight: torch.Tensor):
58
- return load_column_parallel_weight(param, loaded_weight)
59
-
60
-
61
-@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
62
-def quick_gelu(x):
63
- return x * torch.sigmoid(1.702 * x)
64
-
65
-
66
-@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
67
-def gegelu(input, limit: Optional[float] = None):
68
- a_gelu, a_linear = input[..., ::2], input[..., 1::2]
69
- if limit is not None:
70
- a_gelu = torch.where(torch.isinf(a_gelu), a_gelu,
71
- a_gelu.clamp(min=None, max=limit))
72
- a_linear = torch.where(
73
- torch.isinf(a_linear),
74
- a_linear,
75
- a_linear.clamp(min=-limit, max=limit),
76
- )
77
- out_gelu = quick_gelu(a_gelu)
78
- return out_gelu * (a_linear + 1)
79
-
80
-
81
-class Phi3SmallMLP(nn.Module):
82
-
83
- def __init__(
84
- self,
85
- config: PretrainedConfig,
86
- quant_config: Optional[QuantizationConfig] = None,
87
- ) -> None:
88
- super().__init__()
89
- self.config = config
90
- assert (self.config.hidden_act == "gegelu"
91
- ), "Only `gegelu` is supported for the 4.7 series of models .."
92
- self.hidden_size = config.hidden_size
93
- self.gegelu_limit = config.gegelu_limit
94
- self.intermediate_size = config.intermediate_size
95
-
96
- self.up_proj = HeadMajorColumnParallelLinear(
97
- self.hidden_size,
98
- 2 * [self.intermediate_size],
99
- bias=True,
100
- quant_config=quant_config,
101
- )
102
- self.down_proj = RowParallelLinear(
103
- self.intermediate_size,
104
- self.hidden_size,
105
- bias=True,
106
- quant_config=quant_config,
107
- )
108
-
109
- def forward(self, x):
110
- gate_up, _ = self.up_proj(x)
111
- x = gegelu(gate_up)
112
- x, _ = self.down_proj(x)
113
- return x
114
-
115
-
116
-class Phi3SmallSelfAttention(nn.Module):
117
-
118
- def __init__(
119
- self,
120
- config: PretrainedConfig,
121
- layer_idx: int,
122
- cache_config: Optional[CacheConfig] = None,
123
- quant_config: Optional[QuantizationConfig] = None,
124
- prefix: str = "",
125
- ) -> None:
126
- super().__init__()
127
- self.layer_idx = layer_idx
128
- self.config = config
129
- self.sparse_block_size = config.blocksparse_block_size
130
- self.homo_heads = config.blocksparse_homo_head_pattern
131
- self.local_blocks = config.blocksparse_num_local_blocks
132
- self.vert_stride = config.blocksparse_vert_stride
133
-
134
- assert (config.blocksparse_block_size ==
135
- config.blocksparse_triton_kernel_block_size)
136
-
137
- self.hidden_size = config.hidden_size
138
- # Number of Query Heads
139
- self.num_heads = config.num_attention_heads
140
-
141
- self.head_dim = self.hidden_size // self.num_heads
142
- self.tp_size = get_tensor_model_parallel_world_size()
143
- # Number of total Key Value Heads before tensor parallel
144
- self.num_key_value_heads = config.num_key_value_heads
145
- self.num_q_per_kv = self.num_heads // self.num_key_value_heads
146
- if self.tp_size > 1:
147
- assert self.num_key_value_heads % self.tp_size == 0
148
- self.num_kv_heads_per_partition = max(
149
- 1, self.num_key_value_heads // self.tp_size)
150
- self.num_heads_per_partition = self.num_heads // self.tp_size
151
-
152
- self.max_position_embeddings = config.max_position_embeddings
153
- self.rope_embedding_base = config.rope_embedding_base
154
- self.rope_position_scale = config.rope_position_scale
155
- self.is_causal = True
156
-
157
- norm_factor = None
158
- if config.mup_use_scaling:
159
- norm_factor = self.head_dim / config.mup_attn_multiplier
160
- else:
161
- norm_factor = math.sqrt(self.head_dim)
162
- self.scale = 1 / norm_factor
163
-
164
- self.query_key_value = HeadMajorQKVParallelLinear(
165
- self.hidden_size,
166
- self.head_dim,
167
- self.num_heads,
168
- self.num_key_value_heads,
169
- bias=True,
170
- quant_config=quant_config,
171
- )
172
-
173
- self.dense = RowParallelLinear(self.hidden_size,
174
- self.hidden_size,
175
- bias=True,
176
- quant_config=quant_config)
177
-
178
- if getattr(self.config, "rope_scaling", None) is not None:
179
- rope_scaling = self.config.rope_scaling
180
- for key in rope_scaling:
181
- if isinstance(rope_scaling[key], list):
182
- rope_scaling[key] = tuple(rope_scaling[key])
183
-
184
- if "factor" not in rope_scaling:
185
- rope_scaling["factor"] = self.rope_position_scale
186
- else:
187
- rope_scaling = {
188
- "rope_type": "linear",
189
- "factor": self.rope_position_scale,
190
- }
191
-
192
- self.rotary_emb = get_rope(
193
- self.head_dim,
194
- rotary_dim=self.head_dim,
195
- max_position=self.max_position_embeddings,
196
- base=self.rope_embedding_base,
197
- rope_scaling=rope_scaling,
198
- )
199
-
200
- # blocksparse params
201
- self.blocksparse_block_size = config.blocksparse_block_size
202
- self.blocksparse_num_local_blocks = config.blocksparse_num_local_blocks
203
- self.blocksparse_vert_stride = config.blocksparse_vert_stride
204
-
205
- use_dense_attn = (getattr(self.config,
206
- "dense_attention_every_n_layers", None)
207
- and (self.layer_idx + 1) %
208
- self.config.dense_attention_every_n_layers == 0)
209
-
210
- bs_params = None
211
- if not use_dense_attn:
212
- bs_params = {
213
- 'max_seqlen': self.max_position_embeddings,
214
- 'num_heads': self.num_heads_per_partition,
215
- "num_kv_heads": self.num_kv_heads_per_partition,
216
- "block_size": self.sparse_block_size,
217
- "local_blocks": self.local_blocks,
218
- "vert_stride": self.vert_stride,
219
- "homo_head": self.homo_heads
220
- }
221
-
222
- self.attn = Attention(self.num_heads_per_partition,
223
- self.head_dim,
224
- self.scale,
225
- num_kv_heads=self.num_kv_heads_per_partition,
226
- cache_config=cache_config,
227
- quant_config=quant_config,
228
- blocksparse_params=bs_params,
229
- prefix=f"{prefix}.attn")
230
-
231
- def forward(
232
- self,
233
- positions: torch.Tensor,
234
- hidden_states: torch.Tensor,
235
- ) -> tuple[torch.Tensor, Optional[torch.Tensor],
236
- Optional[tuple[torch.Tensor]]]:
237
- qkv, _ = self.query_key_value(hidden_states)
238
-
239
- qkv = qkv.view(qkv.shape[:-1] +
240
- (-1, (self.num_q_per_kv + 2), self.head_dim))
241
- q, k, v = qkv.split([self.num_q_per_kv, 1, 1], dim=-2)
242
-
243
- # NOTE: this is required by RotaryEmbed, which indeed does not have to
244
- # TODO: allow 3D QK for rotary forward
245
- q = q.reshape(-1, self.head_dim * self.num_heads_per_partition)
246
- k = k.reshape(-1, self.head_dim * self.num_kv_heads_per_partition)
247
- v = v.reshape(-1, self.head_dim * self.num_kv_heads_per_partition)
248
-
249
- q, k = self.rotary_emb(positions, q, k)
250
- attn_output = self.attn(q, k, v)
251
- output, _ = self.dense(attn_output)
252
-
253
- return output
254
-
255
-
256
-class Phi3SmallDecoderLayer(nn.Module):
257
-
258
- def __init__(
259
- self,
260
- config: PretrainedConfig,
261
- layer_idx: int,
262
- cache_config: Optional[CacheConfig] = None,
263
- quant_config: Optional[QuantizationConfig] = None,
264
- prefix: str = "",
265
- ):
266
- super().__init__()
267
- self.hidden_size = config.hidden_size
268
- self.self_attn = Phi3SmallSelfAttention(config,
269
- layer_idx,
270
- cache_config=cache_config,
271
- quant_config=quant_config,
272
- prefix=f"{prefix}.self_attn")
273
- self.mlp = Phi3SmallMLP(config, quant_config)
274
-
275
- self.input_layernorm = nn.LayerNorm(config.hidden_size,
276
- eps=config.layer_norm_epsilon)
277
- self.post_attention_layernorm = nn.LayerNorm(
278
- config.hidden_size, eps=config.layer_norm_epsilon)
279
-
280
- def forward(
281
- self,
282
- positions: torch.Tensor,
283
- hidden_states: torch.Tensor,
284
- ) -> torch.Tensor:
285
- residual = hidden_states
286
- hidden_states = self.input_layernorm(hidden_states)
287
-
288
- hidden_states = self.self_attn(
289
- positions=positions,
290
- hidden_states=hidden_states,
291
- )
292
- hidden_states = residual + hidden_states
293
-
294
- residual = hidden_states
295
- hidden_states = self.post_attention_layernorm(hidden_states)
296
- hidden_states = self.mlp(hidden_states)
297
- hidden_states = residual + hidden_states
298
- return hidden_states
299
-
300
-
301
-class Phi3SmallModel(nn.Module):
302
-
303
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
304
- super().__init__()
305
-
306
- config = vllm_config.model_config.hf_config
307
- cache_config = vllm_config.cache_config
308
- quant_config = vllm_config.quant_config
309
-
310
- self.config = config
311
- self.embed_tokens = VocabParallelEmbedding(config.vocab_size,
312
- config.hidden_size)
313
- self.mup_embedding_multiplier = config.mup_embedding_multiplier
314
- self.start_layer, self.end_layer, self.layers = make_layers(
315
- config.num_hidden_layers,
316
- lambda prefix: Phi3SmallDecoderLayer(config,
317
- int(prefix.split('.')[-1]),
318
- cache_config,
319
- quant_config,
320
- prefix=prefix),
321
- prefix=f"{prefix}.layers")
322
-
323
- self.final_layernorm = nn.LayerNorm(config.hidden_size,
324
- eps=config.layer_norm_epsilon)
325
- self.make_empty_intermediate_tensors = (
326
- make_empty_intermediate_tensors_factory(["hidden_states"],
327
- config.hidden_size))
328
-
329
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
330
- return self.embed_tokens(input_ids)
331
-
332
- def forward(
333
- self,
334
- input_ids: torch.LongTensor,
335
- positions: Optional[torch.LongTensor],
336
- intermediate_tensors: Optional[IntermediateTensors],
337
- inputs_embeds: Optional[torch.Tensor],
338
- ) -> Union[torch.Tensor, IntermediateTensors]:
339
- if get_pp_group().is_first_rank:
340
- if inputs_embeds is not None:
341
- hidden_states = inputs_embeds
342
- else:
343
- hidden_states = self.get_input_embeddings(input_ids)
344
- if (self.mup_embedding_multiplier is not None
345
- and self.mup_embedding_multiplier > 0.0):
346
- hidden_states = hidden_states * self.mup_embedding_multiplier
347
- else:
348
- assert intermediate_tensors
349
- hidden_states = intermediate_tensors["hidden_states"]
350
- for layer in self.layers[self.start_layer:self.end_layer]:
351
- hidden_states = layer(positions, hidden_states)
352
- if not get_pp_group().is_last_rank:
353
- return IntermediateTensors({"hidden_states": hidden_states})
354
- hidden_states = self.final_layernorm(hidden_states)
355
- return hidden_states
356
-
357
- def load_weights(self, weights: Iterable[tuple[str,
358
- torch.Tensor]]) -> set[str]:
359
- params_dict = dict(self.named_parameters())
360
- loaded_params: set[str] = set()
361
- for name, loaded_weight in weights:
362
- if name.endswith(".bias") and name not in params_dict:
363
- continue
364
- if is_pp_missing_parameter(name, self):
365
- continue
366
- param = params_dict[name]
367
- weight_loader = getattr(param, "weight_loader",
368
- default_weight_loader)
369
- weight_loader(param, loaded_weight)
370
- loaded_params.add(name)
371
- return loaded_params
372
-
373
-
374
-class Phi3SmallForCausalLM(nn.Module, SupportsPP):
375
- _tied_weights_keys = ["lm_head.weight"]
376
-
377
- hf_to_vllm_mapper = WeightsMapper(
378
- orig_to_new_suffix={"rotary_emb.inv_freq": None})
379
-
380
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
381
- super().__init__()
382
- config = vllm_config.model_config.hf_config
383
- quant_config = vllm_config.quant_config
384
- self.config = config
385
- self.quant_config = quant_config
386
- self.model = Phi3SmallModel(vllm_config=vllm_config,
387
- prefix=maybe_prefix(prefix, "model"))
388
- self.vocab_size = config.vocab_size
389
- self.mup_width_multiplier = config.mup_width_multiplier
390
- self.lm_head = ParallelLMHead(
391
- self.vocab_size,
392
- config.hidden_size,
393
- org_num_embeddings=config.vocab_size,
394
- padding_size=DEFAULT_VOCAB_PADDING_SIZE,
395
- quant_config=quant_config,
396
- )
397
- if self.config.tie_word_embeddings:
398
- self.lm_head.weight = self.model.embed_tokens.weight
399
- self.logits_processor = LogitsProcessor(config.vocab_size)
400
- self.make_empty_intermediate_tensors = (
401
- self.model.make_empty_intermediate_tensors)
402
-
403
- # tokens in tiktoken but not used
404
- if hasattr(config, 'dummy_token_indices'):
405
- device = self.lm_head.weight.device
406
- self.register_buffer('dummy_token_indices',
407
- torch.LongTensor(
408
- config.dummy_token_indices).to(device),
409
- persistent=False)
410
- else:
411
- self.dummy_token_indices = None
412
-
413
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
414
- return self.model.get_input_embeddings(input_ids)
415
-
416
- def set_input_embeddings(self, value):
417
- self.model.embed_tokens = value
418
-
419
- def get_output_embeddings(self):
420
- return self.lm_head
421
-
422
- def set_output_embeddings(self, value):
423
- self.lm_head = value
424
-
425
- def set_decoder(self, decoder):
426
- self.model = decoder
427
-
428
- def get_decoder(self):
429
- return self.model
430
-
431
- def compute_logits(
432
- self,
433
- hidden_states: torch.Tensor,
434
- sampling_metadata: SamplingMetadata,
435
- ) -> Optional[torch.Tensor]:
436
- logits = self.logits_processor(self.lm_head, hidden_states,
437
- sampling_metadata)
438
- if self.dummy_token_indices is not None and logits is not None:
439
- logits.index_fill_(-1, self.dummy_token_indices, -torch.inf)
440
- logits = logits / self.mup_width_multiplier
441
- return logits
442
-
443
- def forward(
444
- self,
445
- input_ids: torch.LongTensor,
446
- positions: Optional[torch.LongTensor],
447
- intermediate_tensors: Optional[IntermediateTensors] = None,
448
- inputs_embeds: Optional[torch.Tensor] = None,
449
- ) -> Union[torch.Tensor, IntermediateTensors]:
450
- output_hidden_states = self.model(
451
- input_ids=input_ids,
452
- positions=positions,
453
- intermediate_tensors=intermediate_tensors,
454
- inputs_embeds=inputs_embeds,
455
- )
456
- output_hidden_states = output_hidden_states
457
- return output_hidden_states
458
-
459
- def load_weights(self, weights: Iterable[tuple[str,
460
- torch.Tensor]]) -> set[str]:
461
- loader = AutoWeightsLoader(
462
- self,
463
- skip_prefixes=(["lm_head.weight"]
464
- if self.config.tie_word_embeddings else None))
465
- return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/registry.py
-1
@@ -110,7 +110,6 @@ _TEXT_GENERATION_MODELS = {
110
"PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),
111
"PhiForCausalLM": ("phi", "PhiForCausalLM"),
112
"Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
113
- "Phi3SmallForCausalLM": ("phi3_small", "Phi3SmallForCausalLM"),
113
"PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
114
"Phi4FlashForCausalLM": ("phi4flash", "Phi4FlashForCausalLM"),
115
"Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
vllm/platforms/interface.py
-1
@@ -57,7 +57,6 @@ class _Backend(enum.Enum):
57
PALLAS = enum.auto()
58
PALLAS_VLLM_V1 = enum.auto()
59
IPEX = enum.auto()
60
- BLOCK_SPARSE_FLASH_ATTN = enum.auto()
60
DUAL_CHUNK_FLASH_ATTN = enum.auto()
61
DIFFERENTIAL_FLASH_ATTN = enum.auto()
62
NO_ATTENTION = enum.auto()
vllm/v1/attention/backends/cpu_attn.py
+1
-5
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass
4
-from typing import Any, Optional
4
+from typing import Optional
5
6
import numpy as np
7
import torch
@@ -443,7 +443,6 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
443
alibi_slopes: Optional[list[float]],
444
sliding_window: Optional[int],
445
kv_cache_dtype: str,
446
- blocksparse_params: Optional[dict[str, Any]] = None,
446
logits_soft_cap: Optional[float] = None,
447
attn_type: str = AttentionType.DECODER,
448
kv_sharing_target_layer_name: Optional[str] = None,
@@ -451,9 +450,6 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
450
) -> None:
451
if kv_sharing_target_layer_name is not None:
452
raise NotImplementedError("KV sharing is not supported in V0.")
454
- if blocksparse_params is not None:
455
- raise ValueError(
456
- "Torch SPDA does not support block-sparse attention.")
453
if logits_soft_cap is not None:
454
logger.warning_once("Torch SPDA does not support logits soft cap. "
455
"Outputs may be slightly off.")
vllm/v1/attention/backends/flash_attn.py
+1
-5
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer with FlashAttention."""
4
from dataclasses import dataclass
5
-from typing import Any, ClassVar, Optional
5
+from typing import ClassVar, Optional
6
7
import numpy as np
8
import torch
@@ -349,15 +349,11 @@ class FlashAttentionImpl(AttentionImpl):
349
alibi_slopes: Optional[list[float]],
350
sliding_window: Optional[int],
351
kv_cache_dtype: str,
352
- blocksparse_params: Optional[dict[str, Any]] = None,
352
logits_soft_cap: Optional[float] = None,
353
attn_type: AttentionType = AttentionType.DECODER,
354
kv_sharing_target_layer_name: Optional[str] = None,
355
use_irope: bool = False,
356
) -> None:
358
- if blocksparse_params is not None:
359
- raise ValueError(
360
- "FlashAttention does not support block-sparse attention.")
357
self.num_heads = num_heads
358
self.head_size = head_size
359
self.scale = float(scale)
vllm/v1/attention/backends/flashinfer.py
+1
-2
@@ -4,7 +4,7 @@
4
from __future__ import annotations
5
6
from dataclasses import dataclass
7
-from typing import TYPE_CHECKING, Any, Optional
7
+from typing import TYPE_CHECKING, Optional
8
9
import torch
10
from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,
@@ -490,7 +490,6 @@ class FlashInferImpl(AttentionImpl):
490
alibi_slopes: Optional[list[float]],
491
sliding_window: Optional[int],
492
kv_cache_dtype: str,
493
- blocksparse_params: Optional[dict[str, Any]] = None,
493
logits_soft_cap: Optional[float] = None,
494
attn_type: AttentionType = AttentionType.DECODER,
495
kv_sharing_target_layer_name: Optional[int] = None,
vllm/v1/attention/backends/flex_attention.py
+1
-6
@@ -3,7 +3,7 @@
3
"""Attention layer with FlashAttention."""
4
from collections import defaultdict
5
from dataclasses import dataclass
6
-from typing import Any, Optional
6
+from typing import Optional
7
8
import torch
9
from torch.nn.attention.flex_attention import (BlockMask, _mask_mod_signature,
@@ -342,15 +342,10 @@ class FlexAttentionImpl(AttentionImpl):
342
alibi_slopes: Optional[list[float]],
343
sliding_window: Optional[int],
344
kv_cache_dtype: str,
345
- blocksparse_params: Optional[dict[str, Any]] = None,
345
logits_soft_cap: Optional[float] = None,
346
attn_type: AttentionType = AttentionType.DECODER,
347
kv_sharing_target_layer_name: Optional[str] = None,
348
) -> None:
350
- if blocksparse_params is not None:
351
- # TODO we should support this :think
352
- raise ValueError(
353
- "FlashAttention does not support block-sparse attention.")
349
self.num_heads = num_heads
350
self.head_size = head_size
351
self.scale = float(scale)
vllm/v1/attention/backends/mla/common.py
+1
-2
@@ -190,7 +190,7 @@ return curr_o @ W_O
190
import functools
191
from abc import abstractmethod
192
from dataclasses import dataclass, field
193
-from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar, Union
193
+from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union
194
195
import torch
196
@@ -754,7 +754,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
754
alibi_slopes: Optional[list[float]],
755
sliding_window: Optional[int],
756
kv_cache_dtype: str,
757
- blocksparse_params: Optional[dict[str, Any]],
757
logits_soft_cap: Optional[float],
758
attn_type: str,
759
kv_sharing_target_layer_name: Optional[str],
vllm/v1/attention/backends/mla/cutlass_mla.py
+4
-8
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import os
5
-from typing import Any, Optional
5
+from typing import Optional
6
7
import torch
8
@@ -74,7 +74,6 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
74
alibi_slopes: Optional[list[float]],
75
sliding_window: Optional[int],
76
kv_cache_dtype: str,
77
- blocksparse_params: Optional[dict[str, Any]],
77
logits_soft_cap: Optional[float],
78
attn_type: str,
79
kv_sharing_target_layer_name: Optional[str],
@@ -82,17 +81,14 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
81
**mla_args) -> None:
82
super().__init__(num_heads, head_size, scale, num_kv_heads,
83
alibi_slopes, sliding_window, kv_cache_dtype,
85
- blocksparse_params, logits_soft_cap, attn_type,
84
+ logits_soft_cap, attn_type,
85
kv_sharing_target_layer_name, **mla_args)
86
88
- unsupported_features = [
89
- alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
90
- ]
87
+ unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
88
if any(unsupported_features):
89
raise NotImplementedError(
90
"CutlassMLAImpl does not support one of the following: "
94
- "alibi_slopes, sliding_window, blocksparse_params, "
95
- "logits_soft_cap")
91
+ "alibi_slopes, sliding_window, logits_soft_cap")
92
93
if attn_type != AttentionType.DECODER:
94
raise NotImplementedError("Encoder self-attention and "
vllm/v1/attention/backends/mla/flashmla.py
+4
-8
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from dataclasses import dataclass
5
-from typing import Any, ClassVar, Optional
5
+from typing import ClassVar, Optional
6
7
import torch
8
@@ -119,7 +119,6 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
119
alibi_slopes: Optional[list[float]],
120
sliding_window: Optional[int],
121
kv_cache_dtype: str,
122
- blocksparse_params: Optional[dict[str, Any]],
122
logits_soft_cap: Optional[float],
123
attn_type: str,
124
kv_sharing_target_layer_name: Optional[str],
@@ -127,20 +126,17 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
126
**mla_args) -> None:
127
super().__init__(num_heads, head_size, scale, num_kv_heads,
128
alibi_slopes, sliding_window, kv_cache_dtype,
130
- blocksparse_params, logits_soft_cap, attn_type,
129
+ logits_soft_cap, attn_type,
130
kv_sharing_target_layer_name, **mla_args)
131
132
assert is_flashmla_supported(), \
133
"FlashMLA is not supported on this device"
134
136
- unsupported_features = [
137
- alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
138
- ]
135
+ unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
136
if any(unsupported_features):
137
raise NotImplementedError(
138
"FlashMLAImpl does not support one of the following: "
142
- "alibi_slopes, sliding_window, blocksparse_params, "
143
- "logits_soft_cap")
139
+ "alibi_slopes, sliding_window, logits_soft_cap")
140
141
if attn_type != AttentionType.DECODER:
142
raise NotImplementedError("Encoder self-attention and "
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+4
-8
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from dataclasses import dataclass
5
-from typing import Any, ClassVar, Optional
5
+from typing import ClassVar, Optional
6
7
import torch
8
@@ -167,7 +167,6 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
167
alibi_slopes: Optional[list[float]],
168
sliding_window: Optional[int],
169
kv_cache_dtype: str,
170
- blocksparse_params: Optional[dict[str, Any]],
170
logits_soft_cap: Optional[float],
171
attn_type: str,
172
kv_sharing_target_layer_name: Optional[str],
@@ -175,20 +174,17 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
174
**mla_args) -> None:
175
super().__init__(num_heads, head_size, scale, num_kv_heads,
176
alibi_slopes, sliding_window, kv_cache_dtype,
178
- blocksparse_params, logits_soft_cap, attn_type,
177
+ logits_soft_cap, attn_type,
178
kv_sharing_target_layer_name, **mla_args)
179
assert (num_heads == 16 or num_heads == 128), (
180
f"Aiter MLA only supports 16 or 128 number of heads.\n"
181
f"Provided {num_heads} number of heads.\n"
182
"Try adjusting tensor_parallel_size value.")
184
- unsupported_features = [
185
- alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
186
- ]
183
+ unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
184
if any(unsupported_features):
185
raise NotImplementedError(
186
"Aiter MLA does not support one of the following: "
190
- "alibi_slopes, sliding_window, blocksparse_params, "
191
- "logits_soft_cap")
187
+ "alibi_slopes, sliding_window, logits_soft_cap")
188
189
from aiter import flash_attn_varlen_func
190
self.flash_attn_varlen_func = flash_attn_varlen_func
vllm/v1/attention/backends/mla/triton_mla.py
+4
-8
@@ -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 Any, Optional
4
+from typing import Optional
5
6
import torch
7
@@ -42,7 +42,6 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
42
alibi_slopes: Optional[list[float]],
43
sliding_window: Optional[int],
44
kv_cache_dtype: str,
45
- blocksparse_params: Optional[dict[str, Any]],
45
logits_soft_cap: Optional[float],
46
attn_type: str,
47
kv_sharing_target_layer_name: Optional[str],
@@ -50,17 +49,14 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
49
**mla_args) -> None:
50
super().__init__(num_heads, head_size, scale, num_kv_heads,
51
alibi_slopes, sliding_window, kv_cache_dtype,
53
- blocksparse_params, logits_soft_cap, attn_type,
52
+ logits_soft_cap, attn_type,
53
kv_sharing_target_layer_name, **mla_args)
54
56
- unsupported_features = [
57
- alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
58
- ]
55
+ unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
56
if any(unsupported_features):
57
raise NotImplementedError(
58
"TritonMLAImpl does not support one of the following: "
62
- "alibi_slopes, sliding_window, blocksparse_params, "
63
- "logits_soft_cap")
59
+ "alibi_slopes, sliding_window, logits_soft_cap")
60
61
if attn_type != AttentionType.DECODER:
62
raise NotImplementedError("Encoder self-attention and "
vllm/v1/attention/backends/pallas.py
+1
-7
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from dataclasses import dataclass
5
-from typing import Any, Optional
5
+from typing import Optional
6
7
import torch
8
import torch_xla.core.xla_builder as xb
@@ -132,7 +132,6 @@ class PallasAttentionBackendImpl(AttentionImpl):
132
alibi_slopes: Optional[list[float]],
133
sliding_window: Optional[int],
134
kv_cache_dtype: str,
135
- blocksparse_params: Optional[dict[str, Any]] = None,
135
logits_soft_cap: Optional[float] = None,
136
attn_type: str = AttentionType.DECODER,
137
kv_sharing_target_layer_name: Optional[int] = None,
@@ -142,9 +141,6 @@ class PallasAttentionBackendImpl(AttentionImpl):
141
logger.warning_once(
142
"Using irope in Pallas is not supported yet, it will fall back "
143
"to global attention for long context.")
145
- if blocksparse_params is not None:
146
- raise ValueError("Paged attention Pallas kernel does "
147
- "not support block-sparse attention.")
144
self.num_heads = num_heads
145
self.head_size = head_size
146
self.scale = float(scale)
@@ -158,8 +154,6 @@ class PallasAttentionBackendImpl(AttentionImpl):
154
raise NotImplementedError("Alibi slopes is not supported.")
155
if kv_cache_dtype != "auto":
156
raise NotImplementedError("FP8 KV cache dtype is not supported.")
161
- if blocksparse_params is not None:
162
- raise NotImplementedError("Blocksparse is not supported.")
157
158
if attn_type != AttentionType.DECODER:
159
raise NotImplementedError("Encoder self-attention and "
vllm/v1/attention/backends/rocm_aiter_fa.py
+1
-5
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer with AiterFlashAttention."""
4
from dataclasses import dataclass
5
-from typing import Any, Optional
5
+from typing import Optional
6
7
import torch
8
@@ -334,15 +334,11 @@ class AiterFlashAttentionImpl(AttentionImpl):
334
alibi_slopes: Optional[list[float]],
335
sliding_window: Optional[int],
336
kv_cache_dtype: str,
337
- blocksparse_params: Optional[dict[str, Any]] = None,
337
logits_soft_cap: Optional[float] = None,
338
attn_type: AttentionType = AttentionType.DECODER,
339
kv_sharing_target_layer_name: Optional[int] = None,
340
use_irope: bool = False,
341
) -> None:
343
- if blocksparse_params is not None:
344
- raise ValueError(
345
- "AiterFlashAttention does not support block-sparse attention.")
342
self.num_heads = num_heads
343
self.head_size = head_size
344
self.scale = float(scale)
vllm/v1/attention/backends/triton_attn.py
+1
-5
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer with PagedAttention and Triton prefix prefill."""
4
from dataclasses import dataclass
5
-from typing import Any, ClassVar, Optional
5
+from typing import ClassVar, Optional
6
7
import torch
8
@@ -205,15 +205,11 @@ class TritonAttentionImpl(AttentionImpl):
205
alibi_slopes: Optional[list[float]],
206
sliding_window: Optional[int],
207
kv_cache_dtype: str,
208
- blocksparse_params: Optional[dict[str, Any]] = None,
208
logits_soft_cap: Optional[float] = None,
209
attn_type: AttentionType = AttentionType.DECODER,
210
kv_sharing_target_layer_name: Optional[int] = None,
211
use_irope: bool = False,
212
) -> None:
214
- if blocksparse_params is not None:
215
- raise ValueError(
216
- "TritonAttention does not support block-sparse attention.")
213
self.num_heads = num_heads
214
self.head_size = head_size
215
self.scale = float(scale)