@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3f52fa5aa2

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

[Model] Add support for openPangu moe model (#28775)

Signed-off-by: yuantao <2422264527@qq.com> Signed-off-by: yt0428 <51468697+yt0428@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

yt0428 committed Dec 31, 2025 at 00:11 UTC 3f52fa5aa26d65c12ea81488986e1b8ecc4dcdc6

11 files changed +1136 -7

docs/models/supported_models.md

@@ -433,6 +433,7 @@ th {

435 | `PanguEmbeddedForCausalLM` |openPangu-Embedded-7B | `FreedomIntelligence/openPangu-Embedded-7B-V1.1` | ✅︎ | ✅︎ |

436 +| `PanguProMoEV2ForCausalLM` |openpangu-pro-moe-v2 | | ✅︎ | ✅︎ |

437 | `PanguUltraMoEForCausalLM` |openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ |

438 | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |

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

tests/models/registry.py

@@ -402,6 +402,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

402 "PanguEmbeddedForCausalLM": _HfExamplesInfo(

403 "FreedomIntelligence/openPangu-Embedded-7B-V1.1", trust_remote_code=True

404 ),

405 + "PanguProMoEV2ForCausalLM": _HfExamplesInfo(

406 + "",

407 + trust_remote_code=True,

408 + is_available_online=False,

409 + ),

410 "PanguUltraMoEForCausalLM": _HfExamplesInfo(

411 "FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1",

412 trust_remote_code=True,

vllm/attention/backends/registry.py

@@ -42,6 +42,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):

42 """

44 FLASH_ATTN = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"

45 + FLASH_ATTN_DIFFKV = (

46 + "vllm.v1.attention.backends.flash_attn_diffkv.FlashAttentionDiffKVBackend"

47 + )

48 TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend"

49 ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"

50 ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"

vllm/attention/layer.py

+11 -4

@@ -170,6 +170,7 @@ class Attention(nn.Module, AttentionLayerBase):

170 attn_type: str = AttentionType.DECODER,

171 kv_sharing_target_layer_name: str | None = None,

172 attn_backend: type[AttentionBackend] | None = None,

173 + head_size_v: int | None = None,

174 **extra_impl_args,

175 ) -> None:

176 """

@@ -217,6 +218,7 @@ class Attention(nn.Module, AttentionLayerBase):

218

219 self.num_heads = num_heads

220 self.head_size = head_size

221 + self.head_size_v = self.head_size if head_size_v is None else head_size_v

222 self.num_kv_heads = num_kv_heads

223 self.sliding_window = sliding_window

224 self.has_sink = extra_impl_args.get("sinks") is not None

@@ -274,8 +276,7 @@ class Attention(nn.Module, AttentionLayerBase):

276 kv_sharing_target_layer_name,

277 **extra_impl_args,

278 )

277 - backend_name = self.attn_backend.get_name()

278 - self.backend = AttentionBackendEnum.__members__.get(backend_name)

279 + self.backend = AttentionBackendEnum[self.attn_backend.get_name()]

280 self.dtype = dtype

281

282 # For cuda-alike (CUDA and ROCM) and cpu platforms, we control how

@@ -355,6 +356,10 @@ class Attention(nn.Module, AttentionLayerBase):

356 query, _ = self.query_quant(query, self._q_scale)

357

358 if self.use_output:

359 + if output_shape is None:

360 + output_shape = torch.Size(

361 + (*query.shape[:-1], self.num_heads * self.head_size_v)

362 + )

363 output_shape = output_shape if output_shape is not None else query.shape

364 output = torch.empty(output_shape, dtype=output_dtype, device=query.device)

365 hidden_size = output_shape[-1]

@@ -362,11 +367,11 @@ class Attention(nn.Module, AttentionLayerBase):

367 # NOTE(woosuk): We do this outside the custom op to minimize the

368 # CPU overheads from the non-CUDA-graph regions.

369 query = query.view(-1, self.num_heads, self.head_size)

365 - output = output.view(-1, self.num_heads, self.head_size)

370 + output = output.view(-1, self.num_heads, self.head_size_v)

371 if key is not None:

372 key = key.view(-1, self.num_kv_heads, self.head_size)

373 if value is not None:

369 - value = value.view(-1, self.num_kv_heads, self.head_size)

374 + value = value.view(-1, self.num_kv_heads, self.head_size_v)

375 if self.use_direct_call:

376 forward_context: ForwardContext = get_forward_context()

377 attn_metadata = forward_context.attn_metadata

@@ -452,6 +457,7 @@ class Attention(nn.Module, AttentionLayerBase):

457 block_size=block_size,

458 num_kv_heads=self.num_kv_heads,

459 head_size=self.head_size,

460 + head_size_v=self.head_size_v,

461 dtype=self.kv_cache_torch_dtype,

462 )

463

@@ -794,6 +800,7 @@ def unified_attention_with_output(

800 output_block_scale: torch.Tensor | None = None,

801 ) -> None:

802 attn_metadata, self, kv_cache = get_attention_context(layer_name)

803 +

804 self.impl.forward(

805 self,

806 query,

vllm/attention/layers/static_sink_attention.py new

+254

@@ -0,0 +1,254 @@

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

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

3 +import functools

4 +

5 +import torch

6 +

7 +from vllm.attention.backends.abstract import (

8 + AttentionBackend,

9 + AttentionMetadata,

10 + AttentionType,

11 +)

12 +from vllm.attention.layer import Attention

13 +from vllm.attention.ops.triton_reshape_and_cache_flash import (

14 + triton_reshape_and_cache_flash_diffkv,

15 +)

16 +from vllm.attention.selector import get_attn_backend

17 +from vllm.config import CacheConfig, VllmConfig

18 +from vllm.forward_context import ForwardContext, get_forward_context

19 +from vllm.logger import init_logger

20 +from vllm.model_executor.custom_op import CustomOp

21 +from vllm.utils.math_utils import cdiv

22 +from vllm.utils.torch_utils import direct_register_custom_op

23 +from vllm.v1.attention.backends.utils import (

24 + CommonAttentionMetadata,

25 + subclass_attention_backend,

26 +)

27 +from vllm.v1.kv_cache_interface import (

28 + AttentionSpec,

29 + KVCacheSpec,

30 + SinkFullAttentionSpec,

31 +)

32 +

33 +logger = init_logger(__name__)

34 +

35 +

36 +@functools.lru_cache

37 +def create_static_sink_attention_backend(

38 + underlying_attn_backend: type[AttentionBackend],

39 + sink_len: int = 0,

40 +) -> type[AttentionBackend]:

41 + prefix = "StaticSink_"

42 + underlying_builder = underlying_attn_backend.get_builder_cls()

43 +

44 + class StaticSinkAttentionBuilder(underlying_builder): # type: ignore

45 + def __init__(

46 + self,

47 + kv_cache_spec: AttentionSpec,

48 + layer_names: list[str],

49 + vllm_config: VllmConfig,

50 + device: torch.device,

51 + ):

52 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)

