| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Weight permutations for FlashMLA's mega-attention kernel. |
| 4 | |
| 5 | The kernel reads Q with 16-element head-dim chunks interleaved across heads and |
| 6 | writes O with 32-element chunks interleaved across the 8 heads of a ``wo_a`` |
| 7 | group. Every permutation here satisfies ``fused = standard[perm]``, and is |
| 8 | applied once to ``wq_b`` rows and ``wo_a`` columns at load time so the |
| 9 | surrounding GEMMs produce and consume the kernel's layouts directly -- no |
| 10 | per-step shuffle. |
| 11 | """ |
| 12 | |
| 13 | import torch |
| 14 | |
| 15 | HEAD_DIM = 512 |
| 16 | Q_CHUNK = 16 |
| 17 | O_CHUNK = 32 |
| 18 | WV_GROUP_SIZE = 8 |
| 19 | |
| 20 | |
| 21 | def q_fused_permutation(num_heads: int, head_dim: int = HEAD_DIM) -> torch.Tensor: |
| 22 | """``fused[(d // 16) * (H * 16) + h * 16 + d % 16] = standard[h * D + d]``.""" |
| 23 | h = torch.arange(num_heads).view(num_heads, 1) |
| 24 | d = torch.arange(head_dim).view(1, head_dim) |
| 25 | fused_index = (d // Q_CHUNK) * (num_heads * Q_CHUNK) + h * Q_CHUNK + d % Q_CHUNK |
| 26 | perm = torch.empty(num_heads * head_dim, dtype=torch.long) |
| 27 | perm[fused_index.reshape(-1)] = torch.arange(num_heads * head_dim) |
| 28 | return perm |
| 29 | |
| 30 | |
| 31 | def o_fused_permutation( |
| 32 | heads_per_group: int = WV_GROUP_SIZE, head_dim: int = HEAD_DIM |
| 33 | ) -> torch.Tensor: |
| 34 | """``fused[(c * G + h) * 32 + j] = standard[h * D + c * 32 + j]`` per group.""" |
| 35 | h = torch.arange(heads_per_group).view(-1, 1, 1) |
| 36 | c = torch.arange(head_dim // O_CHUNK).view(1, -1, 1) |
| 37 | j = torch.arange(O_CHUNK).view(1, 1, -1) |
| 38 | fused_index = (c * heads_per_group + h) * O_CHUNK + j |
| 39 | perm = torch.empty(heads_per_group * head_dim, dtype=torch.long) |
| 40 | perm[fused_index.reshape(-1)] = torch.arange(heads_per_group * head_dim) |
| 41 | return perm |
| 42 | |
| 43 | |
| 44 | def o_fused_chunk_permutation( |
| 45 | heads_per_group: int = WV_GROUP_SIZE, head_dim: int = HEAD_DIM |
| 46 | ) -> torch.Tensor: |
| 47 | """Per-32-element-chunk form of :func:`o_fused_permutation`, for scales.""" |
| 48 | return o_fused_permutation(heads_per_group, head_dim)[::O_CHUNK] // O_CHUNK |
| 49 | |
| 50 | |
| 51 | def permute_q_to_fused(q: torch.Tensor) -> torch.Tensor: |
| 52 | """``[N, H, D]`` standard layout -> the same shape in the fused layout.""" |
| 53 | n, h, d = q.shape |
| 54 | perm = q_fused_permutation(h, d).to(q.device) |
| 55 | return q.reshape(n, h * d)[:, perm].view(n, h, d) |
| 56 | |
| 57 | |
| 58 | def _bytes_view(t: torch.Tensor) -> torch.Tensor: |
| 59 | """Byte view of a 1-byte-element tensor, so fp8/ue8m0 can be gathered.""" |
| 60 | return t.view(torch.uint8) if t.element_size() == 1 else t |
| 61 | |
| 62 | |
| 63 | def permute_wq_b_( |
| 64 | weight: torch.Tensor, weight_scale: torch.Tensor, num_local_heads: int |
| 65 | ) -> None: |
| 66 | """Permute an MXFP8 ``wq_b`` shard's rows and per-row scales, in place.""" |
| 67 | head_dim = weight.shape[0] // num_local_heads |
| 68 | perm = q_fused_permutation(num_local_heads, head_dim).to(weight.device) |
| 69 | for t in (weight, weight_scale): |
| 70 | b = _bytes_view(t) |
| 71 | b.copy_(b[perm]) |
| 72 | |
| 73 | |
| 74 | def permute_wo_a_( |
| 75 | weight: torch.Tensor, |
| 76 | weight_scale: torch.Tensor, |
| 77 | heads_per_group: int = WV_GROUP_SIZE, |
| 78 | ) -> None: |
| 79 | """Permute an MXFP8 ``wo_a`` shard's input columns and scales, in place.""" |
| 80 | head_dim = weight.shape[1] // heads_per_group |
| 81 | perm = o_fused_permutation(heads_per_group, head_dim).to(weight.device) |
| 82 | w = _bytes_view(weight) |
| 83 | w.copy_(w[:, perm]) |
| 84 | chunk_perm = o_fused_chunk_permutation(heads_per_group, head_dim) |
| 85 | s = _bytes_view(weight_scale) |
| 86 | s.copy_(s[:, chunk_perm.to(weight.device)]) |