eec3c36a47 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -8,11 +8,7 @@ from typing import ClassVar

8 import torch

9

10 from vllm._aiter_ops import rocm_aiter_ops

11 -from vllm.config import (

12 - VllmConfig,

13 - get_current_vllm_config_or_none,

14 - get_layers_from_vllm_config,

15 -)

11 +from vllm.config import VllmConfig, get_layers_from_vllm_config

12 from vllm.config.cache import CacheDType

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.attention import Attention

@@ -42,43 +38,9 @@ from vllm.v1.attention.backends.utils import (

38 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states

39 from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout

40

45 -_PA_GLUON_MAX_QUERY_LEN = 4

46 -_PA_GLUON_MAX_QUERY_GROUP_SIZE = 64

47 -# Query group sizes the gluon paged-attention decode kernel is validated for: 8, 16

48 -_PA_GLUON_QUERY_GROUP_SIZES = (8, 16)

49 -

50 -# The kernel is only validated for this head size and kernel block size.

51 -_PA_GLUON_HEAD_SIZE = 128

52 -_PA_GLUON_BLOCK_SIZE = 128

53 -

54 -

55 -def _pa_gluon_supports(num_heads_q: int, num_heads_kv: int, head_size: int) -> bool:

56 - """Whether the head config can use the PA decode gluon kernel.

57 -

58 - Requires the shuffle KV cache layout (the kernel reads K/V in that layout)

59 - plus a head config the kernel is validated for. Both the advertised kernel

60 - block sizes and the decode dispatch go through this, so a config can never

61 - be offered a 128-token page that gluon will then decline to serve.

62 - """

63 - return (

64 - rocm_aiter_ops.is_shuffle_kv_cache_enabled()

65 - and num_heads_kv > 0

66 - and num_heads_q % num_heads_kv == 0

67 - and num_heads_q // num_heads_kv in _PA_GLUON_QUERY_GROUP_SIZES

68 - and head_size == _PA_GLUON_HEAD_SIZE

69 - )

70 -

71 -

41 _PARTITION_SIZE_ROCM = 256

42 _CP_TOKENS_PER_ITER_ROCM = 32 * 1024

43 if current_platform.is_rocm():

75 - from aiter.ops.triton.gluon.pa_decode_gluon import (

76 - get_recommended_splits,

77 - )

78 - from aiter.ops.triton.gluon.pa_decode_gluon import (

79 - pa_decode_gluon as _pa_decode_gluon,

80 - )

81 -

44 from vllm.triton_utils import tl, triton

45

46 def block_size(x, head_dim):

@@ -168,8 +130,10 @@ if current_platform.is_rocm():

130 # for kv cache layout as

131 # K: [num_blocks, num_head, head_dim // x, page_size, x]

132 # V: [num_blocks, num_head, page_size // x, head_dim, x]

171 - # The sides only sit a dense page apart when they are separate

172 - # planes (H outermost), so take the block stride from the tensor.

133 + # Take the block stride from the tensor, as the NHD branch does:

134 + # the two sides only sit a dense page apart when they are separate

135 + # planes, and models that pin a block-compact layout interleave

136 + # them instead.

137 key_cache_ptr_offset = (

138 key_cache_ptr

139 + block_id * k_cache_stride0

@@ -363,7 +327,6 @@ logger = init_logger(__name__)

327 @dataclass

328 class AiterFlashAttentionDecodeMetadata:

329 max_query_len: int

366 - uniform_query_len: int | None

330

331

332 @dataclass

@@ -598,15 +561,8 @@ class AiterFlashAttentionMetadataBuilder(

561

562 decode_metadata = None

563 if num_decodes > 0:

601 - decode_max_query_len = query_lens_cpu[:num_decodes].max().item()

602 - uniform_query_len = (

603 - decode_max_query_len

604 - if num_decode_tokens == num_decodes * decode_max_query_len

605 - else None

606 - )

564 decode_metadata = AiterFlashAttentionDecodeMetadata(

608 - max_query_len=decode_max_query_len,

609 - uniform_query_len=uniform_query_len,

565 + max_query_len=query_lens_cpu[:num_decodes].max().item(),

566 )

567

568 prefill_metadata = None

@@ -778,27 +734,9 @@ class AiterFlashAttentionMetadataBuilder(

734 """

735 num_reqs = common_attn_metadata.num_reqs

736 num_tokens = common_attn_metadata.num_actual_tokens

781 - max_query_len = common_attn_metadata.max_query_len

782 -

783 - # Uniform-decode assumption does not hold for the

784 - # drafter's first forward after a target step: it inherits the target's

785 - # per-request query lengths, so rows can be longer than gluon's limit or

786 - # ragged. Those batches need the real split, which costs a sync.

787 - # _PA_GLUON_MAX_QUERY_LEN only binds when gluon is the decode consumer,

788 - # so test that rather than the shuffle layout alone.

789 - if _pa_gluon_supports(self.num_heads_q, self.num_heads_kv, self.headdim) and (

790 - max_query_len > _PA_GLUON_MAX_QUERY_LEN

791 - or num_tokens != num_reqs * max_query_len

792 - ):

793 - return self.build(

794 - common_prefix_len=0, common_attn_metadata=common_attn_metadata

795 - )

737

738 decode_metadata = AiterFlashAttentionDecodeMetadata(

798 - max_query_len=max_query_len,

799 - uniform_query_len=(

800 - max_query_len if num_tokens == num_reqs * max_query_len else None

801 - ),

739 + max_query_len=common_attn_metadata.max_query_len,

740 )

741

742 return AiterFlashAttentionMetadata(

@@ -853,21 +791,6 @@ class AiterFlashAttentionBackend(AttentionBackend):

791

792 @staticmethod

793 def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:

856 - if not rocm_aiter_ops.is_shuffle_kv_cache_enabled():

857 - return [16, 32]

858 - # Only gluon serves 128-token pages; the pa_fwd_asm/ll4mi decode

859 - # fallback is limited to 16 and 32. Advertise 128 only when gluon can

860 - # run so selection never picks a page we cannot serve.

861 - vllm_config = get_current_vllm_config_or_none()

862 - if vllm_config is not None and vllm_config.model_config is not None:

863 - mc = vllm_config.model_config

864 - pc = vllm_config.parallel_config

865 - if _pa_gluon_supports(

866 - mc.get_num_attention_heads(pc),

867 - mc.get_num_kv_heads(pc),

868 - mc.get_head_size(),

869 - ):

870 - return [16, 32, 128]

794 return [16, 32]

795

796 @classmethod

@@ -924,6 +847,14 @@ class AiterFlashAttentionBackend(AttentionBackend):

847

848 @classmethod

849 def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:

850 + if rocm_aiter_ops.is_shuffle_kv_cache_enabled():

851 + # pa_fwd_asm strides between pages by a whole dense page, so the two

852 + # head slots customize_spec publishes have to be separate planes

853 + # spanning every block: H outermost. Under LBHNC the slots alternate

854 + # block by block, doubling the stride the kernel assumes. LBHNC

855 + # stays listed for models whose mixed HNC shapes need a

856 + # block-compact layout; those read K/V by stride instead.

857 + return (KVCacheLayout.LHBNC, KVCacheLayout.LBHNC)

858 # K and V come out of the content dim as transposed views rather than

859 # copies, so the head dim may sit on either side of the block dim, but

860 # the layer must stay outermost.

@@ -1358,25 +1289,11 @@ class AiterFlashAttentionImpl(AttentionImpl):

1289 if num_decodes > 0:

1290 assert attn_metadata.decode_metadata is not None

1291 decode_max_query_len = attn_metadata.decode_metadata.max_query_len

1361 - decode_query_len = attn_metadata.decode_metadata.uniform_query_len

1292

1363 - # check if we can use the gluon paged-attention decode kernel

1364 - use_gluon = (

1365 - _pa_gluon_supports(

1366 - self.num_heads, self.num_kv_heads, self.head_size

1367 - )

1368 - and key_cache.shape[1] == _PA_GLUON_BLOCK_SIZE

1369 - and decode_query_len is not None

1370 - and decode_query_len <= _PA_GLUON_MAX_QUERY_LEN

1371 - and decode_query_len * (self.num_heads // self.num_kv_heads)

1372 - <= _PA_GLUON_MAX_QUERY_GROUP_SIZE

1373 - and (decode_query_len == 1 or attn_metadata.causal)

1374 - )

1293 # Use unified_attention for the decodes the paged kernels can't

1376 - # take: sliding window, sinks, or a multi-token batch that gluon

1377 - # declined (pa_fwd_asm and paged_attention_v1 don't support

1378 - # sinks).

1379 - if not use_gluon and (

1294 + # take: sliding window, sinks, or a multi-token batch

1295 + # (pa_fwd_asm and paged_attention_v1 don't support sinks).

1296 + if (

1297 self.sliding_window[0] != -1

1298 or decode_max_query_len > 1

1299 or self.sinks is not None

@@ -1505,145 +1422,62 @@ class AiterFlashAttentionImpl(AttentionImpl):

1422 num_blocks, num_kv_heads, block_size // x, head_size, x

1423 )

1424

1508 - if use_gluon:

1509 - is_fp8_kv = is_quantized_kv_cache(self.kv_cache_dtype)

1510 - # Per-tensor descale, as a float32 [1] tensor.

1511 - k_scale_gluon = (

1512 - layer._k_scale.reshape(1).to(torch.float32)

1513 - if is_fp8_kv

1514 - else None

1515 - )

1516 - v_scale_gluon = (

1517 - layer._v_scale.reshape(1).to(torch.float32)

1518 - if is_fp8_kv

1519 - else None

1520 - )

1521 - compute_type = (

1522 - current_platform.fp8_dtype() if is_fp8_kv else query.dtype

1523 - )

1524 - # The kernel folds the query positions into the group

1525 - # dim, so the intermediate buffers are sized by the

1526 - # combined extent.

1527 - query_group_size = decode_query_len * (

1528 - num_heads // num_kv_heads

1529 - )

1530 -

1531 - sliding_window_int = (

1532 - self.sliding_window[0] + 1

1533 - if self.sliding_window[0] > 0

1534 - else 0

1535 - )

1536 - if sliding_window_int > 0:

1537 - max_context_partition_num = 1

1538 - context_partition_size = 128

1539 - else:

1540 - max_context_partition_num = get_recommended_splits(

1541 - num_decodes, num_kv_heads

1542 - )

1543 - context_partition_size = _PARTITION_SIZE_ROCM

1544 -

1545 - intermediate_shape = (

1546 - num_decodes,

1547 - num_kv_heads,

1548 - max_context_partition_num,

1549 - query_group_size,

1550 - )

1551 - exp_sums = torch.empty(

1552 - intermediate_shape,

1553 - dtype=torch.float32,

1554 - device=query.device,

1555 - )

1556 - max_logits = torch.empty_like(exp_sums)

1557 - temporary_output = torch.empty(

1558 - (*intermediate_shape, head_size),

1559 - dtype=output.dtype,

1560 - device=query.device,

1561 - )

1562 -

1563 - _pa_decode_gluon(

1564 - output=output[:num_decode_tokens],

1565 - query=query[:num_decode_tokens],

1566 - key_cache=new_key_cache,

1567 - value_cache=new_value_cache,

1568 - context_lengths=attn_metadata.seq_lens[:num_decodes].to(

1569 - torch.int32

1570 - ),

1571 - block_tables=attn_metadata.block_table[:num_decodes].to(

1572 - torch.int32

1573 - ),

1574 - softmax_scale=self.scale,

1575 - query_length=decode_query_len,

1576 - max_context_partition_num=max_context_partition_num,

1577 - context_partition_size=context_partition_size,

1578 - compute_type=compute_type,

1579 - query_scale=None,

1580 - key_scale=k_scale_gluon,

1581 - value_scale=v_scale_gluon,

1582 - exp_sums=exp_sums,

1583 - max_logits=max_logits,

1584 - temporary_output=temporary_output,

1585 - alibi_slopes=self.alibi_slopes,

1586 - sinks=self.sinks,

1587 - sliding_window=sliding_window_int,

1588 - ps=True,

1589 - )

1590 - else:

1591 - # Unlike gluon, this kernel derives block addresses

1592 - # arithmetically, so it only reads the right bytes when

1593 - # each side is a dense plane of pages.

1594 - assert new_key_cache.stride(0) == (

1595 - num_kv_heads * head_size * block_size

1596 - ), (

1597 - "paged_attention_common needs the K/V sides as "

1598 - "dense planes; the resolved KV cache layout "

1599 - "interleaves them within a block."

1600 - )

1601 - num_seqs = attn_metadata.seq_lens.shape[0]

1602 - max_num_partitions = (

1603 - attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1

1604 - ) // _PARTITION_SIZE_ROCM

1605 - tmp_out = torch.empty(

1606 - (num_seqs, num_heads, max_num_partitions, head_size),

1607 - dtype=query.dtype,

1608 - device=query.device,

1609 - )

1610 - exp_sums = torch.empty(

1611 - (num_seqs, num_heads, max_num_partitions),

1612 - dtype=torch.float32,

1613 - device=query.device,

1614 - )

1615 - max_logits = torch.empty_like(exp_sums)

1616 - k_qscale = (

1617 - layer._k_scale

1618 - if attn_metadata.k_scale is None

1619 - else attn_metadata.k_scale

1620 - )

1621 - v_qscale = (

1622 - layer._v_scale

1623 - if attn_metadata.v_scale is None

1624 - else attn_metadata.v_scale

1625 - )

1626 - rocm_aiter_ops.paged_attention_common(

1627 - Q=query[:num_decode_tokens],

1628 - K=new_key_cache,

1629 - V=new_value_cache,

1630 - tmp_out=tmp_out,

1631 - max_logits=max_logits,

1632 - exp_sums=exp_sums,

1633 - max_seq_len=attn_metadata.max_seq_len,

1634 - block_tables=attn_metadata.block_table[:num_decodes],

1635 - context_lens=attn_metadata.seq_lens[:num_decodes],

1636 - block_tables_stride0=attn_metadata.block_table[

1637 - :num_decodes

1638 - ].stride(0),

1639 - scale=self.scale,

1640 - K_QScale_hip=k_qscale,

1641 - V_QScale_hip=v_qscale,

1642 - K_QScale_asm=k_qscale,

1643 - V_QScale_asm=v_qscale,

1644 - out_=output[:num_decode_tokens],

1645 - kv_cache_dtype=self.kv_cache_dtype,

1646 - )

1425 + # This kernel derives block addresses arithmetically, so it

1426 + # only reads the right bytes when each side is a dense

1427 + # plane of pages.

1428 + assert new_key_cache.stride(0) == (

1429 + num_kv_heads * head_size * block_size

1430 + ), (

1431 + "paged_attention_common needs the K/V sides as dense "

1432 + "planes; the resolved KV cache layout interleaves them "

1433 + "within a block."

1434 + )

1435 + num_seqs = attn_metadata.seq_lens.shape[0]

1436 + max_num_partitions = (

1437 + attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1

1438 + ) // _PARTITION_SIZE_ROCM

1439 + tmp_out = torch.empty(

1440 + (num_seqs, num_heads, max_num_partitions, head_size),

1441 + dtype=query.dtype,

1442 + device=query.device,

1443 + )

1444 + exp_sums = torch.empty(

1445 + (num_seqs, num_heads, max_num_partitions),

1446 + dtype=torch.float32,

1447 + device=query.device,

1448 + )

1449 + max_logits = torch.empty_like(exp_sums)

1450 + k_qscale = (

1451 + layer._k_scale

1452 + if attn_metadata.k_scale is None

1453 + else attn_metadata.k_scale

1454 + )

1455 + v_qscale = (

1456 + layer._v_scale

1457 + if attn_metadata.v_scale is None

1458 + else attn_metadata.v_scale

1459 + )

1460 + rocm_aiter_ops.paged_attention_common(

1461 + Q=query[:num_decode_tokens],

1462 + K=new_key_cache,

1463 + V=new_value_cache,

1464 + tmp_out=tmp_out,

1465 + max_logits=max_logits,

1466 + exp_sums=exp_sums,

1467 + max_seq_len=attn_metadata.max_seq_len,

1468 + block_tables=attn_metadata.block_table[:num_decodes],

1469 + context_lens=attn_metadata.seq_lens[:num_decodes],

1470 + block_tables_stride0=attn_metadata.block_table[

1471 + :num_decodes

1472 + ].stride(0),

1473 + scale=self.scale,

1474 + K_QScale_hip=k_qscale,

1475 + V_QScale_hip=v_qscale,

1476 + K_QScale_asm=k_qscale,

1477 + V_QScale_asm=v_qscale,

1478 + out_=output[:num_decode_tokens],

1479 + kv_cache_dtype=self.kv_cache_dtype,

1480 + )

1481 else:

1482 _, num_heads, head_size = query.shape

1483 nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8

[ROCm][BugFix] Revert AITER PA gluon decode from ROCM_AITER_FA (#54805)