53 + model_config = vllm_config.model_config

54 + scheduler_config = vllm_config.scheduler_config

55 + self.sink_len = sink_len

56 + self.block_size = vllm_config.cache_config.block_size

57 + self.num_sink_blocks = self.sink_len // vllm_config.cache_config.block_size

58 + self.max_num_blocks = cdiv(

59 + model_config.max_model_len, vllm_config.cache_config.block_size

60 + )

61 + self.block_table_with_sink = torch.zeros(

62 + (

63 + scheduler_config.max_num_seqs,

64 + self.max_num_blocks + self.num_sink_blocks,

65 + ),

66 + device=device,

67 + dtype=torch.int32,

68 + )

69 + self.block_table_with_sink[:, : self.num_sink_blocks] = torch.arange(

70 + 1,

71 + self.num_sink_blocks + 1,

72 + device=device,

73 + dtype=torch.int32,

74 + )

75 +

76 + def build(

77 + self,

78 + common_prefix_len: int,

79 + common_attn_metadata: CommonAttentionMetadata,

80 + fast_build: bool = False,

81 + ) -> AttentionMetadata:

82 + common_attn_metadata.seq_lens[:] = (

83 + common_attn_metadata.seq_lens + self.sink_len

84 + )

85 + common_attn_metadata.seq_lens[

86 + common_attn_metadata.seq_lens == self.sink_len

87 + ] = 0

88 + common_attn_metadata.max_seq_len = (

89 + common_attn_metadata.max_seq_len + self.sink_len

90 + )

91 + max_num_blocks = cdiv(common_attn_metadata.max_seq_len, self.block_size)

92 + num_reqs = common_attn_metadata.num_reqs

93 + self.block_table_with_sink[

94 + :num_reqs, self.num_sink_blocks : self.num_sink_blocks + max_num_blocks

95 + ] = common_attn_metadata.block_table_tensor[:, :max_num_blocks]

96 + common_attn_metadata.block_table_tensor = self.block_table_with_sink[

97 + :num_reqs

98 + ]

99 +

100 + return super().build(common_prefix_len, common_attn_metadata, fast_build)

101 +

102 + attn_backend = subclass_attention_backend(

103 + name_prefix=prefix,

104 + attention_backend_cls=underlying_attn_backend,

105 + builder_cls=StaticSinkAttentionBuilder,

106 + )

107 +

108 + return attn_backend

109 +

110 +

111 +@CustomOp.register("static_sink_attention")

112 +class StaticSinkAttention(Attention, CustomOp):

113 + """

114 + Attention with static sink tokens

115 + """

116 +

117 + def __init__(

118 + self,

119 + num_heads: int,

120 + head_size: int,

121 + scale: float,

122 + sink_len: int,

123 + attn_backend: type[AttentionBackend] | None = None,

124 + cache_config: CacheConfig | None = None,

125 + **kwargs,

126 + ):

127 + dtype = torch.get_default_dtype()

128 +

129 + if cache_config is not None:

130 + kv_cache_dtype = cache_config.cache_dtype

131 + block_size = cache_config.block_size

132 + else:

133 + kv_cache_dtype = "auto"

134 + block_size = 16

135 +

136 + if attn_backend is not None:

137 + underlying_attn_backend = attn_backend

138 + else:

139 + underlying_attn_backend = get_attn_backend(

140 + head_size, dtype, kv_cache_dtype, block_size

141 + )

142 + attn_backend = create_static_sink_attention_backend(

143 + underlying_attn_backend,

144 + sink_len=sink_len,

145 + )

146 + Attention.__init__(

147 + self=self,

148 + num_heads=num_heads,

149 + head_size=head_size,

150 + scale=scale,

151 + cache_config=cache_config,

152 + attn_backend=attn_backend,

153 + **kwargs,

154 + )

155 + CustomOp.__init__(self)

156 +

157 + self.sink_len = sink_len

158 + self.block_size = block_size

159 + self.sink_populated = False

160 + self.sink_key = None

161 + self.sink_value = None

162 +

163 + def update_sink_kv(self, sink_key, sink_value) -> None:

164 + self.sink_key = sink_key

165 + self.sink_value = sink_value

166 +

167 + def forward_native(

168 + self,

169 + query: torch.Tensor,

170 + key: torch.Tensor,

171 + value: torch.Tensor,

172 + output_shape: torch.Size | None = None,

173 + ) -> torch.Tensor:

174 + assert self.sink_key is not None and self.sink_value is not None, (

175 + "sink_key and sink_value have not been prepared"

176 + )

177 + if not self.sink_populated:

178 + forward_context: ForwardContext = get_forward_context()

179 + self_kv_cache = self.kv_cache[forward_context.virtual_engine]

180 + torch.ops.vllm.maybe_populate_sink(self_kv_cache, self.layer_name)

181 +

182 + return super().forward(query, key, value, output_shape)

183 +

184 + def forward_cuda(

185 + self,

186 + query: torch.Tensor,

187 + key: torch.Tensor,

188 + value: torch.Tensor,

189 + output_shape: torch.Size | None = None,

190 + ) -> torch.Tensor:

191 + return self.forward_native(query, key, value, output_shape)

192 +

193 + def forward(self, *args, **kwargs):

194 + return self._forward_method(*args, **kwargs)

195 +

196 + def populate_sink_kv(self, self_kv_cache):

197 + sink_kv_slot_mapping = torch.arange(

198 + self.block_size,

199 + self.sink_len + self.block_size,

200 + device=torch.cuda.current_device(),

201 + dtype=torch.long,

202 + )

203 + triton_reshape_and_cache_flash_diffkv(

204 + self.sink_key,

205 + self.sink_value,

206 + self_kv_cache,

207 + sink_kv_slot_mapping,

208 + self.kv_cache_dtype,

209 + self._k_scale,

210 + self._v_scale,

211 + )

212 + # We only populate the sink_key and sink_value once

213 + self.sink_populated = True

214 +

215 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

216 + # Block size may get updated after model loading, refresh it

217 + block_size = vllm_config.cache_config.block_size

218 + # Should not be called for enc-dec or encoder-only attention.

219 + assert self.attn_type == AttentionType.DECODER

220 +

221 + return SinkFullAttentionSpec(

222 + block_size=block_size,

223 + num_kv_heads=self.num_kv_heads,

224 + head_size=self.head_size,

225 + head_size_v=self.head_size_v,

226 + sink_len=self.sink_len,

227 + dtype=self.kv_cache_torch_dtype,

228 + )

229 +

230 +

231 +def maybe_populate_sink(

232 + self_kv_cache: torch.Tensor,

233 + layer_name: str,

234 +) -> None:

