@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 195 lines 5.51 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Startup compilation of Kimi-K3 vision FA4 kernels."""
4
5 from __future__ import annotations
6
7 import math
8 from collections.abc import Iterator
9 from dataclasses import dataclass
10 from functools import partial
11
12 import torch
13
14 from vllm.model_executor.warmup.cutedsl_warmup import (
15 CuTeDSLCompileUnit,
16 register_cutedsl_warmup_provider,
17 )
18 from vllm.platforms import current_platform
19
20 _FA4_TILE_SIZE = 128
21 _FA4_MAX_SPLITS = 128
22
23
24 @dataclass(frozen=True)
25 class KimiK3VisionFA4WarmupConfig:
26 num_heads: int
27 head_dim: int
28 dtype: torch.dtype
29 max_batch_size: int
30 max_seqlen: int
31
32
33 @dataclass(frozen=True)
34 class _FA4WarmupProbe:
35 batch_size: int
36 max_seqlen: int
37 dispatch_key: tuple[int, bool, int | None]
38
39
40 def _combine_log_max_splits(head_dim: int, num_splits: int) -> int:
41 k_block_size = 64 if head_dim <= 64 else 128
42 tile_m = 8 if k_block_size == 128 else 16
43 min_log_max_splits = 5 if tile_m == 8 else 4
44 return max(math.ceil(math.log2(num_splits)), min_log_max_splits)
45
46
47 def _get_dispatch(
48 config: KimiK3VisionFA4WarmupConfig,
49 *,
50 batch_size: int,
51 max_seqlen: int,
52 num_sms: int,
53 ) -> tuple[int, bool, int | None]:
54 # These are the shape-dependent fields in FA4's forward and combine
55 # compile keys for noncausal varlen MHA. Compile units still pass
56 # num_splits=0 so the production selector makes the final decision.
57 q_stage = 2 if max_seqlen > _FA4_TILE_SIZE else 1
58 effective_q_tile = q_stage * _FA4_TILE_SIZE
59 num_m_blocks = math.ceil(max_seqlen / effective_q_tile)
60 num_n_blocks = math.ceil(max_seqlen / _FA4_TILE_SIZE)
61
62 if num_n_blocks <= 4:
63 num_splits = 1
64 else:
65 total_m_blocks = batch_size * config.num_heads * num_m_blocks
66 num_splits = min(
67 num_sms // total_m_blocks,
68 _FA4_MAX_SPLITS,
69 num_n_blocks,
70 )
71
72 is_split_kv = num_splits > 1
73 combine_key = (
74 _combine_log_max_splits(config.head_dim, num_splits) if is_split_kv else None
75 )
76 return q_stage, is_split_kv, combine_key
77
78
79 def _get_warmup_probes(
80 config: KimiK3VisionFA4WarmupConfig,
81 *,
82 num_sms: int,
83 ) -> tuple[_FA4WarmupProbe, ...]:
84 if config.max_batch_size <= 0 or config.max_seqlen <= 0:
85 return ()
86
87 probes: dict[tuple[int, bool, int | None], _FA4WarmupProbe] = {}
88
89 def add_probe(batch_size: int, max_seqlen: int) -> None:
90 dispatch_key = _get_dispatch(
91 config,
92 batch_size=batch_size,
93 max_seqlen=max_seqlen,
94 num_sms=num_sms,
95 )
96 if dispatch_key not in probes:
97 probes[dispatch_key] = _FA4WarmupProbe(
98 batch_size=batch_size,
99 max_seqlen=max_seqlen,
100 dispatch_key=dispatch_key,
101 )
102
103 add_probe(batch_size=1, max_seqlen=1)
104 if config.max_seqlen > _FA4_TILE_SIZE:
105 add_probe(batch_size=1, max_seqlen=_FA4_TILE_SIZE + 1)
106
107 max_n_blocks = math.ceil(config.max_seqlen / _FA4_TILE_SIZE)
108 for num_n_blocks in range(5, max_n_blocks + 1):
109 max_seqlen = min(
110 (num_n_blocks - 1) * _FA4_TILE_SIZE + 1,
111 config.max_seqlen,
112 )
113 num_m_blocks = math.ceil(max_seqlen / (2 * _FA4_TILE_SIZE))
114 max_split_batch_size = min(
115 config.max_batch_size,
116 num_sms // (2 * config.num_heads * num_m_blocks),
117 )
118 if max_split_batch_size == 0:
119 break
120 for batch_size in range(1, max_split_batch_size + 1):
121 add_probe(batch_size, max_seqlen)
122
123 return tuple(probes.values())
124
125
126 def _iter_compile_units(
127 config: KimiK3VisionFA4WarmupConfig,
128 ) -> Iterator[CuTeDSLCompileUnit]:
129 device_id = current_platform.current_device()
130 num_sms = current_platform.num_compute_units(device_id)
131 for probe in _get_warmup_probes(config, num_sms=num_sms):
132 yield CuTeDSLCompileUnit(
133 name="kimi_k3_vision_fa4",
134 key=("kimi_k3_vision_fa4", config, probe.dispatch_key),
135 compile=partial(_compile, config, probe),
136 )
137
138
139 def _compile(
140 config: KimiK3VisionFA4WarmupConfig,
141 probe: _FA4WarmupProbe,
142 ) -> None:
143 from vllm.vllm_flash_attn import flash_attn_varlen_func
144
145 device = current_platform.current_device()
146 total_tokens = probe.batch_size * probe.max_seqlen
147 qkv = torch.empty(
148 3,
149 total_tokens,
150 config.num_heads,
151 config.head_dim,
152 device=device,
153 dtype=config.dtype,
154 )
155 cu_seqlens = torch.arange(
156 0,
157 total_tokens + 1,
158 probe.max_seqlen,
159 device=qkv.device,
160 dtype=torch.int32,
161 )
162 flash_attn_varlen_func(
163 qkv[0],
164 qkv[1],
165 qkv[2],
166 cu_seqlens_q=cu_seqlens,
167 cu_seqlens_k=cu_seqlens,
168 max_seqlen_q=probe.max_seqlen,
169 max_seqlen_k=probe.max_seqlen,
170 dropout_p=0.0,
171 causal=False,
172 softmax_scale=config.head_dim**-0.5,
173 fa_version=4,
174 num_splits=0,
175 )
176
177
178 class _WarmupProvider:
179 def __init__(self) -> None:
180 self.configs: set[KimiK3VisionFA4WarmupConfig] = set()
181
182 def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:
183 return tuple(
184 unit for config in self.configs for unit in _iter_compile_units(config)
185 )
186
187
188 _PROVIDER = _WarmupProvider()
189
190
191 def register_kimi_k3_vision_fa4_warmup(
192 config: KimiK3VisionFA4WarmupConfig,
193 ) -> None:
194 _PROVIDER.configs.add(config)
195 register_cutedsl_warmup_provider(_PROVIDER)