@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 183dad7a85

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, 14:45 UTC

[Attention] Update to lastest FA3 code (#13111)

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

Lucas Wilkinson committed Apr 17, 2025 at 18:14 UTC 183dad7a85487dbb351c43c11d2180a6108d5448

5 files changed +241 -118

cmake/external_projects/vllm_flash_attn.cmake

+1 -1

@@ -38,7 +38,7 @@ else()

38 FetchContent_Declare(

39 vllm-flash-attn

40 GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git

41 - GIT_TAG dc9d410b3e2d6534a4c70724c2515f4def670a22

41 + GIT_TAG 0a721daebe4fa7149f06ecf3d3eabeb6dcd0f1fa

42 GIT_PROGRESS TRUE

43 # Don't share the vllm-flash-attn build between build types

44 BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

vllm/attention/backends/mla/common.py

+92 -90

@@ -1043,8 +1043,8 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1043 self.q_proj = q_proj

1044 self.kv_b_proj = kv_b_proj

1045 self.o_proj = o_proj

1046 - self.triton_fa_func = triton_attention

1046

1047 + self.triton_fa_func = triton_attention

1048 # Handle the differences between the flash_attn_varlen from flash_attn

1049 # and the one from vllm_flash_attn. The former is used on RoCM and the

1050 # latter has an additional parameter to control FA2 vs FA3

@@ -1055,6 +1055,70 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1055 functools.partial(flash_attn_varlen_func,

1056 fa_version=self.vllm_flash_attn_version)

1057

1058 + # For MLA the v head dim is smaller than qk head dim so we pad out

1059 + # v with 0s to match the qk head dim for attention backends that do

1060 + # not support different headdims

1061 + # We don't need to pad V if we are on a hopper system with FA3

1062 + self._pad_v = self.vllm_flash_attn_version is None or not (

1063 + self.vllm_flash_attn_version == 3

1064 + and current_platform.get_device_capability()[0] == 9)

1065 +

1066 + def _flash_attn_varlen_diff_headdims(self, q, k, v, softmax_scale,

1067 + return_softmax_lse, **kwargs):

1068 + maybe_padded_v = v

1069 + if self._pad_v:

1070 + maybe_padded_v = torch.nn.functional.pad(

1071 + v, [0, q.shape[-1] - v.shape[-1]], value=0)

1072 +

1073 + if is_hip and envs.VLLM_USE_TRITON_FLASH_ATTN \

1074 + and not return_softmax_lse:

1075 + attn_out = self.triton_fa_func(

1076 + q,

1077 + k,

1078 + maybe_padded_v,

1079 + **kwargs,

1080 + )

1081 + if is_vllm_fa:

1082 + attn_out = self.flash_attn_varlen_func(

1083 + q=q,

1084 + k=k,

1085 + v=maybe_padded_v,

1086 + return_softmax_lse=return_softmax_lse,

1087 + softmax_scale=softmax_scale,

1088 + **kwargs,

1089 + )

1090 + else:

1091 + # Use return_attn_probs instead of return_softmax_lse for RoCM

1092 + attn_out = self.flash_attn_varlen_func(

1093 + q=q,

1094 + k=k,

1095 + v=maybe_padded_v,

1096 + return_attn_probs=return_softmax_lse,

1097 + softmax_scale=softmax_scale,

1098 + **kwargs,

1099 + )

1100 +

1101 + # Unpack the output if there is multiple results,

1102 + # triton always returns (output, softmax_lse),

1103 + # vllm_flash_attn returns (output, softmax_lse) when

1104 + # `return_softmax_lse = True`

1105 + # flash_attn (RoCM) returns (output, softmax_lse, ...) when

1106 + # `return_attn_probs = True`

1107 + rest = None

1108 + if isinstance(attn_out, tuple):

1109 + attn_out, *rest = attn_out

1110 +

1111 + # unpad if necessary

1112 + if self._pad_v:

1113 + attn_out = attn_out[..., :v.shape[-1]]

1114 +

1115 + # Remain consistent with old `flash_attn_varlen_func` where there

1116 + # is only one output tensor if `return_softmax_lse` is False.

1117 + if return_softmax_lse:

1118 + assert rest is not None

1119 + return attn_out, rest[0]

1120 + return attn_out

1121 +

1122 def _v_up_proj_and_o_proj(self, x):

1123 # Convert from (B, N, L) to (N, B, L)

1124 x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)

@@ -1176,40 +1240,19 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1240 k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))),

1241 dim=-1)