235 + forward_context: ForwardContext = get_forward_context()

236 + self = forward_context.no_compile_layers[layer_name]

237 + if self.sink_populated or self_kv_cache.numel() == 0:

238 + return

239 + self.populate_sink_kv(self_kv_cache)

240 +

241 +

242 +def maybe_populate_sink_fake(

243 + self_kv_cache: torch.Tensor,

244 + layer_name: str,

245 +) -> None:

246 + return

247 +

248 +

249 +direct_register_custom_op(

250 + op_name="maybe_populate_sink",

251 + op_func=maybe_populate_sink,

252 + mutates_args=["self_kv_cache"],

253 + fake_impl=maybe_populate_sink_fake,

254 +)

vllm/attention/ops/triton_reshape_and_cache_flash.py

+171

@@ -182,3 +182,174 @@ def triton_reshape_and_cache_flash(

182 num_warps=num_warps,

183 num_stages=num_stages,

184 )

185 +

186 +

187 +@triton.jit

188 +def reshape_and_cache_kernel_flash_diffkv(

189 + key_ptr, # [num_tokens, num_heads, head_size]

190 + value_ptr, # [num_tokens, num_heads, head_size_v]

191 + kv_cache_ptr, # [num_blocks, block_size, num_heads, head_size + head_size_v]

192 + slot_mapping_ptr, # [num_tokens]

193 + k_scale, # float32

194 + v_scale, # float32

195 + # strides

196 + key_stride: tl.int64,

197 + value_stride: tl.int64,

198 + block_stride: tl.int64,

199 + page_stride: tl.int64,

200 + num_heads: tl.constexpr,

201 + head_size_k: tl.constexpr,

202 + head_size_v: tl.constexpr,

203 + block_size: tl.constexpr,

204 + # FP8 flags

205 + FP8_KV_CACHE: tl.constexpr,

206 + # tune parameters

207 + TILE_SIZE: tl.constexpr,

208 +):

209 + token_idx = tl.program_id(axis=0)

210 + slot_idx = tl.load(slot_mapping_ptr + token_idx).to(tl.int64)

211 + if slot_idx < 0:

212 + # Padding token that should be ignored.

213 + return

214 +

215 + tile_i = tl.program_id(axis=1)

216 + tile_offs = tl.arange(0, TILE_SIZE)

217 +

218 + block_idx = slot_idx // block_size

219 + block_offset = slot_idx % block_size

220 +

221 + src_key_idx = token_idx * key_stride + tile_i * head_size_k

222 + src_value_idx = token_idx * value_stride + tile_i * head_size_v

223 +

224 + tgt_idx = (

225 + block_idx * block_stride

226 + + block_offset * page_stride

227 + + tile_i * (head_size_k + head_size_v)

228 + )

229 +

230 + # [TILE_SIZE]

231 + key_load = tl.load(key_ptr + src_key_idx + tile_offs, mask=tile_offs < head_size_k)

232 + if FP8_KV_CACHE:

233 + # tl.store will do the correct implicit cast to fp8,

234 + # based on the key_cache_ptr.dtype.element_ty

235 + key_tile = key_load if key_load.dtype.is_fp8() else key_load / tl.load(k_scale)

236 + else:

237 + key_tile = key_load

238 +

239 + # [TILE_SIZE]

240 + value_load = tl.load(

241 + value_ptr + src_value_idx + tile_offs, mask=tile_offs < head_size_v

242 + )

243 + if FP8_KV_CACHE:

244 + if value_load.dtype.is_fp8():

245 + value_tile = value_load

246 + else:

247 + # tl.store will do the correct implicit cast to fp8,

248 + # based on the value_cache_ptr.dtype.element_ty

249 + value_tile = value_load / tl.load(v_scale)

250 + else:

251 + value_tile = value_load

252 +

253 + tl.store(

254 + kv_cache_ptr + tgt_idx + tile_offs,

255 + key_tile,

256 + mask=tile_offs < head_size_k,

257 + )

258 + tl.store(

259 + kv_cache_ptr + tgt_idx + head_size_k + tile_offs,

260 + value_tile,

261 + mask=tile_offs < head_size_v,

262 + )

263 + return

264 +

265 +

266 +def triton_reshape_and_cache_flash_diffkv(

267 + key: torch.Tensor, # [num_tokens, num_heads, head_size]

268 + value: torch.Tensor, # [num_tokens, num_heads, head_size_v]

269 + # [num_blocks, block_size, num_heads, head_size + head_size_v]

270 + kv_cache: torch.Tensor,

271 + slot_mapping: torch.Tensor, # [num_tokens]

272 + kv_cache_dtype: str, # "auto", "fp8"

273 + k_scale: torch.Tensor, # float32

274 + v_scale: torch.Tensor, # float32

275 +):

276 + num_heads = key.shape[1]

277 + head_size_k = key.shape[2]

278 + head_size_v = value.shape[2]

279 + block_size = kv_cache.shape[1]

280 +

281 + k_stride = key.stride()[0]

282 + v_stride = value.stride()[0]

283 + block_stride = kv_cache.stride()[0]

284 + page_stride = kv_cache.stride()[1]

285 +

286 + assert kv_cache_dtype == "auto" or kv_cache_dtype.startswith("fp8"), (

287 + f"unsupported kv_cache_dtype (str), got {kv_cache_dtype}."

288 + )

289 + kv_cache_torch_dtype = (

290 + current_platform.fp8_dtype()

291 + if kv_cache_dtype.startswith("fp8")

292 + else kv_cache.dtype

293 + )

294 +

295 + if kv_cache.dtype != kv_cache_torch_dtype and kv_cache_dtype.startswith("fp8"):

296 + # to avoid erounous implicit cast in triton kernel (tl.store to uint8)

297 + # (e.g. explicit cast to fp8e4m3fnuz is not supported in triton 3.4)

298 + kv_cache = kv_cache.view(kv_cache_torch_dtype)

299 + assert kv_cache_dtype != torch.uint8, (

300 + "explicit fp8 cast and store to "

301 + "uint8 is not supported by triton reshape_and_cache_flash_diffkv"

302 + )

303 +

304 + FP8_KV_CACHE = kv_cache_dtype.startswith("fp8")

305 + assert (not FP8_KV_CACHE) or kv_cache_torch_dtype in [

306 + torch.float8_e4m3fn,

307 + torch.float8_e5m2,

308 + torch.uint8,

309 + torch.float8_e4m3fnuz,

310 + ], (

311 + "unsupported dtype of KV cache tensor, got "

312 + "{kv_cache_torch_dtype}. Supported kv cache dtypes: fp8e4m3fn, "

313 + "fp8e5m2, uint8, bfloat16, float16, float32, fp8e4m3fnuz."

314 + )

315 +

316 + # heuristics instead of autotuning

317 + TILE_SIZE = max(head_size_k, head_size_v)

