@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2c11a738b3

Fork 0

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, 02:15 UTC

[Model] New model support for microsoft/Phi-4-mini-flash-reasoning (#20702)

Signed-off-by: Congcong Chen <congcongchen@microsoft.com>

Congcong Chen committed Jul 12, 2025 at 06:02 UTC 2c11a738b35e6e65731c5ad7774581ca2c33dc6a

22 files changed +1869 -41

csrc/mamba/mamba_ssm/selective_scan_fwd.cu

+25 -24

@@ -312,19 +312,20 @@ void selective_scan_fwd_launch(SSMParamsBase &params, cudaStream_t stream) {

312 // kIsVariableB, kIsVariableC and kHasZ are all set to True to reduce binary size

313 constexpr bool kIsVariableB = true;

314 constexpr bool kIsVariableC = true;

315 - constexpr bool kHasZ = true;

315 BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {

317 - BOOL_SWITCH(params.query_start_loc_ptr != nullptr , kVarlen, [&] {

318 - using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kNRows, kIsEvenLen, kIsVariableB, kIsVariableC, kHasZ, kVarlen, input_t, weight_t>;

319 - constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t);

320 - dim3 grid(params.batch, params.dim / kNRows);

321 - auto kernel = &selective_scan_fwd_kernel<Ktraits>;

322 - if (kSmemSize >= 48 * 1024) {

323 - C10_CUDA_CHECK(cudaFuncSetAttribute(

324 - (void *) kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));

325 - }

326 - kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);

327 - C10_CUDA_KERNEL_LAUNCH_CHECK();

316 + BOOL_SWITCH(params.z_ptr != nullptr , kHasZ, [&] {

317 + BOOL_SWITCH(params.query_start_loc_ptr != nullptr , kVarlen, [&] {

318 + using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kNRows, kIsEvenLen, kIsVariableB, kIsVariableC, kHasZ, kVarlen, input_t, weight_t>;

319 + constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t);

320 + dim3 grid(params.batch, params.dim / kNRows);

321 + auto kernel = &selective_scan_fwd_kernel<Ktraits>;

322 + if (kSmemSize >= 48 * 1024) {

323 + C10_CUDA_CHECK(cudaFuncSetAttribute(

324 + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));

325 + }

326 + kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);

327 + C10_CUDA_KERNEL_LAUNCH_CHECK();

328 + });

329 });

330 });

331 }

@@ -612,19 +613,20 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,

613

614 at::Tensor z, out_z;

615 const bool has_z = z_.has_value();

615 - TORCH_CHECK(has_z, "has_z = False is disabled in favor of reduced binary size")

616 - z = z_.value();

617 - TORCH_CHECK(z.scalar_type() == input_type);

618 - TORCH_CHECK(z.is_cuda());

619 - TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);

620 - if (varlen){

621 - CHECK_SHAPE(z, dim, seqlen);

622 - } else {

623 - CHECK_SHAPE(z, batch_size, dim, seqlen);

616 + if (has_z) {

617 + z = z_.value();

618 + TORCH_CHECK(z.scalar_type() == input_type);

619 + TORCH_CHECK(z.is_cuda());

620 + TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);

621 + if (varlen){

622 + CHECK_SHAPE(z, dim, seqlen);

623 + } else {

624 + CHECK_SHAPE(z, batch_size, dim, seqlen);

625 + }

626 +

627 + out_z = z;

628 }

629

626 - out_z = z;

627 -

630 // Right now u has BHL layout and delta has HBL layout, and we want out to have HBL layout

631 at::Tensor out = delta;

632 TORCH_CHECK(ssm_states.scalar_type() == input_type);

@@ -653,4 +655,3 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,

655 selective_scan_fwd_cuda<input_t, weight_t>(params, stream);

656 });

657 }

656 -

docs/models/supported_models.md

@@ -374,6 +374,7 @@ Specified using `--task generate`.