1242

1179 - # For MLA the v head dim is smaller than qk head dim so we pad

1180 - # out v with 0s to match the qk head dim

1181 - v_padded = torch.nn.functional.pad(v,

1182 - [0, q.shape[-1] - v.shape[-1]],

1183 - value=0)

1184 -

1185 - if is_vllm_fa:

1186 - attn_output, attn_softmax_lse = self.flash_attn_varlen_func(

1187 - q=q,

1188 - k=k,

1189 - v=v_padded,

1190 - cu_seqlens_q=prefill_metadata.query_start_loc,

1191 - cu_seqlens_k=prefill_metadata.context_chunk_cu_seq_lens[i],

1192 - max_seqlen_q=prefill_metadata.max_query_len,

1193 - max_seqlen_k=prefill_metadata.

1194 - context_chunk_max_seq_lens[i],

1195 - softmax_scale=self.scale,

1196 - causal=False, # Context is unmasked

1197 - return_softmax_lse=True,

1198 - )

1199 - else:

1200 - attn_output, attn_softmax_lse, _ = self.flash_attn_varlen_func(

1201 - q=q,

1202 - k=k,

1203 - v=v_padded,

1204 - cu_seqlens_q=prefill_metadata.query_start_loc,

1205 - cu_seqlens_k=prefill_metadata.context_chunk_cu_seq_lens[i],

1206 - max_seqlen_q=prefill_metadata.max_query_len,

1207 - max_seqlen_k=prefill_metadata.

1208 - context_chunk_max_seq_lens[i],

1209 - softmax_scale=self.scale,

1210 - causal=False, # Context is unmasked

1211 - return_attn_probs=True,

1212 - )

1243 + attn_output, attn_softmax_lse = \

1244 + self._flash_attn_varlen_diff_headdims(

1245 + q=q,

1246 + k=k,

1247 + v=v,

1248 + cu_seqlens_q=prefill_metadata.query_start_loc,

1249 + cu_seqlens_k=prefill_metadata.context_chunk_cu_seq_lens[i],

1250 + max_seqlen_q=prefill_metadata.max_query_len,

1251 + max_seqlen_k=prefill_metadata.context_chunk_max_seq_lens[i],

1252 + softmax_scale=self.scale,

1253 + causal=False, # Context is unmasked

1254 + return_softmax_lse=True,

1255 + )

1256

1257 if output is None:

1258 output = attn_output

@@ -1252,58 +1295,22 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1295

1296 k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)

1297

1255 - # For MLA the v head dim is smaller than qk head dim so we pad out

1256 - # v with 0s to match the qk head dim

1257 - v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],

1258 - value=0)

1259 -

1260 - if is_hip and envs.VLLM_USE_TRITON_FLASH_ATTN and not has_context:

1261 - output = self.triton_fa_func(

1262 - q,

1263 - k,

1264 - v_padded,

1265 - None,

1266 - prefill_metadata.query_start_loc,

1267 - prefill_metadata.query_start_loc,

1268 - prefill_metadata.max_prefill_seq_len,

1269 - prefill_metadata.max_prefill_seq_len,

1270 - True, # causal

1271 - self.scale,

1272 - None, # attn_mask is None unless applying ALiBi mask

1273 - )

1274 - ## triton flash attention always return 2 objects

1275 - if not has_context:

1276 - output = output[0]

1277 - elif is_vllm_fa:

1278 - output = self.flash_attn_varlen_func(

1279 - q=q,

1280 - k=k,

1281 - v=v_padded,

1282 - cu_seqlens_q=prefill_metadata.query_start_loc,

1283 - cu_seqlens_k=prefill_metadata.query_start_loc,

1284 - max_seqlen_q=prefill_metadata.max_prefill_seq_len,

1285 - max_seqlen_k=prefill_metadata.max_prefill_seq_len,

1286 - softmax_scale=self.scale,

1287 - causal=True,

1288 - return_softmax_lse=has_context,

1289 - )

1290 - else:

1291 - output = self.flash_attn_varlen_func(

1292 - q=q,

1293 - k=k,

1294 - v=v_padded,

1295 - cu_seqlens_q=prefill_metadata.query_start_loc,

1296 - cu_seqlens_k=prefill_metadata.query_start_loc,

1297 - max_seqlen_q=prefill_metadata.max_prefill_seq_len,

1298 - max_seqlen_k=prefill_metadata.max_prefill_seq_len,

1299 - softmax_scale=self.scale,

1300 - causal=True,

1301 - return_attn_probs=has_context,

1302 - )

