|
7 | 7 | from flag_gems.runtime import torch_device_fn |
8 | 8 |
|
9 | 9 | from .. import runtime |
| 10 | +from .flash_api import mha_fwd, mha_varlan_fwd |
| 11 | +from .flash_kernel import keep |
10 | 12 |
|
11 | 13 | logger = logging.getLogger(__name__) |
12 | 14 |
|
@@ -122,18 +124,9 @@ def _attn_fwd_inner( |
122 | 124 | return acc, l_i, m_i |
123 | 125 |
|
124 | 126 |
|
125 | | -def early_config_prune(configs, nargs, **kwargs): |
126 | | - return list(filter(lambda cfg: cfg.kwargs["BLOCK_N"] <= nargs["HEAD_DIM"], configs)) |
127 | | - |
128 | | - |
129 | 127 | @triton.autotune( |
130 | | - configs=runtime.get_tuned_config("attention"), |
| 128 | + configs=list(filter(keep, runtime.get_tuned_config("attention"))), |
131 | 129 | key=["KV_CTX", "HEAD_DIM"], |
132 | | - prune_configs_by={ |
133 | | - "early_config_prune": early_config_prune, |
134 | | - "perf_model": None, |
135 | | - "top_k": 1.0, |
136 | | - }, |
137 | 130 | ) |
138 | 131 | @triton.jit |
139 | 132 | def _attn_fwd( |
@@ -400,3 +393,223 @@ def scaled_dot_product_attention( |
400 | 393 | HAS_ATTN_MASK=HAS_ATTN_MASK, # |
401 | 394 | ) |
402 | 395 | return o |
| 396 | + |
| 397 | + |
| 398 | +def flash_attention_forward( |
| 399 | + query, |
| 400 | + key, |
| 401 | + value, |
| 402 | + cumulative_sequence_length_q, |
| 403 | + cumulative_sequence_length_k, |
| 404 | + max_q, |
| 405 | + max_k, |
| 406 | + dropout_p, |
| 407 | + is_causal, |
| 408 | + return_debug_mask, |
| 409 | + *, |
| 410 | + scale=None, |
| 411 | + softcap=0.0, |
| 412 | + window_size_left=None, |
| 413 | + window_size_right=None, |
| 414 | + seqused_k=None, |
| 415 | + alibi_slopes=None, |
| 416 | + disable_splitkv=False, |
| 417 | +): |
| 418 | + logger.debug("GEMS FLASH_ATTENTION_FORWARD") |
| 419 | + assert ( |
| 420 | + cumulative_sequence_length_q is None and cumulative_sequence_length_k is None |
| 421 | + ), "varlen is not supported yet." |
| 422 | + |
| 423 | + HEAD_DIM_Q, HEAD_DIM_K = query.shape[-1], key.shape[-1] |
| 424 | + HEAD_DIM_V = value.shape[-1] |
| 425 | + assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V |
| 426 | + assert HEAD_DIM_K in {16, 32, 64, 128, 256} |
| 427 | + |
| 428 | + softmax_scale = scale or 1.0 / (HEAD_DIM_K**0.5) |
| 429 | + if window_size_left is not None: |
| 430 | + non_null_window_left = window_size_left |
| 431 | + else: |
| 432 | + non_null_window_left = -1 |
| 433 | + if window_size_right is not None: |
| 434 | + non_null_window_right = window_size_right |
| 435 | + else: |
| 436 | + non_null_window_right = -1 |
| 437 | + |
| 438 | + out = torch.empty_like(query) |
| 439 | + if cumulative_sequence_length_q is not None: |
| 440 | + out, q, k, v, lse, philox_seed, philox_offset, p = mha_varlan_fwd( |
| 441 | + query, |
| 442 | + key, |
| 443 | + value, |
| 444 | + out, |
| 445 | + cumulative_sequence_length_q, |
| 446 | + cumulative_sequence_length_k, |
| 447 | + seqused_k, |
| 448 | + None, |
| 449 | + None, # block_table |
| 450 | + alibi_slopes, |
| 451 | + max_q, |
| 452 | + max_k, |
| 453 | + dropout_p, |
| 454 | + scale, |
| 455 | + False, |
| 456 | + is_causal, |
| 457 | + non_null_window_left, |
| 458 | + non_null_window_right, |
| 459 | + softcap, |
| 460 | + return_debug_mask and dropout_p > 0, |
| 461 | + None, |
| 462 | + ) |
| 463 | + else: |
| 464 | + out, q, k, v, lse, philox_seed, philox_offset, p = mha_fwd( |
| 465 | + query, |
| 466 | + key, |
| 467 | + value, |
| 468 | + out, |
| 469 | + alibi_slopes, |
| 470 | + dropout_p, |
| 471 | + softmax_scale, |
| 472 | + is_causal, |
| 473 | + non_null_window_left, |
| 474 | + non_null_window_right, |
| 475 | + softcap, |
| 476 | + return_debug_mask, |
| 477 | + disable_splitkv=disable_splitkv, |
| 478 | + ) |
| 479 | + |
| 480 | + return (out, lse, philox_seed, philox_offset, p) |
| 481 | + |
| 482 | + |
| 483 | +# Adapted from https://github.com/vllm-project/flash-attention/blob/main/vllm_flash_attn/flash_attn_interface.py |
| 484 | +def maybe_contiguous(x): |
| 485 | + return x.contiguous() if x is not None and x.stride(-1) != 1 else x |
| 486 | + |
| 487 | + |
| 488 | +def flash_attn_varlen_func( |
| 489 | + q, |
| 490 | + k, |
| 491 | + v, |
| 492 | + max_seqlen_q, |
| 493 | + cu_seqlens_q, |
| 494 | + max_seqlen_k, |
| 495 | + cu_seqlens_k=None, # only used for non-paged prefill |
| 496 | + seqused_k=None, |
| 497 | + q_v=None, |
| 498 | + dropout_p=0.0, |
| 499 | + softmax_scale=None, |
| 500 | + causal=False, |
| 501 | + window_size=None, |
| 502 | + softcap=0.0, # 0.0 means deactivated |
| 503 | + alibi_slopes=None, |
| 504 | + deterministic=False, |
| 505 | + return_attn_probs=False, |
| 506 | + block_table=None, |
| 507 | + return_softmax_lse=False, |
| 508 | + out=None, |
| 509 | + fa_version: int = 2, |
| 510 | +): |
| 511 | + """dropout_p should be set to 0.0 during evaluation |
| 512 | + Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads |
| 513 | + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. |
| 514 | + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head |
| 515 | + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. |
| 516 | +
|
| 517 | + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. |
| 518 | + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: |
| 519 | + 1 1 1 1 0 |
| 520 | + 1 1 1 1 1 |
| 521 | + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: |
| 522 | + 0 0 |
| 523 | + 0 0 |
| 524 | + 0 0 |
| 525 | + 1 0 |
| 526 | + 1 1 |
| 527 | + If the row of the mask is all zero, the output will be zero. |
| 528 | +
|
| 529 | + If window_size != (-1, -1), implements sliding window local attention. Query at position i |
| 530 | + will only attend to keys between |
| 531 | + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. |
| 532 | +
|
| 533 | + Arguments: |
| 534 | + q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. |
| 535 | + k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. |
| 536 | + v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. |
| 537 | + cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths |
| 538 | + of the sequences in the batch, used to index into q. |
| 539 | + cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths |
| 540 | + of the sequences in the batch, used to index into kv. |
| 541 | + max_seqlen_q: int. Maximum query sequence length in the batch. |
| 542 | + max_seqlen_k: int. Maximum key sequence length in the batch. |
| 543 | + dropout_p: float. Dropout probability. |
| 544 | + softmax_scale: float. The scaling of QK^T before applying softmax. |
| 545 | + Default to 1 / sqrt(headdim). |
| 546 | + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). |
| 547 | + window_size: (left, right). If not (-1, -1), implements sliding window local attention. |
| 548 | + softcap: float. Anything > 0 activates softcapping attention. |
| 549 | + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of |
| 550 | + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) |
| 551 | + is added to the attention score of query i and key j. |
| 552 | + deterministic: bool. Whether to use the deterministic implementation of the backward pass, |
| 553 | + which is slightly slower and uses more memory. The forward pass is always deterministic. |
| 554 | + return_attn_probs: bool. Whether to return the attention probabilities. This option is for |
| 555 | + testing only. The returned probabilities are not guaranteed to be correct |
| 556 | + (they might not have the right scaling). |
| 557 | + Return: |
| 558 | + out: (total, nheads, headdim). |
| 559 | + softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The |
| 560 | + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax |
| 561 | + normalization factor). |
| 562 | + """ |
| 563 | + assert ( |
| 564 | + cu_seqlens_k is not None or seqused_k is not None |
| 565 | + ), "cu_seqlens_k or seqused_k must be provided" |
| 566 | + assert ( |
| 567 | + cu_seqlens_k is None or seqused_k is None |
| 568 | + ), "cu_seqlens_k and seqused_k cannot be provided at the same time" |
| 569 | + assert ( |
| 570 | + block_table is None or seqused_k is not None |
| 571 | + ), "seqused_k must be provided if block_table is provided" |
| 572 | + |
| 573 | + if softmax_scale is None: |
| 574 | + softmax_scale = q.shape[-1] ** (-0.5) |
| 575 | + # custom op does not support non-tuple input |
| 576 | + if window_size is None: |
| 577 | + real_window_size = (-1, -1) |
| 578 | + else: |
| 579 | + assert len(window_size) == 2 |
| 580 | + real_window_size = (window_size[0], window_size[1]) |
| 581 | + q, k, v = [maybe_contiguous(x) for x in (q, k, v)] |
| 582 | + |
| 583 | + dummy_cu_seqlens_k = torch.empty_like(cu_seqlens_q) |
| 584 | + |
| 585 | + assert fa_version == 2, "Only FA2 is implemented." |
| 586 | + |
| 587 | + out = torch.empty_like(q) |
| 588 | + |
| 589 | + out, q, k, v, softmax_lse, *_ = mha_varlan_fwd( |
| 590 | + q, |
| 591 | + k, |
| 592 | + v, |
| 593 | + out, |
| 594 | + cu_seqlens_q, |
| 595 | + # cu_seqlens_k not used since we use seqused_k, but flash_api.cpp |
| 596 | + # still wants it so we pass all zeros |
| 597 | + dummy_cu_seqlens_k if cu_seqlens_k is None else cu_seqlens_k, |
| 598 | + seqused_k, |
| 599 | + None, |
| 600 | + block_table, |
| 601 | + alibi_slopes, |
| 602 | + max_seqlen_q, |
| 603 | + max_seqlen_k, |
| 604 | + dropout_p, |
| 605 | + softmax_scale, |
| 606 | + False, |
| 607 | + causal, |
| 608 | + real_window_size[0], |
| 609 | + real_window_size[1], |
| 610 | + softcap, |
| 611 | + return_softmax_lse and dropout_p > 0, |
| 612 | + None, |
| 613 | + ) |
| 614 | + |
| 615 | + return (out, softmax_lse) if return_softmax_lse else out |
0 commit comments