374 | `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. | ✅︎ | ✅︎ | ✅︎ |

375 | `Phi3SmallForCausalLM` | Phi-3-Small | `microsoft/Phi-3-small-8k-instruct`, `microsoft/Phi-3-small-128k-instruct`, etc. | | ✅︎ | ✅︎ |

376 | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |

378 | `PersimmonForCausalLM` | Persimmon | `adept/persimmon-8b-base`, `adept/persimmon-8b-chat`, etc. | | ✅︎ | ✅︎ |

380 | `QWenLMHeadModel` | Qwen | `Qwen/Qwen-7B`, `Qwen/Qwen-7B-Chat`, etc. | ✅︎ | ✅︎ | ✅︎ |

tests/models/registry.py

@@ -248,6 +248,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

248 "Phi3SmallForCausalLM": _HfExamplesInfo("microsoft/Phi-3-small-8k-instruct",

249 trust_remote_code=True,

250 v0_only=True),

251 + "Phi4FlashForCausalLM": _HfExamplesInfo("microsoft/Phi-4-mini-flash-reasoning", # noqa: E501

252 + trust_remote_code=True,

253 + v0_only=True,

254 + max_model_len=10240),

255 "PhiMoEForCausalLM": _HfExamplesInfo("microsoft/Phi-3.5-MoE-instruct",

256 trust_remote_code=True),

257 "Plamo2ForCausalLM": _HfExamplesInfo("pfnet/plamo-2-1b",

tests/models/test_initialization.py

@@ -103,6 +103,9 @@ def test_can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch):

103 _initialize_kv_caches_v1), monkeypatch.context() as m):

104 if model_info.v0_only:

105 m.setenv("VLLM_USE_V1", "0")

106 + if model_arch == "Phi4FlashForCausalLM":

107 + # Phi4FlashForCausalLM only supports DIFFERENTIAL_FLASH_ATTN backend

108 + m.setenv("VLLM_ATTENTION_BACKEND", "DIFFERENTIAL_FLASH_ATTN")

109 LLM(

110 model_info.default,

111 tokenizer=model_info.tokenizer,

tests/test_utils.py

+25

@@ -458,6 +458,31 @@ def test_bind_kv_cache():

458 assert ctx['layers.2.self_attn'].kv_cache[0] is kv_cache[2]

459 assert ctx['layers.3.self_attn'].kv_cache[0] is kv_cache[3]

460

461 +def test_bind_kv_cache_kv_sharing():

462 + from vllm.attention import Attention

463 +

464 + ctx = {

465 + 'layers.0.self_attn': Attention(32, 128, 0.1),

466 + 'layers.1.self_attn': Attention(32, 128, 0.1),

467 + 'layers.2.self_attn': Attention(32, 128, 0.1),

468 + 'layers.3.self_attn': Attention(32, 128, 0.1),

469 + }

470 + kv_cache = [

471 + torch.zeros((1, )),

472 + torch.zeros((1, )),

473 + torch.zeros((1, )),

474 + torch.zeros((1, )),

475 + ]

476 + shared_kv_cache_layers = {

477 + 'layers.2.self_attn': 'layers.1.self_attn',

478 + 'layers.3.self_attn': 'layers.0.self_attn'

479 + }

480 + bind_kv_cache(ctx, [kv_cache], shared_kv_cache_layers)

481 + assert ctx['layers.0.self_attn'].kv_cache[0] is kv_cache[0]

482 + assert ctx['layers.1.self_attn'].kv_cache[0] is kv_cache[1]

483 + assert ctx['layers.2.self_attn'].kv_cache[0] is kv_cache[1]

484 + assert ctx['layers.3.self_attn'].kv_cache[0] is kv_cache[0]

485 +

486 def test_bind_kv_cache_non_attention():

487 from vllm.attention import Attention

488

vllm/attention/backends/blocksparse_attn.py

+2 -1

@@ -308,7 +308,8 @@ class BlocksparseFlashAttentionImpl(AttentionImpl):

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.")

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.")

vllm/attention/backends/differential_flash_attn.py new

+1000

@@ -0,0 +1,1000 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""" An implementation of https://arxiv.org/pdf/2410.05258 """

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

8 +

9 +import torch

10 +from einops import rearrange

11 +

12 +from vllm import _custom_ops as ops

13 +# yapf conflicts with isort for this block

14 +# yapf: disable

15 +from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

16 + AttentionLayer,

17 + AttentionMetadata,

18 + AttentionMetadataBuilder,

19 + AttentionType,

20 + is_quantized_kv_cache)

21 +from vllm.attention.backends.flash_attn import FlashAttentionBackend

22 +# yapf: enable

23 +from vllm.attention.backends.utils import (PAD_SLOT_ID, CommonAttentionState,

24 + compute_slot_mapping,

25 + compute_slot_mapping_start_idx,

26 + is_all_cross_attn_metadata_set,

27 + is_all_encoder_attn_metadata_set,

28 + is_block_tables_empty)

29 +from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,

30 + get_flash_attn_version)

31 +from vllm.logger import init_logger

32 +from vllm.multimodal import MultiModalPlaceholderMap

33 +from vllm.utils import async_tensor_h2d, make_tensor_with_pad

34 +from vllm.vllm_flash_attn import (flash_attn_varlen_func,

35 + flash_attn_with_kvcache)

36 +

37 +if TYPE_CHECKING:

38 + from vllm.worker.model_runner import (ModelInputForGPUBuilder,

39 + ModelInputForGPUWithSamplingMetadata)

40 +

41 +logger = init_logger(__name__)

42 +

43 +

44 +class DifferentialFlashAttentionBackend(AttentionBackend):

45 + accept_output_buffer = False

46 +

47 + @staticmethod

48 + def get_supported_head_sizes() -> List[int]:

49 + return [32, 64, 96, 128, 160, 192, 224, 256]

50 +

51 + @staticmethod

52 + def get_kv_cache_shape(

53 + num_blocks: int,

54 + block_size: int,

55 + num_kv_heads: int,

56 + head_size: int,

57 + ) -> Tuple[int, ...]:

58 + if block_size % 16 != 0:

59 + raise ValueError("Block size must be a multiple of 16.")

60 + assert num_kv_heads % 2 == 0, "num_kv_heads must be divisible by 2"

61 + return (2, 2, num_blocks, block_size, num_kv_heads // 2, head_size)

62 +

63 + @staticmethod

64 + def get_name() -> str:

65 + return "DIFFERENTIAL_FLASH_ATTN"

66 +

67 + @staticmethod

68 + def get_impl_cls() -> Type["DifferentialFlashAttentionImpl"]:

69 + return DifferentialFlashAttentionImpl

70 +

71 + @staticmethod

72 + def get_metadata_cls() -> Type["DifferentialFlashAttentionMetadata"]:

73 + return DifferentialFlashAttentionMetadata

74 +

75 + @staticmethod

76 + def get_builder_cls() -> Type["DifferentialFlashAttentionMetadataBuilder"]:

77 + return DifferentialFlashAttentionMetadataBuilder

78 +

79 + @staticmethod

80 + def get_state_cls() -> Type["CommonAttentionState"]:

81 + return CommonAttentionState

82 +

83 + @staticmethod

84 + def swap_blocks(

85 + src_kv_cache: torch.Tensor,

86 + dst_kv_cache: torch.Tensor,

87 + src_to_dst: torch.Tensor,

88 + ) -> None:

89 + src_key_cache = src_kv_cache[0]

90 + dst_key_cache = dst_kv_cache[0]

91 + ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst)

92 + src_value_cache = src_kv_cache[1]

93 + dst_value_cache = dst_kv_cache[1]

94 + ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst)

95 +

96 + @staticmethod

97 + def copy_blocks(

98 + kv_caches: List[torch.Tensor],

99 + src_to_dists: torch.Tensor,

100 + ) -> None:

101 + key_caches = [kv_cache[0] for kv_cache in kv_caches]

102 + value_caches = [kv_cache[1] for kv_cache in kv_caches]

103 +

104 + ops.copy_blocks(key_caches, value_caches, src_to_dists)

105 +

106 +

107 +@dataclass

108 +class DifferentialFlashAttentionMetadata(AttentionMetadata):

109 + """Metadata for FlashAttentionBackend.

110 +

111 + NOTE: Any python object stored here is not updated when it is

112 + cuda-graph replayed. If you have values that need to be changed

113 + dynamically, it should be stored in tensor. The tensor has to be

114 + updated from `CUDAGraphRunner.forward` API.

115 + """

116 + # (batch_size,). The sequence length per sequence. Sequence length means

117 + # the computed tokens + new tokens None if it is a decoding.

118 + seq_lens: Optional[List[int]]

119 + # seq_lens stored as a tensor.

120 + seq_lens_tensor: Optional[torch.Tensor]

121 +

122 + # NOTE(sang): Definition of context_len, query_len, and seq_len.

123 + # |---------- N-1 iteration --------|

124 + # |---------------- N iteration ---------------------|

125 + # |- tokenA -|......................|-- newTokens ---|

126 + # |---------- context_len ----------|

127 + # |-------------------- seq_len ---------------------|

128 + # |-- query_len ---|

129 +

130 + # Maximum sequence length among prefill batch. 0 if there are decoding

131 + # requests only.

132 + max_prefill_seq_len: int

133 + # Maximum sequence length among decode batch. 0 if there are prefill

134 + # requests only.

135 + max_decode_seq_len: int

136 + # (batch_size,) A tensor of context lengths (tokens that are computed

137 + # so far).

138 + context_lens_tensor: Optional[torch.Tensor]

139 +

140 + # (batch_size, max_blocks_per_seq).

141 + # Block addresses per sequence. (Seq id -> list of physical block)

142 + # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks

143 + # in the kv cache. Each block can contain up to block_size tokens.

144 + # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph

145 + # captured.

146 + block_tables: Optional[torch.Tensor]

147 +

148 + # Whether or not if cuda graph is enabled.

149 + # Cuda-graph is currently enabled for decoding only.

150 + # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.

151 +

152 + use_cuda_graph: bool

153 +

154 + # Maximum query length in the batch.

155 + max_query_len: Optional[int] = None

156 +

157 + # Max number of query tokens among request in the batch.

158 + max_decode_query_len: Optional[int] = None

159 +

160 + # (batch_size + 1,). The cumulative subquery lengths of the sequences in

161 + # the batch, used to index into subquery. E.g., if the subquery length

162 + # is [4, 6], it is [0, 4, 10].

163 + query_start_loc: Optional[torch.Tensor] = None

164 + # (batch_size + 1,). The cumulative sequence lengths of the sequences in

165 + # the batch, used to index into sequence. E.g., if the sequence length is

166 + # [4, 6], it is [0, 4, 10].

167 + seq_start_loc: Optional[torch.Tensor] = None

168 +

169 + _cached_prefill_metadata: Optional[

170 + "DifferentialFlashAttentionMetadata"] = None

171 + _cached_decode_metadata: Optional[

172 + "DifferentialFlashAttentionMetadata"] = None

173 +

174 + # Begin encoder attn & enc/dec cross-attn fields...

175 +

176 + # Encoder sequence lengths representation

177 + encoder_seq_lens: Optional[List[int]] = None

178 + encoder_seq_lens_tensor: Optional[torch.Tensor] = None

179 + # (batch_size + 1,). The cumulative sequence lengths of the sequences in

180 + # the batch, used to index into sequence. E.g., if the sequence length is

181 + # [4, 6], it is [0, 4, 10].

182 + encoder_seq_start_loc: Optional[torch.Tensor] = None

183 + # Maximum sequence length among encoder sequences

184 + max_encoder_seq_len: Optional[int] = None

185 + # Number of tokens input to encoder

186 + num_encoder_tokens: Optional[int] = None

187 +

188 + # Cross-attention memory-mapping data structures: slot mapping

189 + # and block tables

190 + cross_slot_mapping: Optional[torch.Tensor] = None

191 + cross_block_tables: Optional[torch.Tensor] = None

192 +

193 + # Cross-layer shared attention block tables

194 + cross_layer_shared_block_tables: Optional[torch.Tensor] = None

195 +

196 + @property

197 + def is_all_encoder_attn_metadata_set(self):

198 + '''

199 + All attention metadata required for encoder attention is set.

200 + '''

201 + return is_all_encoder_attn_metadata_set(self)

202 +

203 + @property

204 + def is_all_cross_attn_metadata_set(self):

205 + '''

206 + All attention metadata required for enc/dec cross-attention is set.

207 +

208 + Superset of encoder attention required metadata.

209 + '''

210 + return is_all_cross_attn_metadata_set(self)

211 +

212 + @property

213 + def prefill_metadata(

214 + self) -> Optional["DifferentialFlashAttentionMetadata"]:

215 + if self.num_prefills == 0:

216 + return None

217 +

218 + if self._cached_prefill_metadata is not None:

219 + return self._cached_prefill_metadata

220 +

221 + assert ((self.seq_lens is not None)

222 + or (self.encoder_seq_lens is not None))

223 + assert ((self.seq_lens_tensor is not None)

224 + or (self.encoder_seq_lens_tensor is not None))

225 +

226 + # Compute some attn_metadata fields which default to None

227 + query_start_loc = (None if self.query_start_loc is None else

228 + self.query_start_loc[:self.num_prefills + 1])

229 + slot_mapping = (None if self.slot_mapping is None else

230 + self.slot_mapping[:self.num_prefill_tokens])

231 + seq_lens = (None if self.seq_lens is None else

232 + self.seq_lens[:self.num_prefills])

233 + seq_lens_tensor = (None if self.seq_lens_tensor is None else

234 + self.seq_lens_tensor[:self.num_prefills])

235 + seq_start_loc = (None if self.seq_start_loc is None else

236 + self.seq_start_loc[:self.num_prefills + 1])

237 + context_lens_tensor = (None if self.context_lens_tensor is None else

238 + self.context_lens_tensor[:self.num_prefills])

239 + block_tables = (None if self.block_tables is None else

240 + self.block_tables[:self.num_prefills])

241 + cross_layer_shared_block_tables = (

242 + None if self.cross_layer_shared_block_tables is None else

243 + self.cross_layer_shared_block_tables[:self.num_prefills])

244 +

245 + self._cached_prefill_metadata = DifferentialFlashAttentionMetadata(

246 + num_prefills=self.num_prefills,

247 + num_prefill_tokens=self.num_prefill_tokens,

248 + num_decode_tokens=0,

249 + slot_mapping=slot_mapping,

250 + multi_modal_placeholder_index_maps=self.

251 + multi_modal_placeholder_index_maps,

252 + enable_kv_scales_calculation=self.enable_kv_scales_calculation,

253 + seq_lens=seq_lens,

254 + seq_lens_tensor=seq_lens_tensor,

255 + max_query_len=self.max_query_len,

256 + max_prefill_seq_len=self.max_prefill_seq_len,

257 + max_decode_query_len=0,

258 + max_decode_seq_len=0,

259 + query_start_loc=query_start_loc,

260 + seq_start_loc=seq_start_loc,

261 + context_lens_tensor=context_lens_tensor,

262 + block_tables=block_tables,

263 + cross_layer_shared_block_tables=cross_layer_shared_block_tables,

264 + use_cuda_graph=False,

265 + # Begin encoder & cross attn fields below...

266 + encoder_seq_lens=self.encoder_seq_lens,

267 + encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,

268 + encoder_seq_start_loc=self.encoder_seq_start_loc,

269 + max_encoder_seq_len=self.max_encoder_seq_len,

270 + cross_slot_mapping=self.cross_slot_mapping,

271 + cross_block_tables=self.cross_block_tables)

272 + return self._cached_prefill_metadata

273 +

274 + @property

275 + def decode_metadata(

276 + self) -> Optional["DifferentialFlashAttentionMetadata"]:

277 + if self.num_decode_tokens == 0:

278 + return None

279 +

280 + if self._cached_decode_metadata is not None:

281 + return self._cached_decode_metadata

282 + assert ((self.seq_lens_tensor is not None)

283 + or (self.encoder_seq_lens_tensor is not None))

284 +

285 + # Compute some attn_metadata fields which default to None

286 + slot_mapping = (None if self.slot_mapping is None else

287 + self.slot_mapping[self.num_prefill_tokens:])

288 + seq_lens_tensor = (None if self.seq_lens_tensor is None else

289 + self.seq_lens_tensor[self.num_prefills:])

290 + block_tables = (None if self.block_tables is None else

291 + self.block_tables[self.num_prefills:])

292 + cross_layer_shared_block_tables = (

293 + None if self.cross_layer_shared_block_tables is None else

294 + self.cross_layer_shared_block_tables[self.num_prefills:])

295 + self._cached_decode_metadata = DifferentialFlashAttentionMetadata(

296 + num_prefills=0,

297 + num_prefill_tokens=0,

298 + num_decode_tokens=self.num_decode_tokens,

299 + slot_mapping=slot_mapping,

300 + multi_modal_placeholder_index_maps=None,

301 + enable_kv_scales_calculation=True,

302 + seq_lens=None,

303 + seq_lens_tensor=seq_lens_tensor,

304 + max_decode_query_len=self.max_decode_query_len,

305 + max_query_len=self.max_query_len,

306 + max_prefill_seq_len=0,

307 + max_decode_seq_len=self.max_decode_seq_len,

308 + # Batch may be composed of prefill|decodes, adjust query start

309 + # indices to refer to the start of decodes. E.g.

310 + # in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].

311 + query_start_loc=(self.query_start_loc[self.num_prefills:] -

312 + self.query_start_loc[self.num_prefills])

313 + if self.query_start_loc is not None else None,

314 + seq_start_loc=self.seq_start_loc[self.num_prefills:]

315 + if self.seq_start_loc is not None else None,

316 + context_lens_tensor=None,

317 + block_tables=block_tables,

318 + cross_layer_shared_block_tables=cross_layer_shared_block_tables,

319 + use_cuda_graph=self.use_cuda_graph,

320 + # Begin encoder & cross attn fields below...

321 + encoder_seq_lens=self.encoder_seq_lens,

322 + encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,

323 + encoder_seq_start_loc=self.encoder_seq_start_loc,

324 + max_encoder_seq_len=self.max_encoder_seq_len,

325 + cross_slot_mapping=self.cross_slot_mapping,

326 + cross_block_tables=self.cross_block_tables)

327 + return self._cached_decode_metadata

328 +

329 + def advance_step(self,

330 + model_input: "ModelInputForGPUWithSamplingMetadata",

331 + sampled_token_ids: Optional[torch.Tensor],

332 + block_size: int,

333 + num_seqs: int,

334 + num_queries: int,

335 + turn_prefills_into_decodes: bool = False):

336 + """

337 + Update metadata in-place to advance one decode step.

338 + """

339 + # When using cudagraph, the num_seqs is padded to the next captured

340 + # batch sized, but num_queries tracks the actual number of requests in

341 + # the batch. For --enforce-eager mode, num_seqs == num_queries

342 + if num_seqs != num_queries:

343 + assert num_seqs > num_queries

344 + assert self.use_cuda_graph

345 +

346 + if turn_prefills_into_decodes:

347 + # When Multi-Step is enabled with Chunked-Prefill, prefills and

348 + # decodes are scheduled together. In the first step, all the

349 + # prefills turn into decodes. This update reflects that

350 + # conversion.

351 + assert self.num_decode_tokens + self.num_prefills == num_seqs

352 + self.num_decode_tokens += self.num_prefills

353 + self.num_prefills = 0

354 + self.num_prefill_tokens = 0

355 + self.max_prefill_seq_len = 0

356 + self.max_query_len = 1

357 +

358 + self.slot_mapping = self.slot_mapping[:num_seqs]

359 + else:

360 + assert self.seq_lens is not None

361 + assert self.max_decode_seq_len == max(self.seq_lens)

362 +

363 + assert self.num_prefills == 0

364 + assert self.num_prefill_tokens == 0

365 + assert self.num_decode_tokens == num_seqs

366 + assert self.slot_mapping.shape == (num_seqs, )

367 +

368 + assert self.seq_lens is not None

369 + assert len(self.seq_lens) == num_seqs

370 + assert self.seq_lens_tensor is not None

371 + assert self.seq_lens_tensor.shape == (num_seqs, )

372 + assert self.max_query_len == 1

373 + assert self.max_prefill_seq_len == 0

374 +

375 + assert self.query_start_loc is not None

376 + assert self.query_start_loc.shape == (num_queries + 1, )

377 + assert self.seq_start_loc is not None

378 + assert self.seq_start_loc.shape == (num_seqs + 1, )

379 +

380 + assert self.context_lens_tensor is not None

381 + assert self.context_lens_tensor.shape == (num_queries, )

382 +

383 + assert self.block_tables is not None

384 + assert self.block_tables.shape[0] == num_seqs

385 +

386 + # Update query lengths. Note that we update only queries and not seqs,

387 + # since tensors may be padded due to captured cuda graph batch size

388 + for i in range(num_queries):

389 + self.seq_lens[i] += 1

390 + self.max_decode_seq_len = max(self.seq_lens)

391 +

392 + ops.advance_step_flashattn(num_seqs=num_seqs,

393 + num_queries=num_queries,

394 + block_size=block_size,

395 + input_tokens=model_input.input_tokens,

396 + sampled_token_ids=sampled_token_ids,

397 + input_positions=model_input.input_positions,

398 + seq_lens=self.seq_lens_tensor,

399 + slot_mapping=self.slot_mapping,

400 + block_tables=self.block_tables)

401 +

402 +

403 +class DifferentialFlashAttentionMetadataBuilder(

404 + AttentionMetadataBuilder[DifferentialFlashAttentionMetadata]):

405 +

406 + def __init__(self, input_builder: "ModelInputForGPUBuilder"):

407 + self.input_builder = input_builder

408 + self.runner = input_builder.runner

409 + self.sliding_window = input_builder.sliding_window

410 + self.block_size = input_builder.block_size

411 +

412 + def prepare(self):

413 + self.slot_mapping: List[int] = []

414 + self.prefill_seq_lens: List[int] = []

415 + self.context_lens: List[int] = []

416 + self.block_tables: List[List[int]] = []

417 + self.cross_layer_shared_block_tables: List[List[int]] = []

418 + self.curr_seq_lens: List[int] = []

419 + self.multimodal_placeholder_maps: Dict[

420 + str,

421 + MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap)

422 + self.num_prefills = 0

423 + self.num_prefill_tokens = 0

424 + self.num_decode_tokens = 0

425 + self.has_prefix_cache_hit = False

426 +

427 + def _add_seq_group(

428 + self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",

429 + chunked_prefill_enabled: bool, prefix_cache_hit: bool):

430 + """Add a sequence group to the metadata. Specifically update/append

431 + 1. context length.

432 + 2. block table.

433 + 3. slot mapping.

434 + """

435 + # TODO: add support for chunked prefill and prefix caching.

436 + assert not chunked_prefill_enabled, \

437 + "chunked prefill is not supported for now"

438 + assert not prefix_cache_hit, "prefix caching is not supported for now"

439 +

440 + is_prompt = inter_data.is_prompt

441 + block_tables = inter_data.block_tables

442 +

443 + for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,

444 + curr_sliding_window_block) in zip(

445 + inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],

446 + inter_data.orig_seq_lens, inter_data.seq_lens,

447 + inter_data.query_lens, inter_data.context_lens,

448 + inter_data.curr_sliding_window_blocks):

449 + self.context_lens.append(context_len)

450 +

451 + if is_prompt:

452 + mm_maps = inter_data.multi_modal_placeholder_maps

453 + if mm_maps:

454 + for modality, placeholders in mm_maps.items():

455 + self.multimodal_placeholder_maps[modality].extend(

456 + placeholders)

457 +

458 + self.num_prefills += 1

459 + self.num_prefill_tokens += token_len

460 + self.prefill_seq_lens.append(seq_len)

461 + else:

462 + self.num_decode_tokens += query_len

463 + self.curr_seq_lens.append(curr_seq_len)

464 +

465 + # Compute block table.

466 + # TODO(sang): Combine chunked prefill and prefix caching by

467 + # only allowing multiple of block_size chunk size.

468 + # NOTE: This only works for oooooooxxx style attention.

469 + block_table = []

470 + if prefix_cache_hit:

471 + # NOTE(woosuk): For flash-attn, the block table should

472 + # include the entries for the incoming prefill tokens.

473 + block_table = block_tables[seq_id]

474 + elif ((chunked_prefill_enabled or not is_prompt)

475 + and block_tables is not None):

476 + if curr_sliding_window_block == 0:

477 + block_table = block_tables[seq_id]

478 + else:

479 + block_table = block_tables[seq_id][

480 + -curr_sliding_window_block:]

481 + self.block_tables.append(block_table)

482 +

483 + cross_layer_shared_block_table = []

484 + if prefix_cache_hit:

485 + cross_layer_shared_block_table = block_tables[seq_id]

486 + elif block_tables is not None:

487 + if curr_sliding_window_block == 0:

488 + cross_layer_shared_block_table = block_tables[seq_id]

489 + else:

490 + cross_layer_shared_block_table = block_tables[seq_id][

491 + -curr_sliding_window_block:]

492 + self.cross_layer_shared_block_tables.append(

493 + cross_layer_shared_block_table)

494 +

495 + # Compute slot mapping.

496 + is_profile_run = is_block_tables_empty(block_tables)

497 + start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,

498 + context_len,

499 + self.sliding_window)

500 + compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,

501 + seq_len, context_len, start_idx,

502 + self.block_size, inter_data.block_tables)

503 +

504 + def _get_graph_runner_block_tables(self, num_seqs: int,

505 + block_tables: List[List[int]],

506 + graph_block_tables) -> torch.Tensor:

507 + # The shape of graph_block_tables is

508 + # [max batch size, max context len // block size].

509 + # max_batch_size, max_blocks = self.runner.graph_block_tables.shape

510 + max_batch_size, max_blocks = graph_block_tables.shape

511 + assert max_batch_size >= num_seqs

512 +

513 + # graph_block_tables = self.runner.graph_block_tables[:num_seqs]

514 + graph_block_tables = graph_block_tables[:num_seqs]

515 + for i, block_table in enumerate(block_tables):

516 + if block_table:

517 + num_blocks = len(block_table)

518 + if num_blocks <= max_blocks:

519 + graph_block_tables[i, :num_blocks] = block_table

520 + else:

521 + # It may be possible to have more blocks allocated due

522 + # to lookahead slots of multi-step, however, they are

523 + # not used anyway, so can be safely ignored.

524 + graph_block_tables[

525 + i, :max_blocks] = block_table[:max_blocks]

526 +

527 + return torch.from_numpy(graph_block_tables).to(

528 + device=self.runner.device, non_blocking=True)

529 +

530 + def build(self, seq_lens: List[int], query_lens: List[int],

531 + cuda_graph_pad_size: int, batch_size: int):

532 + """Build attention metadata with on-device tensors.

533 +

534 + Args:

535 + seq_lens: The maybe padded sequence lengths of the input sequences.

536 + query_lens: The query lengths of the input sequences.

537 + cuda_graph_pad_size: The padding size for cuda graph.

538 + -1 if cuda graph is not used.

539 + batch_size: The maybe padded batch size.

540 + """

541 + prefix_cache_hit = any([

542 + inter_data.prefix_cache_hit

543 + for inter_data in self.input_builder.inter_data_list

544 + ])

545 + for inter_data in self.input_builder.inter_data_list:

546 + self._add_seq_group(inter_data,

547 + self.input_builder.chunked_prefill_enabled,

548 + prefix_cache_hit)

549 +

550 + device = self.runner.device

551 + use_captured_graph = cuda_graph_pad_size != -1

552 +

553 + max_query_len = max(query_lens)

554 + decode_query_lens = query_lens[self.num_prefills:]

555 + if len(decode_query_lens) > 0:

556 + max_decode_query_len = max(decode_query_lens)

557 + else:

558 + max_decode_query_len = 1

559 + max_prefill_seq_len = max(self.prefill_seq_lens, default=0)

560 + max_decode_seq_len = max(self.curr_seq_lens, default=0)

561 + num_decode_tokens = self.num_decode_tokens

562 + query_start_loc = list(accumulate(query_lens, initial=0))

563 + seq_start_loc = list(accumulate(seq_lens, initial=0))

564 +

565 + num_seqs = len(seq_lens)

566 + if use_captured_graph:

567 + self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)

568 + self.block_tables.extend([] * cuda_graph_pad_size)

569 +

570 + self.cross_layer_shared_block_tables.extend([] *

571 + cuda_graph_pad_size)

572 +

573 + num_decode_tokens = batch_size - self.num_prefill_tokens

574 + block_tables = self._get_graph_runner_block_tables(

575 + num_seqs, self.block_tables, self.runner.graph_block_tables)

576 + cross_layer_shared_block_tables = \

577 + self._get_graph_runner_block_tables(

578 + num_seqs, self.cross_layer_shared_block_tables,

579 + self.runner.cross_layer_shared_graph_block_tables)

580 + else:

581 + block_tables = make_tensor_with_pad(

582 + self.block_tables,

583 + pad=0,

584 + dtype=torch.int,

585 + device=device,

586 + )

587 + cross_layer_shared_block_tables = make_tensor_with_pad(

588 + self.cross_layer_shared_block_tables,

589 + pad=0,

590 + dtype=torch.int,

591 + device=device,

592 + )

593 + assert max_query_len > 0, ("query_lens: {}".format(query_lens))

594 +

595 + assert device is not None

596 + context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int,

597 + device, self.runner.pin_memory)

598 + seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,

599 + self.runner.pin_memory)

600 + slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long,

601 + device, self.runner.pin_memory)

602 + query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,

603 + device,

604 + self.runner.pin_memory)

605 + seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32,

606 + device, self.runner.pin_memory)

607 + placeholder_index_maps = {

608 + modality: placeholder_map.index_map()

609 + for modality, placeholder_map in

610 + self.multimodal_placeholder_maps.items()

611 + }

612 +

613 + return DifferentialFlashAttentionMetadata(

614 + num_prefills=self.num_prefills,

615 + slot_mapping=slot_mapping_tensor,

616 + num_prefill_tokens=self.num_prefill_tokens,

617 + num_decode_tokens=num_decode_tokens,

618 + seq_lens=seq_lens,

619 + multi_modal_placeholder_index_maps=placeholder_index_maps,

620 + enable_kv_scales_calculation=True,

621 + seq_lens_tensor=seq_lens_tensor,

622 + max_query_len=max_query_len,

623 + max_decode_query_len=max_decode_query_len,

624 + max_prefill_seq_len=max_prefill_seq_len,

625 + max_decode_seq_len=max_decode_seq_len,

626 + query_start_loc=query_start_loc_tensor,

627 + seq_start_loc=seq_start_loc_tensor,

628 + context_lens_tensor=context_lens_tensor,

629 + block_tables=block_tables,

630 + cross_layer_shared_block_tables=cross_layer_shared_block_tables,

631 + use_cuda_graph=use_captured_graph,

632 + )

633 +

634 +

635 +class DifferentialFlashAttentionImpl(AttentionImpl):

636 + """

637 + If the input tensors contain prompt tokens, the layout is as follows:

638 + |<--------------- num_prefill_tokens ----------------->|

640 +

641 + Otherwise, the layout is as follows:

642 + |<----------------- num_decode_tokens ------------------>|

644 +

645 + Generation tokens can contain padding when cuda-graph is used.

646 + Currently, prompt tokens don't contain any padding.

647 +

648 + The prompts might have different lengths, while the generation tokens

649 + always have length 1.

650 +

651 + If chunked prefill is enabled, prefill tokens and decode tokens can be

652 + batched together in a flattened 1D query.

653 +

654 + |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->|

655 + |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->|

656 +

657 + Currently, cuda graph is disabled for chunked prefill, meaning there's no

658 + padding between prefill and decode tokens.

659 + """

660 +

661 + def __init__(

662 + self,

663 + num_heads: int,

664 + head_size: int,

665 + scale: float,

666 + num_kv_heads: int,

667 + alibi_slopes: Optional[List[float]],

668 + sliding_window: Optional[int],

669 + kv_cache_dtype: str,

670 + blocksparse_params: Optional[Dict[str, Any]] = None,

671 + logits_soft_cap: Optional[float] = None,

672 + attn_type: str = AttentionType.DECODER,

673 + kv_sharing_target_layer_name: Optional[str] = None,

674 + use_irope: bool = False,

675 + differential_flash_attention_config: Optional[Dict[str, Any]] = None,

676 + ) -> None:

677 + if differential_flash_attention_config is None:

678 + differential_flash_attention_config = {}

679 + self.differential_flash_attention_config = \

680 + differential_flash_attention_config

681 + self.used_shared_kv_cache = kv_sharing_target_layer_name is not None

682 + 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.")

686 + if use_irope:

687 + logger.warning(

688 + "Using irope in V0 is not supported yet, it will fall back "

689 + "to global attention for long context.")

690 + self.num_heads = num_heads

691 + self.head_size = head_size

692 + self.scale = float(scale)

693 + self.num_kv_heads = num_kv_heads

694 + if alibi_slopes is not None:

695 + alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)

696 + self.alibi_slopes = alibi_slopes

697 + self.sliding_window = ((sliding_window - 1,

698 + 0) if sliding_window is not None else (-1, -1))

699 + self.kv_cache_dtype = kv_cache_dtype

700 + self.vllm_flash_attn_version = get_flash_attn_version(

701 + requires_alibi=self.alibi_slopes is not None)

702 + if is_quantized_kv_cache(self.kv_cache_dtype) and (

703 + not self.kv_cache_dtype.startswith("fp8")

704 + or not flash_attn_supports_fp8()):

705 + raise NotImplementedError(

706 + f"FlashAttention does not support {self.kv_cache_dtype} "

707 + "kv-cache on this device "

708 + f"(FA supports fp8 = {flash_attn_supports_fp8()}).")

709 + if logits_soft_cap is None:

710 + # In flash-attn, setting logits_soft_cap as 0 means no soft cap.

711 + logits_soft_cap = 0

712 + self.logits_soft_cap = logits_soft_cap

713 +

714 + assert self.num_heads % self.num_kv_heads == 0

715 + self.num_queries_per_kv = self.num_heads // self.num_kv_heads

716 +

717 + support_head_sizes = FlashAttentionBackend.get_supported_head_sizes()

718 + if head_size not in support_head_sizes:

719 + raise ValueError(

720 + f"Head size {head_size} is not supported by FlashAttention. "

721 + f"Supported head sizes are: {support_head_sizes}.")

722 + self.attn_type = attn_type

723 +

724 + self.lambda_full = None

725 + self.subln = self.differential_flash_attention_config["subln"]

726 +

727 + def split_heads(self, x):

728 + # split by num_heads, the stripe pattern is friendly to tensor parallel.

729 + x = rearrange(x, "... (H two) D -> ... H two D", two=2)

730 + x1 = x[..., 0, :]

731 + x2 = x[..., 1, :]

732 + return x1.contiguous(), x2.contiguous()

733 +

734 + def split_kv_cache(self, x):

735 + # split by num_heads, the stripe pattern is friendly to tensor parallel.

736 + if x.numel() == 0:

737 + return torch.empty(0), torch.empty(0)

738 +

739 + x1, x2 = x[0], x[1]

740 + return x1, x2

741 +

742 + def populate_kv_cache(self, layer: AttentionLayer, key: torch.Tensor,

743 + value: torch.Tensor, kv_cache: torch.Tensor,

744 + attn_metadata: DifferentialFlashAttentionMetadata):

745 + if kv_cache.numel() > 0 and key is not None and value is not None:

746 + updated_slot_mapping = attn_metadata.slot_mapping

747 + torch.ops._C_cache_ops.reshape_and_cache_flash(

748 + key,

749 + value,

750 + kv_cache[0],

751 + kv_cache[1],

752 + updated_slot_mapping.flatten(),

753 + self.kv_cache_dtype,

754 + layer._k_scale,

755 + layer._v_scale,

756 + )

757 +

758 + def forward_generate_kv_cache(

759 + self, query: torch.Tensor, key: Optional[torch.Tensor],

760 + value: Optional[torch.Tensor], k_cache: torch.Tensor,

761 + v_cache: torch.Tensor,

762 + attn_metadata: DifferentialFlashAttentionMetadata) -> torch.Tensor:

763 +

764 + head_size = self.head_size

765 + num_heads = self.num_heads // 2

766 + num_kv_heads = self.num_kv_heads // 2

767 +

768 + query = query.view(-1, num_heads, head_size)

769 + if key is not None:

770 + assert value is not None

771 + key = key.view(-1, num_kv_heads, head_size)

772 + value = value.view(-1, num_kv_heads, head_size)

773 + else:

774 + assert value is None

775 +

776 + num_prefill_tokens = attn_metadata.num_prefill_tokens

777 + num_decode_tokens = attn_metadata.num_decode_tokens

778 + assert key.shape[

779 + 0] == num_prefill_tokens + num_decode_tokens, "key shape mismatch"

780 + assert value.shape[

781 + 0] == num_prefill_tokens + num_decode_tokens, "value shape mismatch"

782 +

783 + output = torch.empty_like(query)

784 + # Query for decode. KV is not needed because it is already cached.

785 + decode_query = query[num_prefill_tokens:]

786 + # QKV for prefill.

787 + query = query[:num_prefill_tokens]

788 + if key is not None and value is not None:

789 + key = key[:num_prefill_tokens]

790 + value = value[:num_prefill_tokens]

791 +

792 + assert query.shape[0] == num_prefill_tokens, "query shape mismatch"

793 + assert decode_query.shape[

794 + 0] == num_decode_tokens, "decode query shape mismatch"

795 +

796 + if prefill_meta := attn_metadata.prefill_metadata:

797 + # Prompt run.

798 + if k_cache.numel() == 0 \

799 + or prefill_meta.block_tables is None \

800 + or prefill_meta.block_tables.numel() == 0:

801 + # normal attention

802 + prefill_output = flash_attn_varlen_func(

803 + q=query,

804 + k=key,

805 + v=value,

806 + cu_seqlens_q=prefill_meta.seq_start_loc,

807 + cu_seqlens_k=prefill_meta.seq_start_loc,

808 + max_seqlen_q=prefill_meta.max_prefill_seq_len,

809 + max_seqlen_k=prefill_meta.max_prefill_seq_len,

810 + softmax_scale=self.scale,

811 + causal=True,

812 + window_size=self.sliding_window,

813 + alibi_slopes=self.alibi_slopes,

814 + softcap=self.logits_soft_cap,

815 + )

816 + assert prefill_output.shape == output[:

817 + num_prefill_tokens].shape

818 + output[:num_prefill_tokens] = prefill_output

819 + else:

820 + raise Exception("prefix caching not supported")

821 +

822 + if decode_meta := attn_metadata.decode_metadata:

823 + block_tables_arg = decode_meta.block_tables

824 + try:

825 + output[num_prefill_tokens:] = flash_attn_with_kvcache(

826 + q=decode_query.unsqueeze(1),

827 + k_cache=k_cache,

828 + v_cache=v_cache,

829 + block_table=block_tables_arg,

830 + cache_seqlens=decode_meta.seq_lens_tensor,

831 + softmax_scale=self.scale,

832 + causal=True,

833 + window_size=self.sliding_window,

834 + alibi_slopes=self.alibi_slopes,

835 + softcap=self.logits_soft_cap,

836 + ).squeeze(1)

837 + except Exception as e:

838 + logger.error("Error in PagedAttention.forward_decode: %s",

839 + str(e))

840 + raise e

841 +

842 + # Reshape the output tensor.

843 + return output.view(-1, num_heads, head_size)

844 +

845 + def forward_with_kv_cache_only(

846 + self,

847 + query: torch.Tensor,

848 + k_cache: torch.Tensor,

849 + v_cache: torch.Tensor,

850 + attn_metadata: DifferentialFlashAttentionMetadata,

851 + ):

852 + if not attn_metadata.decode_metadata:

853 + block_tables_arg = attn_metadata.cross_layer_shared_block_tables

854 + else:

855 + block_tables_arg = attn_metadata.block_tables

856 +

857 + output = flash_attn_with_kvcache(

858 + q=query.unsqueeze(1),

859 + k_cache=k_cache,

860 + v_cache=v_cache,

861 + block_table=block_tables_arg,

862 + cache_seqlens=attn_metadata.seq_lens_tensor,

863 + softmax_scale=self.scale,

864 + causal=True,

865 + window_size=self.sliding_window,

866 + alibi_slopes=self.alibi_slopes,

867 + softcap=self.logits_soft_cap,

868 + ).squeeze(1)

869 + return output

870 +

871 + def forward(

872 + self,

873 + layer: AttentionLayer,

874 + q: torch.Tensor,

875 + k: torch.Tensor,

876 + v: torch.Tensor,

877 + kv_cache: torch.Tensor,

878 + attn_metadata: DifferentialFlashAttentionMetadata,

879 + output: Optional[torch.Tensor] = None,

880 + output_scale: Optional[torch.Tensor] = None,

881 + ) -> torch.Tensor:

882 + """Forward pass with FlashAttention.

883 +

884 + Args:

885 + query: shape = [num_tokens, num_heads, head_size]

886 + key: shape = [num_tokens, num_kv_heads, head_size]

887 + value: shape = [num_tokens, num_kv_heads, head_size]

888 + output: shape = [num_tokens, num_heads, head_size]

889 + kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]

890 + NOTE: kv_cache will be an empty tensor with shape [0]

891 + for profiling run.

892 + attn_metadata: Metadata for attention.

893 + NOTE: It in-place updates the output tensor.

894 + NOTE: FP8 quantization, flash-attn expect the size of

895 + {q,k,v}_descale to be (num_sequences, num_kv_heads).

896 + We use torch's .expand() to avoid duplicating values

897 + """

898 + if self.lambda_full is None:

899 + self.lambda_init = self.differential_flash_attention_config[

900 + "lambda_init"]

901 + lambda_q1 = self.differential_flash_attention_config["lambda_q1"]

902 + lambda_k1 = self.differential_flash_attention_config["lambda_k1"]

903 + lambda_q2 = self.differential_flash_attention_config["lambda_q2"]

904 + lambda_k2 = self.differential_flash_attention_config["lambda_k2"]

905 + lambda_1 = torch.exp(

906 + torch.sum(lambda_q1 * lambda_k1, dim=-1).float()).type_as(q)

907 + lambda_2 = torch.exp(

908 + torch.sum(lambda_q2 * lambda_k2, dim=-1).float()).type_as(q)

909 + self.lambda_full = lambda_1 - lambda_2 + self.lambda_init

910 +

911 + if not self.used_shared_kv_cache: # need to generate kv-cache

912 + q = q.view(-1, self.num_heads, self.head_size)

913 + k = k.view(-1, self.num_kv_heads, self.head_size)

914 + v = v.view(-1, self.num_kv_heads, self.head_size)

915 +

916 + q1, q2 = self.split_heads(q)

917 + k1, k2 = self.split_heads(k)

918 + v1, v2 = self.split_heads(v)

919 +

920 + # kv_cache shape is (2, 2, num_blocks, block_size, num_kv_heads // 2, head_size) # noqa: E501

921 + # Split by half along the first dimension.

922 + kv_cache1, kv_cache2 = self.split_kv_cache(kv_cache)

923 + assert kv_cache1.is_contiguous(), "kv_cache1 is not contiguous"

924 + assert kv_cache2.is_contiguous(), "kv_cache2 is not contiguous"

925 +

926 + if kv_cache1.numel() != 0:

927 + self.populate_kv_cache(layer, k1, v1, kv_cache1, attn_metadata)

928 + self.populate_kv_cache(layer, k2, v2, kv_cache2, attn_metadata)

929 +

930 + key_cache1, value_cache1 = self.split_kv_cache(kv_cache1)

931 + key_cache2, value_cache2 = self.split_kv_cache(kv_cache2)

932 + else:

933 + key_cache1, value_cache1 = torch.empty(0), torch.empty(0)

934 + key_cache2, value_cache2 = torch.empty(0), torch.empty(0)

935 + attn11 = self.forward_generate_kv_cache(q1, k1, v1, key_cache1,

936 + value_cache1,

937 + attn_metadata)

938 + attn12 = self.forward_generate_kv_cache(q1, k1, v2, key_cache1,

939 + value_cache2,

940 + attn_metadata)

941 + attn11 = attn11.view(q1.shape)

942 + attn12 = attn12.view(q1.shape)

943 + attn1 = torch.cat([attn11, attn12], dim=-1)

944 +

945 + attn21 = self.forward_generate_kv_cache(q2, k2, v1, key_cache2,

946 + value_cache1,

947 + attn_metadata)

948 + attn22 = self.forward_generate_kv_cache(q2, k2, v2, key_cache2,

949 + value_cache2,

950 + attn_metadata)

951 + attn21 = attn21.view(q2.shape)

952 + attn22 = attn22.view(q2.shape)

953 + attn2 = torch.cat([attn21, attn22], dim=-1)

954 +

955 + attn = attn1 - self.lambda_full * attn2

956 + # attn shape (-1, self.num_heads // 2, 2 * self.head_dim)

957 + attn = self.subln(attn)

958 + attn = attn * (1 - self.lambda_init)

959 + # reshape back to 2 * num_head

960 + attn_output = rearrange(attn,

961 + "... H (two D) -> ... (H two) D",

962 + two=2)

963 +

964 + else: # re-use the kv cache, full attention

965 + q = q.view(-1, self.num_heads, self.head_size)

966 + q1, q2 = self.split_heads(q)

967 + # kv_cache shape is (2, num_blocks, block_size, num_kv_heads, head_size) # noqa: E501

968 + kv_cache1, kv_cache2 = self.split_kv_cache(kv_cache)

969 + key_cache1, value_cache1 = kv_cache1[0], kv_cache1[1]

970 + key_cache2, value_cache2 = kv_cache2[0], kv_cache2[1]

971 +

972 + attn11 = self.forward_with_kv_cache_only(q1, key_cache1,

973 + value_cache1,

974 + attn_metadata)

975 + attn12 = self.forward_with_kv_cache_only(q1, key_cache1,

976 + value_cache2,

977 + attn_metadata)

978 + attn11 = attn11.view(q1.shape)

979 + attn12 = attn12.view(q1.shape)

980 + attn1 = torch.cat([attn11, attn12], dim=-1)

981 +

982 + attn21 = self.forward_with_kv_cache_only(q2, key_cache2,

983 + value_cache1,

984 + attn_metadata)

985 + attn22 = self.forward_with_kv_cache_only(q2, key_cache2,

986 + value_cache2,

987 + attn_metadata)

988 + attn21 = attn21.view(q2.shape)

989 + attn22 = attn22.view(q2.shape)

990 + attn2 = torch.cat([attn21, attn22], dim=-1)

991 +

992 + attn = attn1 - self.lambda_full * attn2

993 + attn = self.subln(attn)

994 + attn = attn * (1 - self.lambda_init)

995 + # reshape back to 2 * num_head

996 + attn_output = rearrange(attn,

997 + "... H (two D) -> ... (H two) D",

998 + two=2)

999 + attn_output = attn_output.view(-1, self.num_heads * self.head_size)

1000 + return attn_output

vllm/attention/backends/dual_chunk_flash_attn.py

+2 -1

@@ -295,7 +295,8 @@ class DualChunkFlashAttentionImpl(FlashAttentionImpl):

295 dual_chunk_attention_config: Optional[Dict[str, Any]] = None,

296 ) -> None:

297 if kv_sharing_target_layer_name is not None:

298 - raise NotImplementedError("KV sharing is not supported in V0.")

298 + raise NotImplementedError("KV sharing is not supported in V0 "

299 + "DUAL_CHUNK_FLASH_ATTN backend.")

300 self.num_heads = num_heads

301 self.head_size = head_size

302 self.scale = float(scale)

vllm/attention/backends/flash_attn.py

+2 -1

@@ -622,7 +622,8 @@ class FlashAttentionImpl(AttentionImpl):

622 use_irope: bool = False,

623 ) -> None:

624 if kv_sharing_target_layer_name is not None:

625 - raise NotImplementedError("KV sharing is not supported in V0.")

625 + raise NotImplementedError("KV sharing is not supported in V0 "

626 + "FLASH_ATTN backend.")

627 if blocksparse_params is not None:

628 raise ValueError(

629 "FlashAttention does not support block-sparse attention.")

vllm/attention/backends/flashinfer.py

+2 -1

@@ -1006,7 +1006,8 @@ class FlashInferImpl(AttentionImpl):

1006 use_irope: bool = False,

1007 ) -> None:

1008 if kv_sharing_target_layer_name is not None:

1009 - raise NotImplementedError("KV sharing is not supported in V0.")

1009 + raise NotImplementedError("KV sharing is not supported in V0 "

1010 + "FLASHINFER backend.")

1011 if use_irope:

1012 logger.warning_once(

1013 "Using irope in FlashInfer is not supported yet, it will fall"

vllm/attention/backends/hpu_attn.py

+2 -1

@@ -115,7 +115,8 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):

115 ) -> None:

116 super(AttentionImpl, self).__init__()

117 if kv_sharing_target_layer_name is not None:

118 - raise NotImplementedError("KV sharing is not supported in V0.")

118 + raise NotImplementedError("KV sharing is not supported in V0 "

119 + "HPU_ATTN backend.")

120 if use_irope:

121 logger.warning_once(

122 "Using irope in HPU is not supported yet, it will fall back "

vllm/attention/backends/rocm_flash_attn.py

+2 -1

@@ -501,7 +501,8 @@ class ROCmFlashAttentionImpl(AttentionImpl):

501 use_irope: bool = False,

502 ) -> None:

503 if kv_sharing_target_layer_name is not None:

504 - raise NotImplementedError("KV sharing is not supported in V0.")

504 + raise NotImplementedError("KV sharing is not supported in V0 "

505 + "ROCM_FLASH backend.")

506 if use_irope:

507 logger.warning_once(

508 "Using irope in ROCm Flash Attention is not supported yet, it "

vllm/attention/backends/xformers.py

+2 -1

@@ -394,7 +394,8 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):

394 use_irope: bool = False,

395 ) -> None:

396 if kv_sharing_target_layer_name is not None:

397 - raise NotImplementedError("KV sharing is not supported in V0.")

397 + raise NotImplementedError("KV sharing is not supported in V0 "

398 + "XFORMERS backend.")

399 if blocksparse_params is not None:

400 raise ValueError(

401 "XFormers does not support block-sparse attention.")

vllm/attention/layer.py

-4

@@ -160,10 +160,6 @@ class Attention(nn.Module):

160 self.attn_type = attn_type

161

162 if kv_sharing_target_layer_name is not None:

163 - if not envs.VLLM_USE_V1:

164 - raise NotImplementedError(

165 - "Cross-layer KV sharing is not supported in V0.")

166 -

163 validate_kv_sharing_target(

164 prefix,

165 kv_sharing_target_layer_name,

vllm/model_executor/layers/logits_processor.py

+2 -1

@@ -59,11 +59,12 @@ class LogitsProcessor(nn.Module):

59 hidden_states: torch.Tensor,

60 sampling_metadata: Optional[SamplingMetadata] = None,

61 embedding_bias: Optional[torch.Tensor] = None,

62 + prune_hidden_states: bool = True,

63 ) -> Optional[torch.Tensor]:

64 if self.logits_as_input:

65 logits = hidden_states

66 else:

66 - if sampling_metadata is not None:

67 + if sampling_metadata is not None and prune_hidden_states:

68 hidden_states = _prune_hidden_states(hidden_states,

69 sampling_metadata)

vllm/model_executor/models/phi4flash.py new

+746

@@ -0,0 +1,746 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import math

4 +from collections.abc import Iterable

5 +from typing import Optional, Union

6 +

7 +import torch

8 +import torch.nn as nn

9 +from transformers.activations import ACT2FN

10 +

11 +import vllm.envs as envs

12 +from vllm.attention import Attention, AttentionMetadata, AttentionType

13 +from vllm.attention.selector import _Backend

14 +from vllm.config import CacheConfig, VllmConfig

15 +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

16 +from vllm.forward_context import ForwardContext, get_forward_context

17 +from vllm.logger import init_logger

18 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,

19 + MergedColumnParallelLinear,

20 + RowParallelLinear)

21 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

22 +from vllm.model_executor.layers.mamba.ops.causal_conv1d import (

23 + causal_conv1d_fn, causal_conv1d_update)

24 +from vllm.model_executor.layers.mamba.ops.mamba_ssm import (

25 + selective_scan_fn, selective_state_update)

26 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

27 +from vllm.model_executor.layers.vocab_parallel_embedding import (

28 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

29 +from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,

30 + SupportsV0Only)

31 +from vllm.model_executor.models.mamba_cache import (MambaCacheManager,

32 + MambaCacheParams)

33 +from vllm.model_executor.sampling_metadata import SamplingMetadata

34 +from vllm.sequence import IntermediateTensors

35 +

36 +from .utils import make_layers, maybe_prefix

37 +

38 +logger = init_logger(__name__)

39 +

40 +

41 +class SwiGLUActivation(nn.Module):

42 +

43 + def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:

44 + return x1 * nn.functional.silu(x2)

45 +

46 +

47 +class SambaYMLP(nn.Module):

48 + """Gated Linear Unit.

49 +

50 + Reference:

51 + Language Modeling with Gated Convolutional Networks.

52 + https://arxiv.org/pdf/1612.08083v3.pdf.

53 +

54 + """

55 +

56 + def __init__(self, config):

57 + super().__init__()

58 +

59 + self.config = config

60 + self.fc1 = nn.Linear(config.hidden_size,

61 + 2 * config.intermediate_size,

62 + bias=False)

63 + self.fc2 = nn.Linear(config.intermediate_size,

64 + config.hidden_size,

65 + bias=False)

66 +

67 + self.activation_fn = ACT2FN[config.hidden_act]

68 +

69 + def forward(self, hidden_states):

70 + y = self.fc1(hidden_states)

71 + gate, y = y.chunk(2, dim=-1)

72 + y = y * self.activation_fn(gate)

73 + return self.fc2(y)

74 +

75 +

76 +def get_virtual_engine():

77 + forward_context: ForwardContext = get_forward_context()

78 + return forward_context.virtual_engine

79 +

80 +

81 +class SambaYAttention(nn.Module):

82 +

83 + def __init__(self,

84 + config,

85 + layer_idx: Optional[int] = None,

86 + yoco_cross: bool = False,

87 + cache_config: Optional[CacheConfig] = None,

88 + prefix: str = ""):

89 + super().__init__()

90 + if layer_idx is None:

91 + logger.warning_once(

92 + f"Instantiating {self.__class__.__name__} without passing "

93 + "a `layer_idx` is not recommended and will lead to errors "

94 + "during the forward call if caching is used. Please make "

95 + "sure to provide a `layer_idx` when creating this class.")

96 + self.hidden_size = config.hidden_size

97 + self.num_heads = config.num_attention_heads

98 + self.head_dim = self.hidden_size // self.num_heads

99 + self.num_key_value_heads = config.num_key_value_heads

100 + self.yoco_cross = yoco_cross

101 +

102 + if (self.head_dim * self.num_heads) != self.hidden_size:

103 + raise ValueError("hidden_size must be divisible by num_heads "

104 + f"(got `hidden_size`: {self.hidden_size} and "

105 + f"`num_heads`: {self.num_heads}).")

106 +

107 + op_size = self.num_heads * self.head_dim + 2 * (

108 + self.num_key_value_heads * self.head_dim)

109 + self.out_proj = nn.Linear(self.num_heads * self.head_dim,

110 + self.hidden_size,

111 + bias=True)

112 + if yoco_cross:

113 + self.Wqkv = nn.Linear(self.hidden_size,

114 + self.num_heads * self.head_dim,

115 + bias=True)

116 + else:

117 + self.Wqkv = nn.Linear(self.hidden_size, op_size, bias=True)

118 +

119 + # disable sliding window for the second half of the model

120 + sliding_window = config.interleaved_sliding_window[layer_idx]

121 + if layer_idx >= config.num_hidden_layers // 2:

122 + assert sliding_window is None, \

123 + "sliding_window must be none for the second decoder"

124 + else:

125 + assert sliding_window is not None, \

126 + "sliding_window must be set for the first decoder"

127 +

128 + assert self.num_heads % 2 == 0, 'num_heads should be even'

129 + assert self.num_key_value_heads % 2 == 0, 'num_heads should be even'

130 +

131 + self.lambda_init = self.lambda_init_fn(layer_idx)

132 + self.lambda_q1 = nn.Parameter(

133 + torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,

134 + std=0.1))

135 + self.lambda_k1 = nn.Parameter(

136 + torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,

137 + std=0.1))

138 + self.lambda_q2 = nn.Parameter(

139 + torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,

140 + std=0.1))

141 + self.lambda_k2 = nn.Parameter(

142 + torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,

143 + std=0.1))

144 + self.subln = nn.RMSNorm(2 * self.head_dim,

145 + eps=1e-5,

146 + elementwise_affine=True)

147 +

148 + params = {

149 + 'differential_flash_attention_config': {

150 + 'lambda_init': self.lambda_init,

151 + 'lambda_q1': self.lambda_q1,

152 + 'lambda_k1': self.lambda_k1,

153 + 'lambda_q2': self.lambda_q2,

154 + 'lambda_k2': self.lambda_k2,

155 + "subln": self.subln,

156 + }

157 + }

158 +

159 + if yoco_cross:

160 + kv_shared_layer_index = config.num_hidden_layers // 2 + 1

161 + kv_sharing_target_layer_name = \

162 + f"model.layers.{kv_shared_layer_index}.self_attn.attn"

163 + else:

164 + kv_sharing_target_layer_name = None

165 +

166 + self.attn = Attention(

167 + self.num_heads,

168 + self.head_dim,

169 + self.head_dim**-0.5,

170 + num_kv_heads=self.num_key_value_heads,

171 + cache_config=cache_config,

172 + per_layer_sliding_window=sliding_window,

173 + prefix=f"{prefix}.attn",

174 + attn_type=AttentionType.DECODER,

175 + kv_sharing_target_layer_name=kv_sharing_target_layer_name,

176 + **params)

177 + assert self.attn.backend == _Backend.DIFFERENTIAL_FLASH_ATTN,\

178 + "DIFFERENTIAL_FLASH_ATTN required"

179 +

180 + def lambda_init_fn(self, depth):

181 + return 0.8 - 0.6 * math.exp(-0.3 * depth)

182 +

183 + def forward(

184 + self,

185 + hidden_states: torch.Tensor,

186 + ):

187 +

188 + if not self.yoco_cross: # need to generate kv-cache

189 + qkv = self.Wqkv(hidden_states)

190 + q, k, v = qkv.split([

191 + self.hidden_size, self.num_key_value_heads * self.head_dim,

192 + self.num_key_value_heads * self.head_dim

193 + ],

194 + dim=-1)

195 + attn_output = self.attn(q, k, v)

196 + else: # re-use the kv cache, full attention

197 + q = self.Wqkv(hidden_states)

198 + attn_output = self.attn(q, None, None)

199 + attn_output = attn_output.view(-1, self.num_heads * self.head_dim)

200 + return self.out_proj(attn_output)

201 +

202 +

203 +class Phi4Mamba(nn.Module):

204 +

205 + def __init__(

206 + self,

207 + d_model,

208 + d_state=16,

209 + d_conv=4,

210 + expand=2,

211 + dt_rank="auto",

212 + dt_min=0.001,

213 + dt_max=0.1,

214 + dt_init="random", # difference

215 + dt_scale=1.0, # difference

216 + dt_init_floor=1e-4,

217 + conv_bias=True,

218 + bias=False,

219 + use_fast_path=True, # Fused kernel options

220 + layer_idx=None,

221 + device=None,

222 + dtype=None,

223 + yoco_cross=False,

224 + yoco_kv=False,

225 + ):

226 + factory_kwargs = {"params_dtype": dtype} # difference

227 + super().__init__()

228 + self.yoco_cross = yoco_cross

229 + self.yoco_kv = yoco_kv

230 + self.d_model = d_model

231 + self.d_state = d_state

232 + self.d_conv = d_conv

233 + self.expand = expand

234 + self.d_inner = int(self.expand * self.d_model)

235 + self.dt_rank = math.ceil(self.d_model /

236 + 16) if dt_rank == "auto" else dt_rank

237 + self.use_fast_path = use_fast_path

238 + self.layer_idx = layer_idx

239 + self.swiGluActivation = SwiGLUActivation()

240 + if self.yoco_cross:

241 + self.in_proj = MergedColumnParallelLinear(self.d_model,

242 + [self.d_inner],

243 + bias=bias,

244 + **factory_kwargs)

245 + self.out_proj = RowParallelLinear(self.d_inner,

246 + self.d_model,

247 + bias=bias,

248 + **factory_kwargs)

249 + return

250 + self.conv1d = ColumnParallelLinear(

251 + input_size=d_conv,

252 + output_size=self.d_inner,

253 + bias=conv_bias,

254 + params_dtype=dtype,

255 + )

256 + # unsqueeze to fit conv1d weights shape into the linear weights shape.

257 + # Can't do this in `weight_loader` since it already exists in

258 + # `ColumnParallelLinear` and `set_weight_attrs`

259 + # doesn't allow to override it

260 + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)

261 +

262 + self.in_proj = MergedColumnParallelLinear(

263 + self.d_model,

264 + [self.d_inner] * 2,

265 + bias=bias,

266 + params_dtype=dtype,

267 + )

268 +

269 + # selective projection used to make dt, B and C input dependent

270 + self.x_proj = RowParallelLinear(

271 + self.d_inner,

272 + self.dt_rank + self.d_state * 2,

273 + bias=False,

274 + params_dtype=dtype,

275 + )

276 +

277 + # time step projection (discretization) -

278 + # In the forward we need to apply dt_proj without the bias,

279 + # as the bias is added in the selective scan kernel.

280 + self.dt_proj = ColumnParallelLinear(

281 + self.dt_rank,

282 + self.d_inner,

283 + bias=True,

284 + skip_bias_add=True,

285 + params_dtype=dtype,

286 + )

287 +

288 + # # D "skip" parameter

289 + # self.D = nn.Parameter(torch.ones(self.d_inner)) # Keep in fp32

290 + self.A = nn.Parameter(

291 + torch.empty(

292 + self.d_inner,

293 + self.d_state,

294 + dtype=torch.float32,

295 + ))

296 + self.D = nn.Parameter(torch.ones(self.d_inner, dtype=torch.float32))

297 +

298 + self.out_proj = RowParallelLinear(

299 + self.d_inner,

300 + self.d_model,

301 + bias=bias,

302 + input_is_parallel=True,

303 + params_dtype=dtype,

304 + )

305 + self.activation = "silu"

306 +

307 + def forward(self,

308 + hidden_states: torch.Tensor,

309 + attn_metadata: AttentionMetadata,

310 + mamba_cache_params: MambaCacheParams,

311 + yoco_key_values=None) -> torch.Tensor:

312 +

313 + if self.yoco_cross:

314 + out = self.in_proj(hidden_states)[0]

315 + out = self.swiGluActivation(yoco_key_values, out)

316 + out = self.out_proj(out)

317 + return out[0], yoco_key_values

318 +

319 + # 1. Gated MLP's linear projection

320 + # projected_states = self.in_proj(hidden_states)[0].transpose(-2, -1)

321 + projected_states = self.in_proj(

322 + hidden_states.to(self.in_proj.weight.dtype))[0].transpose(-2, -1)

323 + hidden_states, gate = projected_states.chunk(2, dim=-2)

324 +

325 + # 2. Convolution sequence transformation

326 + conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),

327 + self.conv1d.weight.size(2))

328 +

329 + if attn_metadata.query_start_loc is not None \

330 + and attn_metadata.context_lens_tensor is not None:

331 + # |---------- N-1 iteration --------|

332 + # |---------------- N iteration ---------------------|

333 + # |- tokenA -|......................|-- newTokens ---|

334 + # |---------- context_len ----------|

335 + # |-------------------- seq_len ---------------------|

336 + # |-- query_len ---|

337 + hidden_states = causal_conv1d_fn(

338 + hidden_states,

339 + conv_weights,

340 + self.conv1d.bias,

341 + activation=self.activation,

342 + conv_states=mamba_cache_params.conv_state,

343 + has_initial_state=attn_metadata.context_lens_tensor > 0,

344 + cache_indices=mamba_cache_params.state_indices_tensor,

345 + query_start_loc=attn_metadata.query_start_loc)

346 + else:

347 + hidden_states = causal_conv1d_update(

348 + hidden_states.transpose(0, 1),

349 + mamba_cache_params.conv_state,

350 + conv_weights,

351 + self.conv1d.bias,

352 + self.activation,

353 + conv_state_indices=mamba_cache_params.state_indices_tensor)

354 + hidden_states = hidden_states.transpose(0, 1)

355 +

356 + # 3. State Space Model sequence transformation

357 + # 3.a. input varying initialization of time_step, B and C

358 + ssm_parameters = self.x_proj(hidden_states.transpose(-2, -1))[0]

359 +

360 + time_step, B, C = torch.split(

361 + ssm_parameters,

362 + [self.dt_rank, self.d_state, self.d_state],

363 + dim=-1,

364 + )

365 +

366 + # Note that Jamba normalizes B, C, and time_step here but Mamba doesn't.

367 +

368 + discrete_time_step = self.dt_proj(time_step)[0].transpose(-2, -1)

369 + # 3.c perform the recurrence y ← SSM(A, B, C)(x)

370 + time_proj_bias = (self.dt_proj.bias.float() if hasattr(

371 + self.dt_proj, "bias") else None)

372 +

373 + if attn_metadata.query_start_loc is not None \

374 + and attn_metadata.context_lens_tensor is not None:

375 + scan_outputs = selective_scan_fn(

376 + hidden_states,

377 + mamba_cache_params.ssm_state,

378 + discrete_time_step,

379 + self.A,

380 + B.transpose(-2, -1),

381 + C.transpose(-2, -1),

382 + self.D.float(),

383 + # z,

384 + None if self.yoco_kv else gate,

385 + time_proj_bias,

386 + delta_softplus=True,

387 + cache_indices=mamba_cache_params.state_indices_tensor,

388 + has_initial_state=attn_metadata.context_lens_tensor > 0,

389 + query_start_loc=attn_metadata.query_start_loc)

390 + else:

391 + scan_outputs = selective_state_update(

392 + mamba_cache_params.ssm_state,

393 + hidden_states.transpose(0, 1),

394 + discrete_time_step.transpose(0, 1),

395 + self.A,

396 + B,

397 + C,

398 + self.D,

399 + # z

400 + # gate.transpose(0, 1),

401 + None if self.yoco_kv else gate.transpose(0, 1),

402 + time_proj_bias,

403 + dt_softplus=True,

404 + state_batch_indices=mamba_cache_params.state_indices_tensor)

405 + scan_outputs = scan_outputs.transpose(0, 1)

406 +

407 + # 4. Final linear projection

408 + if self.yoco_kv:

409 + # gate = gate.transpose(-1,-2).contiguous()

410 + yoco_key_values = scan_outputs.transpose(-2, -1)

411 + scan_outputs = self.swiGluActivation(scan_outputs, gate)

412 +

413 + contextualized_states = self.out_proj(scan_outputs.transpose(-2,

414 + -1))[0]

415 +

416 + return contextualized_states, yoco_key_values

417 +

418 +

419 +class SambaYDecoderLayer(nn.Module):

420 +

421 + def __init__(

422 + self,

423 + config,

424 + layer_idx,

425 + cache_config,

426 + prefix: str = "",

427 + ) -> None:

428 + super().__init__()

429 +

430 + self.config = config

431 + self.layer_idx = layer_idx

432 +

433 + self.mlp = SambaYMLP(config)

434 + self.input_layernorm = nn.LayerNorm(config.hidden_size,

435 + eps=config.layer_norm_eps)

436 +

437 + self.yoco_mb = False

438 + self.yoco_cross = False

439 + if layer_idx >= config.num_hidden_layers // 2:

440 + self.yoco_mb = True

441 + self.yoco_cross = (layer_idx

442 + >= (config.num_hidden_layers // 2 + 2))

443 + self.use_mamba = config.mb_per_layer > 0 and \

444 + layer_idx % config.mb_per_layer == 0

445 + if self.use_mamba:

446 + factory_kwargs = {"dtype": None}

447 + self.attn = Phi4Mamba(config.hidden_size,

448 + layer_idx=layer_idx,

449 + yoco_cross=self.yoco_cross,

450 + yoco_kv=self.yoco_mb,

451 + **factory_kwargs)

452 + else:

453 + self.attn = SambaYAttention(config,

454 + layer_idx=layer_idx,

455 + yoco_cross=self.yoco_cross,

456 + cache_config=cache_config,

457 + prefix=f"{prefix}.self_attn")

458 + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size,

459 + eps=config.layer_norm_eps)

460 +

461 + def forward(

462 + self,

463 + hidden_states: torch.Tensor,

464 + positions: torch.Tensor,

465 + attn_metadata: AttentionMetadata,

466 + mamba_cache_params: MambaCacheParams,

467 + ssm_output: Optional[torch.LongTensor] = None,

468 + ) -> Union[torch.Tensor, IntermediateTensors]:

469 + if self.use_mamba:

470 + assert mamba_cache_params is not None

471 + else:

472 + assert mamba_cache_params is None

473 +

474 + residual = hidden_states

475 + hidden_states = self.input_layernorm(

476 + hidden_states.to(dtype=self.input_layernorm.weight.dtype))

477 +

478 + if self.use_mamba:

479 + attn_outputs, ssm_output = self.attn(hidden_states,

480 + attn_metadata,

481 + mamba_cache_params,

482 + yoco_key_values=ssm_output)

483 + residual = residual.to(torch.float32)

484 + else:

485 + attn_outputs = self.attn(hidden_states, )

486 + hidden_states = residual + attn_outputs

487 + residual = hidden_states

488 + hidden_states = self.post_attention_layernorm(

489 + hidden_states.to(dtype=self.post_attention_layernorm.weight.dtype))

490 + hidden_states = self.mlp(hidden_states)

491 + hidden_states = residual + hidden_states

492 +

493 + return hidden_states, ssm_output

494 +

495 +

496 +class SambaYModel(nn.Module):

497 +

498 + def __init__(self,

499 + config,

500 + cache_config=None,

501 + quant_config=None,

502 + lora_config=None,

503 + prefix: str = "") -> None:

504 + super().__init__()

505 + self.config = config

506 + self.vocab_size = config.vocab_size

507 + self.embed_tokens = VocabParallelEmbedding(

508 + self.vocab_size,

509 + config.hidden_size,

510 + org_num_embeddings=config.vocab_size,

511 + )

512 +

513 + # Pipeline parallel is not supported since the second half of

514 + # the layers share the kv cache.

515 + if get_pp_group().world_size != 1:

516 + raise ValueError("Pipeline Parallel not supported")

517 +

518 + self.start_layer, self.end_layer, self.layers = make_layers(

519 + config.num_hidden_layers,

520 + lambda prefix: SambaYDecoderLayer(config,

521 + int(prefix.split('.')[-1]),

522 + cache_config,

523 + prefix=prefix),

524 + prefix=f"{prefix}.layers")

525 + self.final_layernorm = nn.LayerNorm(config.hidden_size,

526 + eps=config.layer_norm_eps)

527 +

528 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

529 + return self.embed_tokens(input_ids)

530 +

531 + def forward(

532 + self,

533 + input_ids: Optional[torch.Tensor],

534 + positions: torch.Tensor,

535 + attn_metadata: AttentionMetadata,

536 + mamba_cache_params: MambaCacheParams,

537 + intermediate_tensors: Optional[IntermediateTensors] = None,

538 + inputs_embeds: Optional[torch.Tensor] = None,

539 + ) -> Union[torch.Tensor, IntermediateTensors]:

540 +

541 + if get_pp_group().is_first_rank:

542 + if inputs_embeds is not None:

543 + hidden_states = inputs_embeds

544 + else:

545 + hidden_states = self.get_input_embeddings(input_ids)

546 + else:

547 + assert intermediate_tensors is not None

548 + hidden_states = intermediate_tensors["hidden_states"]

549 +

550 + mamba_state_idx = 0

551 + ssm_output = None

552 + for i in range(self.start_layer, self.end_layer):

553 + layer = self.layers[i]

554 + if i == self.config.num_hidden_layers // 2 + 2:

555 + # profile run

556 + kv_cache_idx = self.config.num_hidden_layers // 2 + 1

557 + cache_layer = self.layers[kv_cache_idx]

558 + kv_cache = cache_layer.attn.attn.kv_cache

559 + if kv_cache[0].numel() == 0:

560 + break

561 +

562 + # Starting from this layer, we do not need to calculate

563 + # the kv cache since we reuse the kv cache from last layer.

564 + # If in prefill phase, we can <s>prune></s> truncate

565 + # the hidden state to save computation cost.

566 + if attn_metadata.prefill_metadata and not envs.VLLM_USE_V1:

567 + selected_token_indices = torch.cumsum(

568 + attn_metadata.seq_lens_tensor, dim=0) - 1

569 + hidden_states = hidden_states.index_select(

570 + 0, selected_token_indices)

571 + ssm_output = ssm_output.index_select(

572 + 0, selected_token_indices)

573 +

574 + if layer.use_mamba:

575 + if i < self.config.num_hidden_layers // 2 or \

576 + not layer.yoco_cross:

577 + mamba_cache = mamba_cache_params.at_layer_idx(

578 + mamba_state_idx)

579 + mamba_state_idx += 1

580 + else:

581 + mamba_cache = mamba_cache_params.at_layer_idx(

582 + mamba_state_idx - 1)

583 +

584 + hidden_states, ssm_output = layer(hidden_states,

585 + positions,

586 + attn_metadata,

587 + mamba_cache,

588 + ssm_output=ssm_output)

589 + else:

590 + hidden_states, ssm_output = layer(

591 + hidden_states,

592 + positions,

593 + attn_metadata,

594 + None, # mamba_cache_params

595 + ssm_output=ssm_output)

596 +

597 + hidden_states = self.final_layernorm(

598 + hidden_states.to(dtype=self.final_layernorm.weight.dtype))

599 + return hidden_states

600 +

601 +

602 +class Phi4FlashForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):

603 +

604 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

605 + config = vllm_config.model_config.hf_config

606 + cache_config = vllm_config.cache_config

607 + lora_config = vllm_config.lora_config

608 + quant_config = vllm_config.quant_config

609 + scheduler_config = vllm_config.scheduler_config

610 + self.compilation_config = vllm_config.compilation_config

611 + self.vllm_config = vllm_config

612 + # Prefix caching and chunked prefill is not supported for this model.

613 + assert not cache_config.enable_prefix_caching, \

614 + "Phi4flash currently does not support prefix caching"

615 + assert not scheduler_config.chunked_prefill_enabled, \

616 + "Phi4Flash currently does not support prefix caching"

617 + super().__init__()

618 + self.config = config

619 + self.model_config = vllm_config.model_config

620 + self.scheduler_config = scheduler_config

621 + self.model = SambaYModel(config,

622 + cache_config=cache_config,

623 + prefix=maybe_prefix(prefix, "model"))

624 + self.unpadded_vocab_size = config.vocab_size

625 + if lora_config:

626 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size

627 + self.lm_head = ParallelLMHead(

628 + self.unpadded_vocab_size,

629 + config.hidden_size,

630 + org_num_embeddings=config.vocab_size,

631 + padding_size=(

632 + DEFAULT_VOCAB_PADDING_SIZE

633 + # We need bigger padding if using lora for kernel

634 + # compatibility

635 + if not lora_config else lora_config.lora_vocab_padding_size),

636 + quant_config=quant_config,

637 + )

638 + self.embedding_bias = None

639 + # Used to track and store by the Mamba cache between steps.

640 + self.mamba_cache: Optional[MambaCacheManager] = None

641 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

642 + config.vocab_size,

643 + logits_as_input=False)

644 + self.sampler = get_sampler()

645 +

646 + def forward(

647 + self,

648 + input_ids: torch.Tensor,

649 + positions: torch.Tensor,

650 + intermediate_tensors: Optional[IntermediateTensors] = None,

651 + inputs_embeds: Optional[torch.Tensor] = None,

652 + **kwargs,

653 + ) -> Union[torch.Tensor, IntermediateTensors]:

654 + if self.mamba_cache is None:

655 + num_mamba_layers = self.config.num_hidden_layers \

656 + // 2 // self.config.mb_per_layer + 1

657 + self.mamba_cache = MambaCacheManager(

658 + self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,

659 + *self._get_mamba_cache_shape())

660 + mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

661 +

662 + attn_metadata = get_forward_context().attn_metadata

663 + # input_ids and hidden_states isn't a one-to-one mapping in prefill

664 + # stage due to YOCO optimization.

665 + hidden_states = self.model(input_ids, positions, attn_metadata,

666 + mamba_cache_params, intermediate_tensors,

667 + inputs_embeds)

668 + return hidden_states

669 +

670 + def _get_mamba_cache_shape(

671 + self

672 + ) -> tuple[Optional[tuple[int, int]], Optional[tuple[int, int]]]:

673 + world_size = get_tensor_model_parallel_world_size()

674 + hidden_size = self.config.hidden_size

675 + mamba_expand = self.config.mamba_expand # 2

676 + mamba_d_conv = self.config.mamba_d_conv # 4

677 + mamba_d_state = self.config.mamba_d_state # 16

678 + conv_state_shape = (

679 + mamba_expand * hidden_size // world_size,

680 + mamba_d_conv - 1,

681 + )

682 + temporal_state_shape = (

683 + mamba_expand * hidden_size // world_size,

684 + mamba_d_state,

685 + )

686 + return conv_state_shape, temporal_state_shape

687 +

688 + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):

689 + return self.mamba_cache.copy_inputs_before_cuda_graphs(

690 + input_buffers, **kwargs)

691 +

692 + def get_seqlen_agnostic_capture_inputs(self, batch_size: int):

693 + return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)

694 +

695 + def compute_logits(

696 + self,

697 + hidden_states: torch.Tensor,

698 + sampling_metadata: SamplingMetadata,

699 + ) -> Optional[torch.Tensor]:

700 + # If the shape is the same, it means that we have already

701 + # prune hidden states manually.

702 + prune_hidden_states = hidden_states.size(

703 + 0) != sampling_metadata.selected_token_indices.size(0)

704 + processed_logits = self.logits_processor(

705 + self.lm_head,

706 + hidden_states,

707 + sampling_metadata,

708 + self.embedding_bias,

709 + prune_hidden_states=prune_hidden_states)

710 + return processed_logits

711 +

712 + def sample(

713 + self,

714 + logits: torch.Tensor,

715 + sampling_metadata: SamplingMetadata,

716 + ) -> Optional[SamplerOutput]:

717 + next_tokens = self.sampler(logits, sampling_metadata)

718 + return next_tokens

719 +

720 + def load_weights(

721 + self,

722 + weights: Iterable[tuple[str, torch.Tensor]],

723 + ):

724 + weights = {name: weight for name, weight in weights}

725 + adjusted_weights = {}

726 + for name, weight in weights.items():

727 + if "A_log" in name:

728 + name = name.replace("A_log", "A")

729 + weight = -torch.exp(weight.float())

730 + if "inner_cross_attn." in name:

731 + name = name.replace("inner_cross_attn.", "")

732 + adjusted_weights[name] = weight

733 + adjusted_weights["lm_head.weight"] = weights[

734 + "model.embed_tokens.weight"]

735 + loaded_params: set[str] = set()

736 + for name, param in self.named_parameters():

737 + weight = adjusted_weights.get(name)

738 + if weight is not None and weight.shape != param.shape:

739 + logger.warning("Shape mismatch: %s %s %s", name, weight.shape,

740 + param.shape)

741 + loaded_params.add(name)

742 + missing_keys, unexpected_keys = self.load_state_dict(adjusted_weights,

743 + strict=False)

744 + assert len(unexpected_keys) == 0, f"Unexpected keys: {unexpected_keys}"

745 + assert len(missing_keys) == 0, f"Missing keys: {missing_keys}"

746 + return loaded_params

vllm/model_executor/models/registry.py

@@ -110,6 +110,7 @@ _TEXT_GENERATION_MODELS = {

110 "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),

111 "Phi3SmallForCausalLM": ("phi3_small", "Phi3SmallForCausalLM"),

112 "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),

113 + "Phi4FlashForCausalLM": ("phi4flash", "Phi4FlashForCausalLM"),

114 "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),

115 "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),

116 "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),

vllm/platforms/cuda.py

@@ -316,6 +316,10 @@ class CudaPlatformBase(Platform):

316 logger.info("Using DualChunkFlashAttention backend.")

317 return ("vllm.attention.backends.dual_chunk_flash_attn."

318 "DualChunkFlashAttentionBackend")

319 + elif selected_backend == _Backend.DIFFERENTIAL_FLASH_ATTN:

320 + logger.info("Using DifferentialFlashAttention backend.")

321 + return ("vllm.attention.backends.differential_flash_attn."

322 + "DifferentialFlashAttentionBackend")

323 elif selected_backend == _Backend.FLASH_ATTN:

324 pass

325 elif selected_backend:

vllm/platforms/interface.py

@@ -60,6 +60,7 @@ class _Backend(enum.Enum):

60 IPEX = enum.auto()

61 BLOCK_SPARSE_FLASH_ATTN = enum.auto()

62 DUAL_CHUNK_FLASH_ATTN = enum.auto()

63 + DIFFERENTIAL_FLASH_ATTN = enum.auto()

64 NO_ATTENTION = enum.auto()

65 FLEX_ATTENTION = enum.auto()

vllm/utils/__init__.py

+15 -3

@@ -2888,8 +2888,9 @@ def get_mp_context():

2888

2889

2890 def bind_kv_cache(

2891 - ctx: dict[str, Any],

2892 - kv_cache: list[list[torch.Tensor]], # [virtual_engine][layer_index]

2891 + ctx: dict[str, Any],

2892 + kv_cache: list[list[torch.Tensor]], # [virtual_engine][layer_index]

2893 + shared_kv_cache_layers: Optional[dict[str, str]] = None

2894 ) -> None:

2895 # Bind the kv_cache tensor to Attention modules, similar to

2896 # ctx[layer_name].kv_cache[ve]=kv_cache[ve][extract_layer_index(layer_name)]

@@ -2901,12 +2902,17 @@ def bind_kv_cache(

2902 # attention of the same layer (e.g., bart's decoder.layers.1.self_attn

2903 # and decoder.layers.1.encoder_attn) is mapped to the same kv cache

2904 # tensor

2905 + # 5. Some models have attention layers that share kv cache with previous

2906 + # layers, this is specified through shared_kv_cache_layers

2907 + if shared_kv_cache_layers is None:

2908 + shared_kv_cache_layers = {}

2909 from vllm.attention import AttentionType

2910 from vllm.model_executor.models.utils import extract_layer_index

2911 layer_need_kv_cache = [

2912 layer_name for layer_name in ctx

2913 if (hasattr(ctx[layer_name], 'attn_type') and ctx[layer_name].attn_type

2909 - in (AttentionType.DECODER, AttentionType.ENCODER_DECODER))

2914 + in (AttentionType.DECODER, AttentionType.ENCODER_DECODER)) \

2915 + and ctx[layer_name].kv_sharing_target_layer_name is None

2916 ]

2917 layer_index_sorted = sorted(

2918 set(

@@ -2919,6 +2925,12 @@ def bind_kv_cache(

2925 assert len(forward_ctx.kv_cache) == len(kv_cache)

2926 for ve, ve_kv_cache in enumerate(kv_cache):

2927 forward_ctx.kv_cache[ve] = ve_kv_cache[kv_cache_idx]

2928 + if shared_kv_cache_layers is not None:

2929 + for layer_name, target_layer_name in shared_kv_cache_layers.items():

2930 + assert extract_layer_index(target_layer_name) < \

2931 + extract_layer_index(layer_name), \

2932 + "v0 doesn't support interleaving kv sharing"

2933 + ctx[layer_name].kv_cache = ctx[target_layer_name].kv_cache

2934

2935

2936 def run_method(obj: Any, method: Union[str, bytes, Callable], args: tuple[Any],

vllm/worker/model_runner.py

@@ -1112,6 +1112,10 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):

1112 (self.max_batchsize_to_capture, self.get_max_block_per_batch()),

1113 dtype=np.int32)

1114

1115 + self.cross_layer_shared_graph_block_tables = np.zeros(

1116 + (self.max_batchsize_to_capture, self.get_max_block_per_batch()),

1117 + dtype=np.int32)

1118 +

1119 # Attention-free but stateful models like Mamba need a placeholder attn

1120 # backend, as the attention metadata is needed to manage internal state.

1121 # However we must bypass attention selection altogether for some models

vllm/worker/worker.py

+24 -2

@@ -9,7 +9,8 @@ import torch

9 import torch.distributed

11 import vllm.envs as envs

12 -from vllm.config import VllmConfig

12 +from vllm.attention.layer import Attention

13 +from vllm.config import VllmConfig, get_layers_from_vllm_config

14 from vllm.device_allocator.cumem import CuMemAllocator

15 from vllm.distributed import (ensure_model_parallel_initialized,

16 init_distributed_environment,

@@ -345,8 +346,29 @@ class Worker(LocalOrDistributedWorkerBase):

346 self.cache_engine[ve].gpu_cache

347 for ve in range(self.parallel_config.pipeline_parallel_size)

348 ]

349 +

350 + # Layer pairings for cross-layer KV sharing.

351 + # If an Attention layer `layer_name` is in the keys of this dict, it

352 + # means this layer will perform attention using the keys and values

353 + # from the KV cache of `shared_kv_cache_layers[layer_name]`.

354 + shared_kv_cache_layers: dict[str, str] = {}

355 +

356 + attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)

357 +

358 + for layer_name, attn_module in attn_layers.items():

359 + if (kv_tgt_layer :=

360 + attn_module.kv_sharing_target_layer_name) is not None:

361 + # The layer doesn't need its own KV cache and will use that of

362 + # the target layer. We skip creating a KVCacheSpec for it, so

363 + # that KV cache management logic will act as this layer does

364 + # not exist, and doesn't allocate KV cache for the layer. This

365 + # enables the memory saving of cross-layer kv sharing, allowing

366 + # a given amount of memory to accommodate longer context lengths

367 + # or enable more requests to be processed simultaneously.

368 + shared_kv_cache_layers[layer_name] = kv_tgt_layer

369 +

370 bind_kv_cache(self.compilation_config.static_forward_context,

349 - self.gpu_cache)

371 + self.gpu_cache, shared_kv_cache_layers)

372

373 def _warm_up_model(self) -> None:

374 # warm up sizes that are not in cudagraph capture sizes,