1298 + output = self._flash_attn_varlen_diff_headdims(

1299 + q=q,

1300 + k=k,

1301 + v=v,

1302 + cu_seqlens_q=prefill_metadata.query_start_loc,

1303 + cu_seqlens_k=prefill_metadata.query_start_loc,

1304 + max_seqlen_q=prefill_metadata.max_prefill_seq_len,

1305 + max_seqlen_k=prefill_metadata.max_prefill_seq_len,

1306 + softmax_scale=self.scale,

1307 + causal=True,

1308 + return_softmax_lse=has_context,

1309 + )

1310

1311 if has_context:

1312 # ROCm flash_attn_varlen_func will return 3 objects instead of 2

1306 - suffix_output, suffix_lse, *rest = output

1313 + suffix_output, suffix_lse = output

1314 context_output, context_lse = self._compute_prefill_context( \

1315 q, kv_c_and_k_pe_cache, attn_metadata)

1316

@@ -1316,12 +1323,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1323 suffix_lse=suffix_lse,

1324 )

1325

1319 - # slice by `:v.shape[-1]` in order to remove v headdim padding

1320 - output = output\

1321 - .view(-1, self.num_heads, q.shape[-1])[..., :v.shape[-1]]\

1322 - .reshape(-1, self.num_heads * v.shape[-1])

1323 -

1324 - return self.o_proj(output)[0]

1326 + return self.o_proj(output.flatten(start_dim=-2))[0]

1327

1328 @abstractmethod