318 + TILE_SIZE = triton.next_power_of_2(TILE_SIZE)

319 + if current_platform.is_rocm() or current_platform.is_xpu():

320 + num_stages = 4

321 + num_warps = 8

322 + else: # cuda

323 + num_stages = 10

324 + num_warps = 16

325 +

326 + # TODO(ngl): maybe replace with static launch grid to avoid overhead if

327 + # using cudagraphs

328 + grid = lambda meta: (

329 + slot_mapping.shape[0],

330 + num_heads,

331 + )

332 +

333 + reshape_and_cache_kernel_flash_diffkv[grid](

334 + key_ptr=key,

335 + value_ptr=value,

336 + kv_cache_ptr=kv_cache,

337 + slot_mapping_ptr=slot_mapping,

338 + k_scale=k_scale,

339 + v_scale=v_scale,

340 + # strides

341 + key_stride=k_stride,

342 + value_stride=v_stride,

343 + block_stride=block_stride,

344 + page_stride=page_stride,

345 + num_heads=num_heads,

346 + head_size_k=head_size_k,

347 + head_size_v=head_size_v,

348 + block_size=block_size,

349 + # FP8 flags

350 + FP8_KV_CACHE=FP8_KV_CACHE,

351 + # autotune parameters

352 + TILE_SIZE=TILE_SIZE,

353 + num_warps=num_warps,

354 + num_stages=num_stages,

355 + )

vllm/model_executor/models/openpangu.py

+329 -3

@@ -29,13 +29,14 @@ import torch

29 from torch import nn

30 from transformers import PretrainedConfig

32 -from vllm.attention.backends.abstract import AttentionType

33 -from vllm.attention.layer import Attention

32 +from vllm.attention.layer import Attention, AttentionType

33 +from vllm.attention.layers.static_sink_attention import StaticSinkAttention

34 from vllm.compilation.decorators import support_torch_compile

35 from vllm.config import CacheConfig, ParallelConfig, VllmConfig

36 from vllm.distributed import (

37 get_ep_group,

38 get_pp_group,

39 + get_tensor_model_parallel_rank,

40 get_tensor_model_parallel_world_size,

41 get_tp_group,

42 tensor_model_parallel_all_gather,

@@ -77,8 +78,11 @@ from vllm.model_executor.models.utils import (

78 maybe_prefix,

79 sequence_parallel_chunk,

80 )

81 +from vllm.model_executor.utils import set_weight_attrs

82 +from vllm.platforms import current_platform

83 from vllm.sequence import IntermediateTensors

84 from vllm.transformers_utils.config import set_default_rope_theta

85 +from vllm.v1.attention.backends.flash_attn_diffkv import FlashAttentionDiffKVBackend

88 def check_ffn_act_fn(act_fn: str):

@@ -155,7 +159,15 @@ class OpenPanguMoE(nn.Module):

159 quant_config=None,

160 prefix=f"{prefix}.gate",

161 )

158 - self.gate.e_score_correction_bias = None

162 + if (

163 + hasattr(config, "router_enable_expert_bias")

164 + and config.router_enable_expert_bias

165 + ):

166 + self.gate.e_score_correction_bias = nn.Parameter(

167 + torch.empty(self.n_routed_experts, dtype=torch.float32)

168 + )

169 + else:

170 + self.gate.e_score_correction_bias = None

171

172 # Load balancing settings.

173 eplb_config = parallel_config.eplb_config

@@ -530,6 +542,264 @@ class OpenPanguEmbeddedAttention(nn.Module):

542 )

543

544

545 +class OpenPanguSinkAttention(nn.Module):

546 + def __init__(

547 + self,

548 + config: PretrainedConfig,

549 + hidden_size: int,

550 + num_heads: int,

551 + num_kv_heads: int,

552 + rope_parameters: dict[str, Any] | None = None,

553 + max_position_embeddings: int = 8192,

554 + quant_config: QuantizationConfig | None = None,

555 + bias: bool = False,

556 + bias_o_proj: bool = False,

557 + cache_config: CacheConfig | None = None,

558 + prefix: str = "",

559 + attn_type: str = AttentionType.DECODER,

560 + ) -> None:

561 + super().__init__()

562 + layer_idx = extract_layer_index(prefix)

563 + self.hidden_size = hidden_size

564 + self.tp_size = get_tensor_model_parallel_world_size()

565 + self.tp_rank = get_tensor_model_parallel_rank()

566 + self.total_num_heads = num_heads

567 + if self.total_num_heads % self.tp_size != 0:

568 + raise ValueError(

569 + f"total_num_heads {self.total_num_heads} "

570 + f"is not divisible by tp_size {self.tp_size}."

571 + )

572 + self.num_heads = self.total_num_heads // self.tp_size

573 + self.total_num_kv_heads = num_kv_heads

574 + if (

575 + self.total_num_kv_heads > self.tp_size

576 + and self.total_num_kv_heads % self.tp_size != 0

577 + ):

578 + # Number of KV heads is greater than TP size, so we partition

579 + # the KV heads across multiple tensor parallel ranks.

580 + raise ValueError(

581 + "Number of KV heads is greater than TP size, "

582 + f"but total_num_kv_heads {self.total_num_kv_heads} "

583 + f"is not divisible by tp_size {self.tp_size}."

584 + )

585 + elif self.total_num_kv_heads < self.tp_size:

586 + # TODO: Number of KV heads is less than TP size, so we replicate

587 + # the KV heads across multiple tensor parallel ranks.

588 + raise ValueError(

589 + f"Number of KV heads {self.total_num_kv_heads} is less than "

590 + f"TP size {self.tp_size}, KV heads replication is not support yet."

591 + )

