| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Merge AITER segmented MLA split-K partials into output plus natural-log LSE. |
| 4 | |
| 5 | This mirrors AITER's own segment reduction, including its |
| 6 | ``tiles_per_segment = cdiv(seq_len, NUM_SEGMENTS * TILE_SIZE)`` partitioning, so |
| 7 | it has to move together with the ``skip_reduce=True`` call in the AITER MLA |
| 8 | backend. It is a rank-local split-K merge and unrelated to any collective |
| 9 | reduce; the natural-log LSE it returns is what the cross-rank DCP merge |
| 10 | consumes. |
| 11 | """ |
| 12 | |
| 13 | import torch |
| 14 | |
| 15 | from vllm.triton_utils import LOGE2, tl, triton |
| 16 | |
| 17 | |
| 18 | @triton.jit |
| 19 | def _merge_mla_segments_kernel( |
| 20 | out_ptr, |
| 21 | lse_ptr, |
| 22 | segm_output_ptr, |
| 23 | segm_max_ptr, |
| 24 | segm_expsum_ptr, |
| 25 | seq_lens_ptr, |
| 26 | num_query_heads: tl.constexpr, |
| 27 | out_stride0: tl.int64, |
| 28 | out_stride1: tl.int64, |
| 29 | lse_stride0: tl.int64, |
| 30 | TILE_SIZE: tl.constexpr, |
| 31 | KV_LORA_RANK: tl.constexpr, |
| 32 | NUM_SEGMENTS_PER_SEQ: tl.constexpr, |
| 33 | LOGE2: tl.constexpr, |
| 34 | ): |
| 35 | token_idx = tl.program_id(0) |
| 36 | head_idx = tl.program_id(1) |
| 37 | |
| 38 | seq_len = tl.load(seq_lens_ptr + token_idx) |
| 39 | tiles_per_segment = tl.maximum( |
| 40 | 1, tl.cdiv(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE) |
| 41 | ) |
| 42 | active_segments = tl.where( |
| 43 | seq_len > 0, |
| 44 | tl.cdiv(seq_len, tiles_per_segment * TILE_SIZE), |
| 45 | 0, |
| 46 | ) |
| 47 | segment_offsets = tl.arange(0, NUM_SEGMENTS_PER_SEQ) |
| 48 | segment_mask = segment_offsets < active_segments |
| 49 | |
| 50 | segment_base = ( |
| 51 | token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) |
| 52 | + head_idx * NUM_SEGMENTS_PER_SEQ |
| 53 | ) |
| 54 | segment_indices = segment_base + segment_offsets |
| 55 | segment_max = tl.load( |
| 56 | segm_max_ptr + segment_indices, |
| 57 | mask=segment_mask, |
| 58 | other=float("-inf"), |
| 59 | ) |
| 60 | overall_max = tl.max(segment_max) |
| 61 | |
| 62 | segment_expsum = tl.load( |
| 63 | segm_expsum_ptr + segment_indices, |
| 64 | mask=segment_mask, |
| 65 | other=0.0, |
| 66 | ) |
| 67 | segment_weight = tl.where( |
| 68 | segment_mask, |
| 69 | tl.exp2(segment_max - overall_max), |
| 70 | 0.0, |
| 71 | ) |
| 72 | overall_expsum = tl.sum(segment_expsum * segment_weight) |
| 73 | |
| 74 | output_offsets = ( |
| 75 | token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ * KV_LORA_RANK) |
| 76 | + head_idx * (NUM_SEGMENTS_PER_SEQ * KV_LORA_RANK) |
| 77 | + segment_offsets[:, None] * KV_LORA_RANK |
| 78 | + tl.arange(0, KV_LORA_RANK)[None, :] |
| 79 | ) |
| 80 | segment_output = tl.load( |
| 81 | segm_output_ptr + output_offsets, |
| 82 | mask=segment_mask[:, None], |
| 83 | other=0.0, |
| 84 | ) |
| 85 | accumulator = tl.sum(segment_output * segment_weight[:, None], axis=0) |
| 86 | accumulator = tl.where(overall_expsum == 0.0, 0.0, accumulator / overall_expsum) |
| 87 | lse = tl.where( |
| 88 | overall_expsum == 0.0, |
| 89 | float("-inf"), |
| 90 | (overall_max + tl.log2(overall_expsum)) * LOGE2, |
| 91 | ) |
| 92 | |
| 93 | output_indices = ( |
| 94 | token_idx * out_stride0 + head_idx * out_stride1 + tl.arange(0, KV_LORA_RANK) |
| 95 | ) |
| 96 | tl.store(out_ptr + output_indices, accumulator) |
| 97 | tl.store(lse_ptr + token_idx * lse_stride0 + head_idx, lse) |
| 98 | |
| 99 | |
| 100 | def merge_mla_segments_triton( |
| 101 | segm_output: torch.Tensor, |
| 102 | segm_max: torch.Tensor, |
| 103 | segm_expsum: torch.Tensor, |
| 104 | seq_lens: torch.Tensor, |
| 105 | tile_size: int, |
| 106 | out_dtype: torch.dtype, |
| 107 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 108 | """Merge AITER base-2 segment partials into output and natural-log LSE.""" |
| 109 | num_tokens, num_heads, num_segments, kv_lora_rank = segm_output.shape |
| 110 | output = torch.empty( |
| 111 | (num_tokens, num_heads, kv_lora_rank), |
| 112 | dtype=out_dtype, |
| 113 | device=segm_output.device, |
| 114 | ) |
| 115 | lse = torch.empty( |
| 116 | (num_tokens, num_heads), |
| 117 | dtype=torch.float32, |
| 118 | device=segm_output.device, |
| 119 | ) |
| 120 | _merge_mla_segments_kernel[(num_tokens, num_heads)]( |
| 121 | output, |
| 122 | lse, |
| 123 | segm_output, |
| 124 | segm_max, |
| 125 | segm_expsum, |
| 126 | seq_lens, |
| 127 | num_query_heads=num_heads, |
| 128 | out_stride0=output.stride(0), |
| 129 | out_stride1=output.stride(1), |
| 130 | lse_stride0=lse.stride(0), |
| 131 | TILE_SIZE=tile_size, |
| 132 | KV_LORA_RANK=kv_lora_rank, |
| 133 | NUM_SEGMENTS_PER_SEQ=num_segments, |
| 134 | LOGE2=LOGE2, |
| 135 | ) |
| 136 | return output, lse |