1329 def _forward_decode(

vllm/attention/backends/utils.py

+25 -1

@@ -2,8 +2,10 @@

2 """Attention backend utils"""

3 from collections import defaultdict

4 from contextlib import contextmanager

5 +from dataclasses import dataclass

6 from itertools import accumulate

6 -from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Type, TypeVar, Union

7 +from typing import (TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type,

8 + TypeVar, Union)

10 import numpy as np

11 import torch

@@ -11,6 +13,7 @@ import torch

13 from vllm.attention import (AttentionMetadata, AttentionMetadataBuilder,

14 AttentionState)

15 from vllm.attention.backends.abstract import AttentionType

16 +from vllm.config import ModelConfig

17 from vllm.logger import init_logger

18 from vllm.multimodal import MultiModalPlaceholderMap

19 from vllm.utils import async_tensor_h2d, make_tensor_with_pad

@@ -583,3 +586,24 @@ def get_num_prefill_decode_query_kv_tokens(

586

587 return (num_prefill_query_tokens, num_prefill_kv_tokens,

588 num_decode_query_tokens)

589 +

590 +

591 +@dataclass

592 +class MLADims:

593 + q_lora_rank: Optional[int]

594 + kv_lora_rank: int

595 + qk_nope_head_dim: int

596 + qk_rope_head_dim: int

597 + v_head_dim: int

598 +

599 +

600 +def get_mla_dims(model_config: ModelConfig) -> MLADims:

601 + hf_text_config = model_config.hf_text_config

602 +

603 + return MLADims(

604 + q_lora_rank=getattr(hf_text_config, "q_lora_rank", None),

605 + kv_lora_rank=hf_text_config.kv_lora_rank,

606 + qk_nope_head_dim=hf_text_config.qk_nope_head_dim,

607 + qk_rope_head_dim=hf_text_config.qk_rope_head_dim,

608 + v_head_dim=hf_text_config.v_head_dim,

609 + )

vllm/v1/attention/backends/flash_attn.py

+58 -1

@@ -23,7 +23,8 @@ if TYPE_CHECKING:

23 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

25 if current_platform.is_cuda():

26 - from vllm.vllm_flash_attn import flash_attn_varlen_func

26 + from vllm.vllm_flash_attn import (flash_attn_varlen_func,

27 + get_scheduler_metadata)

29 logger = init_logger(__name__)

@@ -93,6 +94,10 @@ class FlashAttentionMetadata:

94 prefix_kv_lens: Optional[torch.Tensor]

95 suffix_kv_lens: Optional[torch.Tensor]

97 + # Optional aot scheduling

98 + scheduler_metadata: Optional[torch.Tensor] = None

99 + prefix_scheduler_metadata: Optional[torch.Tensor] = None

100 +

101 # For logging.

102 num_input_tokens: int = 0 # Number of tokens including padding.

103

@@ -277,7 +282,14 @@ def make_local_attention_virtual_batches(

282 class FlashAttentionMetadataBuilder:

283

284 def __init__(self, runner: "GPUModelRunner"):

285 + model_config = runner.model_config

286 +

287 self.runner = runner

288 + self.aot_schedule = (get_flash_attn_version() == 3)

289 + self.num_heads = model_config.get_num_attention_heads(

290 + runner.parallel_config)

291 + self.headdim = model_config.get_head_size()

292 + self.page_size = self.runner.block_size

293

294 def reorder_batch(self, input_batch: "InputBatch",

295 scheduler_output: "SchedulerOutput") -> bool:

@@ -319,6 +331,24 @@ class FlashAttentionMetadataBuilder:

331 )

332

333 use_cascade = common_prefix_len > 0

334 +

335 + def schedule(cu_query_lens, max_query_len, seqlens, max_seq_len,

336 + causal):

337 + if self.aot_schedule:

338 + return get_scheduler_metadata(

339 + batch_size=num_reqs,

340 + max_seqlen_q=max_query_len,

341 + max_seqlen_k=max_seq_len,

342 + cache_seqlens=seqlens,

343 + num_heads_q=self.num_heads,

344 + num_heads_kv=self.num_heads,

345 + headdim=self.headdim,

346 + page_size=self.page_size,

347 + cu_seqlens_q=cu_query_lens,

348 + causal=causal,

349 + )

350 + return None

351 +

352 if use_cascade:

353 cu_prefix_query_lens = torch.tensor([0, num_actual_tokens],

354 dtype=torch.int32,

@@ -330,10 +360,28 @@ class FlashAttentionMetadataBuilder:

360 common_prefix_len)

361 suffix_kv_lens = torch.from_numpy(suffix_kv_lens).to(

362 self.runner.device)

363 + prefix_scheduler_metadata = schedule(

364 + cu_query_lens=cu_prefix_query_lens,

365 + max_query_len=num_actual_tokens,

366 + seqlens=prefix_kv_lens,

367 + max_seq_len=common_prefix_len,

368 + causal=False)

369 + scheduler_metadata = schedule(cu_query_lens=query_start_loc,

370 + max_query_len=max_query_len,

371 + seqlens=suffix_kv_lens,

372 + max_seq_len=max_seq_len -

373 + common_prefix_len,

374 + causal=True)

375 else:

376 cu_prefix_query_lens = None

377 prefix_kv_lens = None

378 suffix_kv_lens = None

379 + prefix_scheduler_metadata = None

380 + scheduler_metadata = schedule(cu_query_lens=query_start_loc,

381 + max_query_len=max_query_len,

382 + seqlens=seq_lens,

383 + max_seq_len=max_seq_len,

384 + causal=True)

385

386 attn_metadata = FlashAttentionMetadata(

387 num_actual_tokens=num_actual_tokens,

@@ -345,10 +393,12 @@ class FlashAttentionMetadataBuilder:

393 slot_mapping=slot_mapping,

394 use_cascade=use_cascade,

395 common_prefix_len=common_prefix_len,

396 + scheduler_metadata=scheduler_metadata,

397 cu_prefix_query_lens=cu_prefix_query_lens,

398 prefix_kv_lens=prefix_kv_lens,

399 suffix_kv_lens=suffix_kv_lens,

400 local_attn_metadata=local_attn_metadata,

401 + prefix_scheduler_metadata=prefix_scheduler_metadata,

402 )

403 return attn_metadata

404

@@ -515,6 +565,7 @@ class FlashAttentionImpl(AttentionImpl):

565 window_size=self.sliding_window,

566 block_table=block_table,

567 softcap=self.logits_soft_cap,

568 + scheduler_metadata=attn_metadata.scheduler_metadata,

569 fa_version=self.vllm_flash_attn_version,

570 q_descale=layer._q_scale.expand(descale_shape),

571 k_descale=layer._k_scale.expand(descale_shape),

@@ -543,6 +594,8 @@ class FlashAttentionImpl(AttentionImpl):

594 block_table=attn_metadata.block_table,

595 common_prefix_len=attn_metadata.common_prefix_len,

596 fa_version=self.vllm_flash_attn_version,

597 + prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata,

598 + suffix_scheduler_metadata=attn_metadata.scheduler_metadata,

599 q_descale=layer._q_scale,

600 k_descale=layer._k_scale,

601 v_descale=layer._v_scale,

@@ -636,6 +689,8 @@ def cascade_attention(

689 block_table: torch.Tensor,

690 common_prefix_len: int,

691 fa_version: int,

692 + prefix_scheduler_metadata: Optional[torch.Tensor] = None,

693 + suffix_scheduler_metadata: Optional[torch.Tensor] = None,

694 q_descale: Optional[torch.Tensor] = None,

695 k_descale: Optional[torch.Tensor] = None,

696 v_descale: Optional[torch.Tensor] = None,

@@ -667,6 +722,7 @@ def cascade_attention(

722 block_table=block_table[:1],

723 softcap=logits_soft_cap,

724 return_softmax_lse=True,

725 + scheduler_metadata=prefix_scheduler_metadata,

726 fa_version=fa_version,

727 q_descale=q_descale.expand(descale_shape)

728 if q_descale is not None else None,

@@ -693,6 +749,7 @@ def cascade_attention(

749 block_table=block_table[:, num_common_kv_blocks:],

750 softcap=logits_soft_cap,

751 return_softmax_lse=True,

752 + scheduler_metadata=suffix_scheduler_metadata,

753 fa_version=fa_version,

754 q_descale=q_descale.expand(descale_shape)

755 if q_descale is not None else None,

vllm/v1/attention/backends/mla/common.py

+65 -25

@@ -195,6 +195,7 @@ from vllm import _custom_ops as ops

195 from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,

196 AttentionMetadata,

197 MLAAttentionImpl)

198 +from vllm.attention.backends.utils import get_mla_dims

199 from vllm.attention.ops.merge_attn_states import merge_attn_states

200 from vllm.logger import init_logger

201 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

@@ -207,9 +208,11 @@ from vllm.vllm_flash_attn.fa_utils import get_flash_attn_version

208

209 try:

210 from vllm.vllm_flash_attn import flash_attn_varlen_func

211 + is_vllm_fa = True

212 except ImportError:

213 # For rocm use upstream flash attention

214 from flash_attn import flash_attn_varlen_func

215 + is_vllm_fa = False

216

217 if TYPE_CHECKING:

218 from vllm.v1.core.sched.output import SchedulerOutput

@@ -350,6 +353,14 @@ class MLACommonMetadataBuilder(Generic[M]):

353 model_config = runner.model_config

354 cache_config = runner.cache_config

355 self.chunked_prefill_enabled = scheduler_config.chunked_prefill_enabled

356 + self.num_heads = model_config.get_num_attention_heads(

357 + runner.parallel_config)

358 + self.mla_dims = get_mla_dims(model_config)

359 + self.aot_schedule = is_vllm_fa and (get_flash_attn_version() == 3)

360 +

361 + # Dont try to access the runner on AMD

362 + if self.aot_schedule:

363 + self.page_size = self.runner.block_size

364

365 if self.chunked_prefill_enabled:

366 self.chunked_prefill_workspace_size = min(

@@ -375,7 +386,6 @@ class MLACommonMetadataBuilder(Generic[M]):

386 dtype=model_config.dtype,

387 device=runner.device,

388 )

378 - self.page_size = self.runner.block_size

389

390 def reorder_batch(self, input_batch: "InputBatch",

391 scheduler_output: "SchedulerOutput") -> bool:

@@ -464,7 +474,6 @@ class MLACommonMetadataBuilder(Generic[M]):

474

475 seq_lens_cpu = self.runner.seq_lens_cpu[:num_reqs]

476 seq_lens = seq_lens_cpu.to(device, non_blocking=True)

467 - max_query_len = seq_lens_cpu.max().item()

477

478 prefill_metadata = None

479 if self._num_prefills > 0:

@@ -475,6 +484,8 @@ class MLACommonMetadataBuilder(Generic[M]):

484 num_computed_tokens_cpu_tensor[reqs_start:num_reqs]

485 max_context_len_cpu = context_lens_cpu.max().item()

486 num_prefills_with_context_cpu = (context_lens_cpu > 0).sum().item()

487 + prefill_query_start_loc = query_start_loc[

488 + reqs_start:] - query_start_loc[reqs_start]

489

490 chunked_context_metadata = None

491 if self.chunked_prefill_enabled and self._num_prefills > 0 \

@@ -537,8 +548,7 @@ class MLACommonMetadataBuilder(Generic[M]):

548 prefill_metadata = MLACommonPrefillMetadata(

549 input_positions=input_positions[tokens_start:],

550 block_table=block_table[reqs_start:, ...],

540 - query_start_loc=query_start_loc[reqs_start:] -

541 - query_start_loc[reqs_start],

551 + query_start_loc=prefill_query_start_loc,

552 max_query_len=max_query_len,

553 chunked_context=chunked_context_metadata,

554 )

@@ -628,11 +638,56 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

638 # and the one from vllm_flash_attn. The former is used on RoCM and the

639 # latter has an additional parameter to control FA2 vs FA3

640 self.flash_attn_varlen_func = flash_attn_varlen_func

641 + self.vllm_flash_attn_version = get_flash_attn_version()

642 if self.vllm_flash_attn_version is not None:

643 self.flash_attn_varlen_func = \

644 functools.partial(flash_attn_varlen_func,

645 fa_version=self.vllm_flash_attn_version)

646

647 + # For MLA the v head dim is smaller than qk head dim so we pad out

648 + # v with 0s to match the qk head dim for attention backends that do

649 + # not support different headdims

650 + # We don't need to pad V if we are on a hopper system with FA3

651 + self._pad_v = self.vllm_flash_attn_version is None or not (

652 + self.vllm_flash_attn_version == 3

653 + and current_platform.get_device_capability()[0] == 9)

654 +

655 + def _flash_attn_varlen_diff_headdims(self,

656 + q,

657 + k,

658 + v,

659 + return_softmax_lse=False,

660 + softmax_scale=None,

661 + **kwargs):

662 + maybe_padded_v = v

663 + if self._pad_v:

664 + maybe_padded_v = torch.nn.functional.pad(

665 + v, [0, q.shape[-1] - v.shape[-1]], value=0)

666 +

667 + attn_out = self.flash_attn_varlen_func(

668 + q=q,

669 + k=k,

670 + v=maybe_padded_v,

671 + return_softmax_lse=return_softmax_lse,

672 + softmax_scale=softmax_scale,

673 + **kwargs,

674 + )

675 +

676 + # Unpack the output if there is multiple results

677 + lse = None

678 + if isinstance(attn_out, tuple):

679 + attn_out, lse = attn_out[0], attn_out[1]

680 +

681 + # unpad if necessary

682 + if self._pad_v:

683 + attn_out = attn_out[..., :v.shape[-1]]

684 +

685 + # Remain consistent with old `flash_attn_varlen_func` where there

686 + # is only one output tensor if `return_softmax_lse` is False.

687 + if return_softmax_lse:

688 + return attn_out, lse

689 + return attn_out

690 +

691 def _v_up_proj_and_o_proj(self, x):

692 # Convert from (B, N, L) to (N, B, L)

693 x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)

@@ -745,16 +800,11 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

800 k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))),