592 + self.num_kv_heads = max(1, self.total_num_kv_heads // self.tp_size)

593 + self.qk_nope_dim = getattr(config, "qk_nope_dim", None)

594 + self.qk_rope_dim = getattr(config, "qk_rope_dim", None)

595 + self.v_channels = getattr(config, "v_channels", None)

596 + self.head_dim = self.qk_rope_dim + self.qk_nope_dim

597 + self.q_size = self.num_heads * self.head_dim

598 + self.k_size = self.num_kv_heads * self.head_dim

599 + self.v_size = self.num_kv_heads * self.v_channels

600 + self.scaling = self.head_dim**-0.5

601 + self.max_position_embeddings = max_position_embeddings

602 +

603 + self.param_sink_number = getattr(config, "param_sink_number", 0)

604 + self.param_sink_with_value = getattr(config, "param_sink_with_value", False)

605 + self.param_sink_scalar = getattr(config, "param_sink_scalar", None)

606 + self.param_sink_of_head_num = getattr(config, "param_sink_of_head_dim", False)

607 +

608 + self.qkv_proj = MergedColumnParallelLinear(

609 + input_size=hidden_size,

610 + output_sizes=[

611 + self.q_size * self.tp_size,

612 + self.k_size * self.tp_size,

613 + self.v_size * self.tp_size,

614 + ],

615 + bias=bias,

616 + quant_config=quant_config,

617 + prefix=f"{prefix}.qkv_proj",

618 + )

619 +

620 + self.o_proj = RowParallelLinear(

621 + input_size=self.total_num_heads * self.v_channels,

622 + output_size=hidden_size,

623 + bias=bias_o_proj,

624 + quant_config=quant_config,

625 + prefix=f"{prefix}.o_proj",

626 + )

627 +

628 + self.k_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

629 +

630 + self._init_rotary_emb(

631 + config, rope_parameters=rope_parameters, quant_config=quant_config

632 + )

633 +

634 + if hasattr(config, "interleaved_sliding_window"):

635 + interleaved_sliding_window = config.interleaved_sliding_window

636 + if isinstance(interleaved_sliding_window, int):

637 + sliding_window = interleaved_sliding_window

638 + elif isinstance(interleaved_sliding_window, list):

639 + sw_idx = layer_idx % len(interleaved_sliding_window)

640 + sliding_window = interleaved_sliding_window[sw_idx]

641 + else:

642 + raise ValueError(

643 + f"{type(interleaved_sliding_window)} "

644 + "for interleaved_sliding_window is not supported."

645 + )

646 + else:

647 + sliding_window = None

648 +

649 + FlashAttentionDiffKVBackend.set_head_size_v(self.v_channels)

650 + self.attn = StaticSinkAttention(

651 + self.num_heads,

652 + self.head_dim,

653 + self.scaling,

654 + sink_len=self.param_sink_number,

655 + num_kv_heads=self.num_kv_heads,

656 + cache_config=cache_config,

657 + quant_config=quant_config,

658 + per_layer_sliding_window=sliding_window,

659 + attn_type=attn_type,

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

661 + attn_backend=FlashAttentionDiffKVBackend,

662 + head_size_v=self.v_channels,

663 + )

664 +

665 + if self.param_sink_number > 0:

666 + self.param_sink_key = torch.nn.Parameter(

667 + torch.empty(

668 + (

669 + self.param_sink_number,

670 + self.num_kv_heads,

671 + self.head_dim,

672 + ),

673 + device=current_platform.current_device(),

674 + dtype=config.torch_dtype,

675 + )

676 + )

677 + set_weight_attrs(

678 + self.param_sink_key,

679 + {

680 + "output_dim": 1,

681 + "weight_loader": self.weight_loader,

682 + },

683 + )

684 +

685 + if self.param_sink_with_value:

686 + self.param_sink_value = torch.nn.Parameter(

687 + torch.empty(

688 + (

689 + self.param_sink_number,

690 + self.num_kv_heads,

691 + self.v_channels,

692 + ),

693 + device=current_platform.current_device(),

694 + dtype=config.torch_dtype,

695 + )

696 + )

697 + set_weight_attrs(

698 + self.param_sink_value,

699 + {

700 + "output_dim": 1,

701 + "weight_loader": self.weight_loader,

702 + },

703 + )

704 + else:

705 + self.param_sink_value = torch.zeros(

706 + (

707 + self.param_sink_number,

708 + self.num_kv_heads,

709 + self.v_channels,

710 + ),

711 + device=current_platform.current_device(),

712 + dtype=config.torch_dtype,

713 + )

714 + # To enable dummy run with out weight

715 + self.post_weight_load()

716 +

717 + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):

718 + output_dim = getattr(param, "output_dim", None)

719 +

720 + is_sharded_weight = getattr(param, "is_sharded_weight", False)

721 + use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)

722 + # bitsandbytes loads the weights of the specific portion

723 + # no need to narrow

724 + is_sharded_weight = is_sharded_weight or use_bitsandbytes_4bit

725 +

726 + # Special case for GGUF

727 + is_gguf_weight = getattr(param, "is_gguf_weight", False)

728 + is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False)

729 + if is_gguf_weight_type:

730 + param.weight_type = loaded_weight.item()

731 +

732 + # Materialize GGUF UninitializedParameter

733 + if is_gguf_weight and isinstance(param, nn.UninitializedParameter):

734 + final_shape = list(loaded_weight.shape)

735 + if output_dim is not None:

736 + assert final_shape[output_dim] % self.tp_size == 0

737 + final_shape[output_dim] = final_shape[output_dim] // self.tp_size

738 + param.materialize(final_shape, dtype=loaded_weight.dtype)

739 +

740 + param_data = param.data

741 + if output_dim is not None and not is_sharded_weight:

742 + shard_size = param_data.shape[output_dim]

743 + start_idx = self.tp_rank * shard_size

744 + loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)

745 +

746 + # Special case for loading scales off disk, which often do not

747 + # have a shape (such as in the case of AutoFP8).

748 + if len(loaded_weight.shape) == 0:

749 + loaded_weight = loaded_weight.reshape(1)

750 +

751 + assert param_data.shape == loaded_weight.shape

752 + param_data.copy_(loaded_weight)

753 +

754 + def forward(

755 + self,

756 + positions: torch.Tensor,

757 + hidden_states: torch.Tensor,

758 + ) -> torch.Tensor:

759 + qkv, _ = self.qkv_proj(hidden_states)

760 + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)

761 + k = self.k_layernorm(k.view(-1, self.num_kv_heads, self.head_dim))

762 + q, k = self.rotary_emb(positions, q, k)

763 +

764 + q = q.view(-1, self.q_size)

765 + k = k.view(-1, self.k_size)

766 +

