@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 165 lines 5.47 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """ROCm Kimi-K3 specialization of GDN attention metadata.
4
5 The request classification and cudagraph staging intentionally mirror
6 ``GDNAttentionMetadataBuilder``. Only the FLA chunk metadata is built
7 differently on device rather than on the host.
8 """
9
10 from dataclasses import replace
11
12 import torch
13
14 from vllm.logger import init_logger
15 from vllm.model_executor.layers.mamba.checkpoint import (
16 MambaPrefillCheckpointMetadata,
17 )
18 from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
19 from vllm.triton_utils import tl, triton
20 from vllm.utils.math_utils import next_power_of_2
21 from vllm.v1.attention.backends.gdn_attn import (
22 GDNAttentionBackend,
23 GDNAttentionMetadata,
24 GDNAttentionMetadataBuilder,
25 )
26 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
27
28 logger = init_logger(__name__)
29
30 _BLOCK_T = 256
31 _MIN_BLOCK_N = 128
32
33
34 @triton.jit(do_not_specialize=["N"])
35 def _chunk_metadata_kernel(
36 cu_seqlens,
37 chunk_indices,
38 chunk_offsets,
39 N,
40 BT: tl.constexpr,
41 BLOCK_N: tl.constexpr,
42 BLOCK_T: tl.constexpr,
43 ):
44 i_n = tl.program_id(0)
45 offs_n = tl.arange(0, BLOCK_N)
46 is_seq = offs_n < N
47 bos = tl.load(cu_seqlens + offs_n, mask=is_seq, other=0).to(tl.int32)
48 eos = tl.load(cu_seqlens + offs_n + 1, mask=is_seq, other=0).to(tl.int32)
49 nt = tl.where(is_seq, tl.cdiv(eos - bos, BT), 0)
50
51 base = tl.sum(tl.where(offs_n < i_n, nt, 0))
52 num_chunks = tl.sum(tl.where(offs_n == i_n, nt, 0))
53
54 tl.store(chunk_offsets + i_n, base)
55 if i_n == 0:
56 tl.store(chunk_offsets + N, tl.sum(nt))
57
58 for t0 in range(0, num_chunks, BLOCK_T):
59 offs_t = t0 + tl.arange(0, BLOCK_T)
60 mask_t = offs_t < num_chunks
61 row = (base + offs_t) * 2
62 tl.store(chunk_indices + row, tl.full([BLOCK_T], i_n, tl.int32), mask=mask_t)
63 tl.store(chunk_indices + row + 1, offs_t.to(tl.int32), mask=mask_t)
64
65
66 def prepare_chunk_metadata_device(
67 cu_seqlens: torch.Tensor,
68 cu_seqlens_cpu: torch.Tensor,
69 chunk_size: int,
70 ) -> tuple[torch.Tensor, torch.Tensor]:
71 """Build FLA chunk metadata on device, with no host<->device transfer."""
72 num_seqs = cu_seqlens_cpu.numel() - 1
73 seq_lens = cu_seqlens_cpu[1:] - cu_seqlens_cpu[:-1]
74 num_chunks = int(((seq_lens + chunk_size - 1) // chunk_size).sum())
75
76 chunk_indices = torch.empty(
77 num_chunks, 2, dtype=cu_seqlens.dtype, device=cu_seqlens.device
78 )
79 chunk_offsets = torch.empty(
80 num_seqs + 1, dtype=torch.int64, device=cu_seqlens.device
81 )
82 _chunk_metadata_kernel[(num_seqs,)](
83 cu_seqlens,
84 chunk_indices,
85 chunk_offsets,
86 num_seqs,
87 BT=chunk_size,
88 BLOCK_N=max(_MIN_BLOCK_N, next_power_of_2(num_seqs + 1)),
89 BLOCK_T=_BLOCK_T,
90 num_warps=4,
91 )
92 return chunk_indices, chunk_offsets
93
94
95 class KimiK3ROCmKDAMetadataBuilder(GDNAttentionMetadataBuilder):
96 def build( # type: ignore[override]
97 self,
98 common_prefix_len: int,
99 common_attn_metadata: CommonAttentionMetadata,
100 num_accepted_tokens: torch.Tensor | None = None,
101 num_decode_draft_tokens_cpu: torch.Tensor | None = None,
102 fast_build: bool = False,
103 ) -> GDNAttentionMetadata:
104 attn_metadata = super().build(
105 common_prefix_len,
106 common_attn_metadata,
107 num_accepted_tokens,
108 num_decode_draft_tokens_cpu,
109 fast_build,
110 )
111 if attn_metadata.checkpoint is None:
112 return attn_metadata
113 # The shared builder keys checkpoint rows on every request; the ROCm
114 # prefill path indexes them by prefill_query_start_loc, so rebuild with
115 # those rows.
116 return replace(
117 attn_metadata,
118 checkpoint=self._build_checkpoint_metadata(
119 common_attn_metadata, attn_metadata, num_decode_draft_tokens_cpu
120 ),
121 )
122
123 def _build_checkpoint_metadata(
124 self,
125 m: CommonAttentionMetadata,
126 attn_metadata: GDNAttentionMetadata,
127 num_decode_draft_tokens_cpu: torch.Tensor | None,
128 ) -> MambaPrefillCheckpointMetadata | None:
129 if attn_metadata.num_prefills == 0:
130 return None
131 # request_rows must line up with prefill_query_start_loc, either the
132 # prefill tail of a decode-first batch, or every non-spec row
133 # including padding
134 num_decodes = attn_metadata.num_decodes
135 request_rows = list(
136 range(num_decodes, num_decodes + attn_metadata.num_prefills)
137 )
138 if attn_metadata.spec_sequence_masks is not None:
139 assert num_decode_draft_tokens_cpu is not None
140 request_rows = (
141 (num_decode_draft_tokens_cpu < 0).nonzero().flatten().tolist()
142 )
143 return self.checkpoint_builder.build(m, request_rows)
144
145 def _build_chunk_metadata(
146 self,
147 prefill_query_start_loc: torch.Tensor,
148 prefill_query_start_loc_cpu: torch.Tensor,
149 device: torch.device,
150 ) -> tuple[torch.Tensor, torch.Tensor]:
151 return prepare_chunk_metadata_device(
152 prefill_query_start_loc,
153 prefill_query_start_loc_cpu,
154 FLA_CHUNK_SIZE,
155 )
156
157
158 class KimiK3ROCmKDABackend(GDNAttentionBackend):
159 @staticmethod
160 def get_name() -> str:
161 return "KIMI_K3_KDA_ROCM"
162
163 @staticmethod
164 def get_builder_cls() -> type[KimiK3ROCmKDAMetadataBuilder]:
165 return KimiK3ROCmKDAMetadataBuilder