801 dim=-1)

802

748 - # For MLA the v head dim is smaller than qk head dim so we pad

749 - # out v with 0s to match the qk head dim

750 - v_padded = torch.nn.functional.pad(v,

751 - [0, q.shape[-1] - v.shape[-1]],

752 - value=0)

753 -

754 - attn_output, attn_softmax_lse = self.flash_attn_varlen_func(

803 + attn_output, attn_softmax_lse = \

804 + self._flash_attn_varlen_diff_headdims(

805 q=q,

806 k=k,

757 - v=v_padded,

807 + v=v,

808 cu_seqlens_q=prefill_metadata.query_start_loc,

809 cu_seqlens_k=prefill_metadata.chunked_context.cu_seq_lens[i],

810 max_seqlen_q=prefill_metadata.max_query_len,

@@ -801,15 +851,10 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

851

852 k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)

853

804 - # For MLA the v head dim is smaller than qk head dim so we pad out

805 - # v with 0s to match the qk head dim

806 - v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],

807 - value=0)

808 -

809 - output = self.flash_attn_varlen_func(

854 + output = self._flash_attn_varlen_diff_headdims(

855 q=q,

856 k=k,

812 - v=v_padded,

857 + v=v,

858 cu_seqlens_q=attn_metadata.prefill.query_start_loc,

859 cu_seqlens_k=attn_metadata.prefill.query_start_loc,

860 max_seqlen_q=attn_metadata.prefill.max_query_len,

@@ -833,12 +878,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

878 suffix_lse=suffix_lse,

879 )

880

836 - # slice by `:v.shape[-1]` in order to remove v headdim padding

837 - output = output\

838 - .view(-1, self.num_heads, q.shape[-1])[..., :v.shape[-1]]\

839 - .reshape(-1, self.num_heads * v.shape[-1])

840 -

841 - return self.o_proj(output)[0]

881 + return self.o_proj(output.flatten(start_dim=-2))[0]

882

883 @abstractmethod

884 def _forward_decode(