767 + attn_output = self.attn(

768 + q,

769 + k,

770 + v,

771 + output_shape=torch.Size(

772 + [q.shape[0], q.shape[1] // self.head_dim * self.v_channels]

773 + ),

774 + )

775 + output, _ = self.o_proj(attn_output)

776 + return output

777 +

778 + def _init_rotary_emb(

779 + self,

780 + config: PretrainedConfig,

781 + rope_parameters: dict[str, Any] | None,

782 + quant_config: QuantizationConfig | None,

783 + ) -> None:

784 + is_neox_style = False

785 + rope_parameters = {"partial_rotary_factor": self.qk_rope_dim / self.head_dim}

786 +

787 + self.rotary_emb = get_rope(

788 + self.head_dim,

789 + max_position=self.max_position_embeddings,

790 + rope_parameters=rope_parameters,

791 + is_neox_style=is_neox_style,

792 + )

793 +

794 + def post_weight_load(self) -> None:

795 + if hasattr(self, "k_layernorm") and self.k_layernorm is not None:

796 + param_sink_key = self.k_layernorm(self.param_sink_key)

797 + else:

798 + param_sink_key = self.param_sink_key

799 +

800 + self.attn.update_sink_kv(param_sink_key, self.param_sink_value)

801 +

802 +

803 class OpenPanguDecoderLayer(nn.Module):

804 def __init__(

805 self,

@@ -557,6 +827,9 @@ class OpenPanguDecoderLayer(nn.Module):

827 and hasattr(config, "v_head_dim")

828 and hasattr(config, "kv_lora_rank")

829 )

830 + self.use_sink_attention = (

831 + hasattr(config, "param_sink_number") and config.param_sink_number > 0

832 + )

833 if self.use_mla:

834 self.self_attn = OpenPanguMLAAttention(

835 config=config,

@@ -574,6 +847,42 @@ class OpenPanguDecoderLayer(nn.Module):

847 quant_config=quant_config,

848 prefix=f"{prefix}.self_attn",

849 )

850 + elif self.use_sink_attention:

851 + attention_bias = getattr(config, "attention_bias", False) or getattr(

852 + config, "bias", False

853 + )

854 + bias_o_proj = attention_bias

855 + if hasattr(config, "qkv_bias"):

856 + attention_bias = config.qkv_bias

857 + if getattr(config, "is_causal", True):

858 + attn_type = AttentionType.DECODER

859 + else:

860 + raise ValueError(

861 + f"is_causal={config.is_causal} is not support "

862 + "for attention with sink"

863 + )

864 + rope_parameters = getattr(config, "rope_scaling", None)

865 + if rope_parameters is None:

866 + rope_parameters = {

867 + "rope_type": "default",

868 + "rope_theta": config.rope_theta,

869 + }

870 + self.self_attn = OpenPanguSinkAttention(

871 + config=config,

872 + hidden_size=self.hidden_size,

873 + num_heads=config.num_attention_heads,

874 + num_kv_heads=getattr(

875 + config, "num_key_value_heads", config.num_attention_heads

876 + ),

877 + rope_parameters=rope_parameters,

878 + max_position_embeddings=max_position_embeddings,

879 + quant_config=quant_config,

880 + bias=attention_bias,

881 + bias_o_proj=bias_o_proj,

882 + cache_config=cache_config,

883 + prefix=f"{prefix}.self_attn",

884 + attn_type=attn_type,

885 + )

886 else:

887 attention_bias = getattr(config, "attention_bias", False) or getattr(

888 config, "bias", False

@@ -903,6 +1212,10 @@ class OpenPanguModel(nn.Module):

1212 if name.endswith(".bias") and name not in params_dict:

1213 continue

1214 name = maybe_remap_kv_scale_name(name, params_dict)

1215 + if name.endswith("e_score_correction_bias"):

1216 + name = name.replace(

1217 + "e_score_correction_bias", "gate.e_score_correction_bias"

1218 + )

1219 if name is None:

1220 continue

1221 if is_pp_missing_parameter(name, self):

@@ -912,8 +1225,17 @@ class OpenPanguModel(nn.Module):

1225 weight_loader = getattr(param, "weight_loader", default_weight_loader)

1226 weight_loader(param, loaded_weight)

1227 loaded_params.add(name)

1228 +

1229 + self.post_weight_load()

1230 return loaded_params

1231

1232 + def post_weight_load(self) -> None:

1233 + for name, module in self.named_modules():

1234 + if module is self:

1235 + continue

1236 + if hasattr(module, "post_weight_load"):

1237 + module.post_weight_load()

1238 +

1239

1240 class OpenPanguModelBase(nn.Module, SupportsPP, SupportsLoRA):

1241 packed_modules_mapping = {

@@ -1047,3 +1369,7 @@ class PanguEmbeddedForCausalLM(OpenPanguEmbeddedModel):

1369

1370 class PanguUltraMoEForCausalLM(OpenPanguMoEModel):

1371 pass

1372 +

1373 +

1374 +class PanguProMoEV2ForCausalLM(OpenPanguMoEModel):

1375 + pass

vllm/model_executor/models/registry.py

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

164 "OrionForCausalLM": ("orion", "OrionForCausalLM"),

165 "OuroForCausalLM": ("ouro", "OuroForCausalLM"),

166 "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"),

167 + "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"),

168 "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"),

169 "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),

170 "PhiForCausalLM": ("phi", "PhiForCausalLM"),

vllm/v1/attention/backends/flash_attn_diffkv.py new

+269

@@ -0,0 +1,269 @@

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

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

3 +"""Attention layer with FlashAttention."""

4 +

5 +import torch

6 +

7 +from vllm.attention.backends.abstract import AttentionType

8 +from vllm.attention.ops.triton_reshape_and_cache_flash import (

9 + triton_reshape_and_cache_flash_diffkv,

10 +)

11 +from vllm.attention.utils.fa_utils import is_flash_attn_varlen_func_available

12 +

13 +if is_flash_attn_varlen_func_available():

14 + from vllm.attention.utils.fa_utils import flash_attn_varlen_func

15 +from vllm.logger import init_logger

16 +from vllm.v1.attention.backends.utils import get_kv_cache_layout

17 +

18 +from .flash_attn import (

19 + FlashAttentionBackend,

20 + FlashAttentionImpl,

21 + FlashAttentionMetadata,

22 + cascade_attention,

23 +)

24 +

25 +logger = init_logger(__name__)

26 +

27 +

28 +class FlashAttentionDiffKVBackend(FlashAttentionBackend):

29 + # Default to 128 for this backend

30 + head_size_v: int = 128

31 +

32 + @classmethod

33 + def set_head_size_v(cls, head_size_v: int) -> None:

34 + cls.head_size_v = head_size_v

35 +

36 + @staticmethod

37 + def get_name() -> str:

38 + return "FLASH_ATTN_DIFFKV"

39 +

40 + @staticmethod

41 + def get_impl_cls() -> type["FlashAttentionImpl"]:

42 + return FlashAttentionDiffKVImpl

43 +

44 + # Do not modify the interface of get_kv_cache_shape,

45 + # but consider head_size_v when returning result.

46 + @staticmethod

47 + def get_kv_cache_shape(

48 + num_blocks: int,

49 + block_size: int,

50 + num_kv_heads: int,

51 + head_size: int,

52 + cache_dtype_str: str = "auto",

53 + ) -> tuple[int, ...]:

54 + if block_size % 16 != 0:

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

56 + return (

57 + num_blocks,

58 + block_size,

59 + num_kv_heads,

60 + head_size + FlashAttentionDiffKVBackend.head_size_v,

61 + )

62 +

63 + @staticmethod

64 + def get_kv_cache_stride_order(

65 + include_num_layers_dimension: bool = False,

66 + ) -> tuple[int, ...]:

67 + # `stride_order` indicates the permutation that gets

68 + # us from `get_kv_cache_shape` to the actual memory layout we want.

69 + cache_layout = get_kv_cache_layout()

70 + if cache_layout == "NHD" and include_num_layers_dimension:

71 + # (num_blocks, num_layers, block_size,

72 + # num_kv_heads, head_size + head_size_v)

73 + return (1, 0, 2, 3, 4)

74 + elif cache_layout == "NHD":

75 + stride_order = (0, 1, 2, 3)

76 + elif cache_layout == "HND" and include_num_layers_dimension:

77 + # (num_blocks, num_kv_heads, num_layers,

78 + # block_size, head_size + head_size_v)

79 + return (1, 3, 0, 2, 4)

80 + elif cache_layout == "HND":

81 + stride_order = (0, 2, 1, 3)

82 + else:

83 + raise ValueError(f"Unknown cache layout format {cache_layout}.")

84 + return stride_order

85 +

86 +

87 +class FlashAttentionDiffKVImpl(FlashAttentionImpl):

88 + def forward(

89 + self,

90 + layer: torch.nn.Module,

91 + query: torch.Tensor,

92 + key: torch.Tensor,

93 + value: torch.Tensor,

94 + kv_cache: torch.Tensor,

95 + attn_metadata: FlashAttentionMetadata,

96 + output: torch.Tensor | None = None,

97 + output_scale: torch.Tensor | None = None,

98 + output_block_scale: torch.Tensor | None = None,

99 + ) -> torch.Tensor:

100 + """Forward pass with FlashAttention.

101 +

102 + Args:

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

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

105 + value: shape = [num_tokens, num_kv_heads, head_size_v]

106 + kv_cache: shape =

107 + [num_blocks, block_size, num_kv_heads, head_size + head_size_v]

108 + attn_metadata: Metadata for attention.

109 + Returns:

110 + shape = [num_tokens, num_heads * head_size_v]

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

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

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

114 + """

115 + assert output is not None, "Output tensor must be provided."

116 +

117 + if output_scale is not None or output_block_scale is not None:

118 + raise NotImplementedError(

119 + "fused output quantization is not yet supported for FlashAttentionImpl"

120 + )

121 +

122 + if attn_metadata is None:

123 + # Profiling run.

124 + return output.fill_(0)

125 +

126 + attn_type = self.attn_type

127 +

128 + # IMPORTANT!

129 + # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in

130 + # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead

131 + # in this method. For example, `view` and `slice` (or `[:n]`) operations

132 + # are surprisingly slow even in the case they do not invoke any GPU ops.

133 + # Minimize the PyTorch ops in this method as much as possible.

134 + # Whenever making a change in this method, please benchmark the

135 + # performance to make sure it does not introduce any overhead.

136 +

137 + num_actual_tokens = attn_metadata.num_actual_tokens

138 +

139 + # Handle encoder attention differently - no KV cache needed

140 + if attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):

141 + # For encoder attention,

142 + # we use direct Q, K, V tensors without caching

143 + return self._forward_encoder_attention(

144 + query[:num_actual_tokens],

145 + key[:num_actual_tokens],

146 + value[:num_actual_tokens],

147 + output[:num_actual_tokens],

148 + attn_metadata,

149 + layer,

150 + )

151 +

152 + # For decoder and cross-attention, use KV cache as before

153 + # Different head_size for K and V

154 + key_cache = kv_cache[..., : self.head_size]

155 + value_cache = kv_cache[..., self.head_size :]

156 +

157 + # key and value may be None in the case of cross attention. They are

158 + # calculated once based on the output from the encoder and then cached

159 + # in KV cache.

160 + if (

161 + self.kv_sharing_target_layer_name is None

162 + and key is not None

163 + and value is not None

164 + ):

165 + # Reshape the input keys and values and store them in the cache.

166 + # Skip this if sharing KV cache with an earlier attention layer.

167 + # NOTE(woosuk): Here, key and value are padded while slot_mapping is

168 + # not padded. However, we don't need to do key[:num_actual_tokens]

169 + # and value[:num_actual_tokens] because the reshape_and_cache_flash

170 + # op uses the slot_mapping's shape to determine the number of

171 + # actual tokens.

172 +

173 + # kv_cache update for different head_size K and V

174 + triton_reshape_and_cache_flash_diffkv(

175 + key,

176 + value,

177 + kv_cache,

178 + attn_metadata.slot_mapping,

179 + self.kv_cache_dtype,

180 + layer._k_scale,

181 + layer._v_scale,

182 + )

183 +

184 + if self.kv_cache_dtype.startswith("fp8"):

185 + # queries are quantized in the attention layer

186 + dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(

187 + self.kv_cache_dtype

188 + )

189 + key_cache = key_cache.view(dtype)

190 + value_cache = value_cache.view(dtype)

191 +

192 + if not attn_metadata.use_cascade:

193 + cu_seqlens_q = attn_metadata.query_start_loc

194 + seqused_k = attn_metadata.seq_lens

195 + max_seqlen_q = attn_metadata.max_query_len

196 + max_seqlen_k = attn_metadata.max_seq_len

197 + block_table = attn_metadata.block_table

198 + scheduler_metadata = attn_metadata.scheduler_metadata

199 +

200 + descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)

201 +

202 + if self.dcp_world_size > 1:

203 + self._forward_with_dcp(

204 + query[:num_actual_tokens],

205 + key[:num_actual_tokens],

206 + value[:num_actual_tokens],

207 + key_cache,

208 + value_cache,

209 + output[:num_actual_tokens],

210 + attn_metadata,

211 + q_descale=layer._q_scale.expand(descale_shape),

212 + k_descale=layer._k_scale.expand(descale_shape),

213 + v_descale=layer._v_scale.expand(descale_shape),

214 + )

215 + return output

216 + else:

217 + flash_attn_varlen_func(

218 + q=query[:num_actual_tokens],

219 + k=key_cache,

220 + v=value_cache,

221 + out=output[:num_actual_tokens],

222 + cu_seqlens_q=cu_seqlens_q,

223 + max_seqlen_q=max_seqlen_q,

224 + seqused_k=seqused_k,

225 + max_seqlen_k=max_seqlen_k,

226 + softmax_scale=self.scale,

227 + causal=attn_metadata.causal,

228 + alibi_slopes=self.alibi_slopes,

229 + window_size=self.sliding_window,

230 + block_table=block_table,

231 + softcap=self.logits_soft_cap,

232 + scheduler_metadata=scheduler_metadata,

233 + fa_version=self.vllm_flash_attn_version,

234 + q_descale=layer._q_scale.expand(descale_shape),

235 + k_descale=layer._k_scale.expand(descale_shape),

236 + v_descale=layer._v_scale.expand(descale_shape),

237 + num_splits=attn_metadata.max_num_splits,

238 + s_aux=self.sinks,

239 + )

240 + return output

241 +

242 + # Cascade attention (rare case).

243 + cascade_attention(

244 + output[:num_actual_tokens],

245 + query[:num_actual_tokens],

246 + key_cache,

247 + value_cache,

248 + cu_query_lens=attn_metadata.query_start_loc,

249 + max_query_len=attn_metadata.max_query_len,

250 + cu_prefix_query_lens=attn_metadata.cu_prefix_query_lens,

251 + prefix_kv_lens=attn_metadata.prefix_kv_lens,

252 + suffix_kv_lens=attn_metadata.suffix_kv_lens,

253 + max_kv_len=attn_metadata.max_seq_len,

254 + softmax_scale=self.scale,

255 + alibi_slopes=self.alibi_slopes,

256 + sliding_window=self.sliding_window,

257 + logits_soft_cap=self.logits_soft_cap,

258 + block_table=attn_metadata.block_table,

259 + common_prefix_len=attn_metadata.common_prefix_len,

260 + max_num_splits=attn_metadata.max_num_splits,

261 + fa_version=self.vllm_flash_attn_version,

262 + prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata,

263 + suffix_scheduler_metadata=attn_metadata.scheduler_metadata,

264 + q_descale=layer._q_scale,

265 + k_descale=layer._k_scale,

266 + v_descale=layer._v_scale,

267 + s_aux=self.sinks,

268 + )

269 + return output

vllm/v1/core/single_type_kv_cache_manager.py

+26

@@ -15,6 +15,7 @@ from vllm.v1.kv_cache_interface import (

15 KVCacheSpec,

16 MambaSpec,

17 MLAAttentionSpec,

18 + SinkFullAttentionSpec,

19 SlidingWindowSpec,

20 )

21 from vllm.v1.request import Request

@@ -882,6 +883,30 @@ class CrossAttentionManager(SingleTypeKVCacheManager):

883 raise NotImplementedError("CrossAttentionManager does not support caching")

884

885

886 +class SinkFullAttentionManager(FullAttentionManager):

887 + def __init__(

888 + self,

889 + kv_cache_spec: SinkFullAttentionSpec,

890 + block_pool: BlockPool,

891 + enable_caching: bool,

892 + kv_cache_group_id: int,

893 + dcp_world_size: int = 1,

894 + pcp_world_size: int = 1,

895 + ):

896 + super().__init__(

897 + kv_cache_spec,

898 + block_pool,

899 + enable_caching,

900 + kv_cache_group_id,

901 + dcp_world_size,

902 + pcp_world_size,

903 + )

904 + sink_len = kv_cache_spec.sink_len

905 + assert sink_len is not None and sink_len > 0 and sink_len % self.block_size == 0

906 + num_sink_block = sink_len // self.block_size

907 + self.sink_blocks = self.block_pool.free_block_queue.popleft_n(num_sink_block)

908 +

909 +

910 spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {

911 FullAttentionSpec: FullAttentionManager,

912 MLAAttentionSpec: FullAttentionManager,

@@ -889,6 +914,7 @@ spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {

914 ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,

915 MambaSpec: MambaManager,

916 CrossAttentionSpec: CrossAttentionManager,

917 + SinkFullAttentionSpec: SinkFullAttentionManager,

918 }

919

920

vllm/v1/kv_cache_interface.py

+66

@@ -89,12 +89,18 @@ class FullAttentionSpec(AttentionSpec):

89 In this case, we use FullAttentionSpec and record the sliding window size.

90 """

92 + head_size_v: int | None = None

93 +

94 sliding_window: int | None = None

95 """

96 Default to None for not using sliding window attention.

97 """

98 attention_chunk_size: int | None = None

100 + def __post_init__(self):

101 + if self.head_size_v is None:

102 + object.__setattr__(self, "head_size_v", self.head_size)

103 +

104 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

105 max_model_len = vllm_config.model_config.max_model_len

106 dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size

@@ -142,6 +148,7 @@ class FullAttentionSpec(AttentionSpec):

148 block_size=specs[0].block_size,

149 num_kv_heads=specs[0].num_kv_heads,

150 head_size=specs[0].head_size,

151 + head_size_v=specs[0].head_size_v,

152 dtype=specs[0].dtype,

153 sliding_window=cls.merge_window_sizes(sliding_window),

154 attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),

@@ -160,6 +167,15 @@ class FullAttentionSpec(AttentionSpec):

167 )

168 return merged_spec

169

170 + @property

171 + def page_size_bytes(self) -> int:

172 + return (

173 + self.block_size

174 + * self.num_kv_heads

175 + * (self.head_size + self.head_size_v)

176 + * get_dtype_size(self.dtype)

177 + )

178 +

179

180 @dataclass(frozen=True)

181 class MLAAttentionSpec(FullAttentionSpec):

@@ -287,6 +303,56 @@ class CrossAttentionSpec(AttentionSpec):

303 return cdiv(max_encoder_len, self.block_size) * self.page_size_bytes

304

305

306 +@dataclass(frozen=True)

307 +class SinkFullAttentionSpec(FullAttentionSpec):

308 + sink_len: int | None = None

309 +

310 + @classmethod

311 + def merge(cls, specs: list[Self]) -> Self:

312 + """

313 + Merge a list of FullAttentionSpec objects into a single

314 + FullAttentionSpec object.

315 + """

316 + assert all(isinstance(spec, FullAttentionSpec) for spec in specs), (

317 + "All attention layers in the same KV cache group must be FullAttentionSpec."

318 + )

319 +

320 + sliding_window = set(

321 + spec.sliding_window for spec in specs if spec.sliding_window is not None

322 + )

323 + attention_chunk_size = set(

324 + spec.attention_chunk_size

325 + for spec in specs

326 + if spec.attention_chunk_size is not None

327 + )

328 + assert not any(isinstance(spec, MLAAttentionSpec) for spec in specs), (

329 + "MLAAttentionSpec should be merged in MLAAttentionSpec.merge"

330 + )

331 + merged_spec = cls(

332 + block_size=specs[0].block_size,

333 + num_kv_heads=specs[0].num_kv_heads,

334 + head_size=specs[0].head_size,

335 + head_size_v=specs[0].head_size_v,

336 + sink_len=specs[0].sink_len,

337 + dtype=specs[0].dtype,

338 + sliding_window=cls.merge_window_sizes(sliding_window),

339 + attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),

340 + )

341 + for spec in specs:

342 + for f in fields(AttentionSpec):

343 + assert getattr(spec, f.name) == getattr(merged_spec, f.name), (

344 + "All attention layers in the same KV cache group must have "

345 + "the same attention spec."

346 + )

347 + assert (merged_spec.sliding_window is not None) + (

348 + merged_spec.attention_chunk_size is not None

349 + ) <= 1, (

350 + "Model with both sliding window layers and chunked local attention "

351 + "layers is not supported."

352 + )

353 + return merged_spec

354 +

355 +

356 @dataclass(frozen=True)

357 class UniformTypeKVCacheSpecs(KVCacheSpec):

358 """