已合并
harden sort_chunks_by_idx memory access #3316
guofanfeng23创建于 3月11日
harden sort_chunks_by_idx memory access #3316
已合并
共 3 个文件变更+94-294
| @@ -3,10 +3,21 @@ | |||
| 3 | 3 | ||
| 4 | from typing import Tuple | 4 | from typing import Tuple |
| 5 | import torch | 5 | import torch |
| 6 | +import functools | ||
| 6 | import triton | 7 | import triton |
| 7 | import triton.language as tl | 8 | import triton.language as tl |
| 8 | 9 | ||
| 9 | 10 | ||
| 11 | + | ||
| 12 | +def get_vector_num() -> int: | ||
| 13 | + from triton.runtime import driver | ||
| 14 | + import torch_npu | ||
| 15 | + | ||
| 16 | + device = torch_npu.npu.current_device() | ||
| 17 | + properties = driver.active.utils.get_device_properties(device) | ||
| 18 | + return properties["num_vectorcore"] | ||
| 19 | + | ||
| 20 | + | ||
| 10 | 21 | ||
| 11 | configs=[ | 22 | configs=[ |
| 12 | triton.Config({'SUB_BLOCK_SIZE': 8}), | 23 | triton.Config({'SUB_BLOCK_SIZE': 8}), |
| @@ -37,31 +48,37 @@ def _make_chunk_sort_map_kernel( | |||
| 37 | IDX_LOAD_WIDTH: tl.constexpr, | 48 | IDX_LOAD_WIDTH: tl.constexpr, |
| 38 | ): | 49 | ): |
| 39 | pid = tl.program_id(0) | 50 | pid = tl.program_id(0) |
| 51 | + pid_start = pid * BLOCK_SIZE | ||
| 40 | 52 | ||
| 41 | load_split_offset = tl.arange(0, IDX_LOAD_WIDTH) | 53 | load_split_offset = tl.arange(0, IDX_LOAD_WIDTH) |
| 42 | - input_split_sizes = tl.load( | 54 | + input_split_sizes = tl.load(split_sizes_ptr + load_split_offset, mask=load_split_offset < num_splits, other=0).to( |
| 43 | - split_sizes_ptr + load_split_offset, mask=load_split_offset < num_splits, other=0 | 55 | + tl.float32 |
| 44 | - ).to(tl.float32) | 56 | + ) |
| 45 | input_split_sizes_cumsum = tl.cumsum(input_split_sizes) | 57 | input_split_sizes_cumsum = tl.cumsum(input_split_sizes) |
| 46 | 58 | ||
| 47 | for off in range(0, BLOCK_SIZE, SUB_BLOCK_SIZE): | 59 | for off in range(0, BLOCK_SIZE, SUB_BLOCK_SIZE): |
| 48 | - token_offsets = pid * BLOCK_SIZE + off + tl.arange(0, SUB_BLOCK_SIZE) | 60 | + block_offsets = off + tl.arange(0, SUB_BLOCK_SIZE) |
| 61 | + token_offsets = pid_start + block_offsets | ||
| 49 | token_offsets_cmp = token_offsets.to(tl.float32) | 62 | token_offsets_cmp = token_offsets.to(tl.float32) |
| 50 | 63 | ||
| 51 | input_split_sizes_mask = tl.where(input_split_sizes_cumsum[None, :] <= token_offsets_cmp[:, None], 1, 0) | 64 | input_split_sizes_mask = tl.where(input_split_sizes_cumsum[None, :] <= token_offsets_cmp[:, None], 1, 0) |
| 52 | input_chunk_indices = tl.sum(input_split_sizes_mask, axis=-1) | 65 | input_chunk_indices = tl.sum(input_split_sizes_mask, axis=-1) |
| 53 | cumsum_mask = input_chunk_indices < (num_splits + 1) | 66 | cumsum_mask = input_chunk_indices < (num_splits + 1) |
| 54 | - input_split_sizes_presums = tl.load(cumsum_ptr + input_chunk_indices, mask=cumsum_mask, other=0) | 67 | + safe_indices = tl.where(cumsum_mask, input_chunk_indices, 0) |
| 68 | + input_split_sizes_presums = tl.load(cumsum_ptr + safe_indices, mask=cumsum_mask, other=0) | ||
| 55 | 69 | ||
| 56 | inv_mask = input_chunk_indices < num_splits | 70 | inv_mask = input_chunk_indices < num_splits |
| 57 | - output_chunk_indices = tl.load(inverse_sorted_indices_ptr + input_chunk_indices, mask=inv_mask, other=0) | 71 | + safe_inv_indices = tl.where(inv_mask, input_chunk_indices, 0) |
| 72 | + output_chunk_indices = tl.load(inverse_sorted_indices_ptr + safe_inv_indices, mask=inv_mask, other=0) | ||
| 58 | 73 | ||
| 59 | output_chunk_mask = output_chunk_indices < (num_splits + 1) | 74 | output_chunk_mask = output_chunk_indices < (num_splits + 1) |
| 60 | - output_presums = tl.load(output_cumsum_ptr + output_chunk_indices, mask=output_chunk_mask, other=0) | 75 | + safe_out_indices = tl.where(output_chunk_mask, output_chunk_indices, 0) |
| 76 | + output_presums = tl.load(output_cumsum_ptr + safe_out_indices, mask=output_chunk_mask, other=0) | ||
| 61 | 77 | ||
| 62 | dst_rows = output_presums + token_offsets_cmp - input_split_sizes_presums | 78 | dst_rows = output_presums + token_offsets_cmp - input_split_sizes_presums |
| 63 | - store_mask = token_offsets < num_tokens | 79 | + store_mask = (block_offsets < BLOCK_SIZE) & (token_offsets < num_tokens) |
| 64 | - tl.store(dst_rows_ptr + token_offsets, dst_rows, mask=store_mask) | 80 | + safe_token_offsets = tl.where(store_mask, token_offsets, 0) |
| 81 | + tl.store(dst_rows_ptr + safe_token_offsets, dst_rows, mask=store_mask) | ||
| 65 | 82 | ||
| 66 | 83 | ||
| 67 | def make_chunk_sort_map( | 84 | def make_chunk_sort_map( |
| @@ -85,7 +102,7 @@ def make_chunk_sort_map( | |||
| 85 | Number of splits of split_sizes and sorted_indices. | 102 | Number of splits of split_sizes and sorted_indices. |
| 86 | """ | 103 | """ |
| 87 | row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device="npu") | 104 | row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device="npu") |
| 88 | - num_blocks = min(48, num_tokens) | 105 | + num_blocks = min(get_vector_num(), num_tokens) |
| 89 | block_size = triton.cdiv(num_tokens, num_blocks) | 106 | block_size = triton.cdiv(num_tokens, num_blocks) |
| 90 | grid = (num_blocks, 1, 1) | 107 | grid = (num_blocks, 1, 1) |
| 91 | 108 | ||
| @@ -94,11 +111,13 @@ def make_chunk_sort_map( | |||
| 94 | cumsum[1:] = split_sizes.cumsum(dim=0, dtype=torch.int32) | 111 | cumsum[1:] = split_sizes.cumsum(dim=0, dtype=torch.int32) |
| 95 | 112 | ||
| 96 | inverse_sorted_indices = torch.empty(num_splits, dtype=torch.int32, device="npu") | 113 | inverse_sorted_indices = torch.empty(num_splits, dtype=torch.int32, device="npu") |
| 97 | - inverse_sorted_indices.scatter_(dim=0, index=sorted_indices, src=torch.arange(num_splits, dtype=torch.int32, device="npu")) | 114 | + inverse_sorted_indices.scatter_( |
| 115 | + dim=0, index=sorted_indices, src=torch.arange(num_splits, dtype=torch.int32, device="npu") | ||
| 116 | + ) | ||
| 98 | 117 | ||
| 99 | output_split_sizes = split_sizes[sorted_indices] | 118 | output_split_sizes = split_sizes[sorted_indices] |
| 100 | output_cumsum = torch.empty(split_sizes.size(0) + 1, dtype=torch.int32, device="npu") | 119 | output_cumsum = torch.empty(split_sizes.size(0) + 1, dtype=torch.int32, device="npu") |
| 101 | - output_cumsum[0] = 0 | 120 | + output_cumsum[0] = 0 |
| 102 | output_cumsum[1:] = output_split_sizes.cumsum(dim=0, dtype=torch.int32) | 121 | output_cumsum[1:] = output_split_sizes.cumsum(dim=0, dtype=torch.int32) |
| 103 | 122 | ||
| 104 | _make_chunk_sort_map_kernel[grid]( | 123 | _make_chunk_sort_map_kernel[grid]( |
| @@ -144,11 +163,11 @@ def _sort_chunks_by_map_kernel( | |||
| 144 | 163 | ||
| 145 | for i in range(pid_start, pid_end): | 164 | for i in range(pid_start, pid_end): |
| 146 | if FORWARD: | 165 | if FORWARD: |
| 147 | - src_row = i | 166 | + src_row = i.to(tl.int64) |
| 148 | dst_row = tl.load(row_id_map_ptr + i).to(tl.int64) | 167 | dst_row = tl.load(row_id_map_ptr + i).to(tl.int64) |
| 149 | else: | 168 | else: |
| 150 | src_row = tl.load(row_id_map_ptr + i).to(tl.int64) | 169 | src_row = tl.load(row_id_map_ptr + i).to(tl.int64) |
| 151 | - dst_row = i | 170 | + dst_row = i.to(tl.int64) |
| 152 | 171 | ||
| 153 | current_offset = tl.arange(0, hidden_size) | 172 | current_offset = tl.arange(0, hidden_size) |
| 154 | input_offsets = src_row * stride_input_token + current_offset * stride_input_hidden | 173 | input_offsets = src_row * stride_input_token + current_offset * stride_input_hidden |
| @@ -196,7 +215,7 @@ def sort_chunks_by_map( | |||
| 196 | else: | 215 | else: |
| 197 | permuted_probs = None | 216 | permuted_probs = None |
| 198 | 217 | ||
| 199 | - num_blocks = 48 | 218 | + num_blocks = min(get_vector_num(), num_tokens) |
| 200 | block_size = triton.cdiv(num_tokens, num_blocks) | 219 | block_size = triton.cdiv(num_tokens, num_blocks) |
| 201 | grid = (num_blocks, 1, 1) | 220 | grid = (num_blocks, 1, 1) |
| 202 | _sort_chunks_by_map_kernel[grid]( | 221 | _sort_chunks_by_map_kernel[grid]( |
| @@ -219,7 +238,7 @@ def sort_chunks_by_map( | |||
| 219 | ) | 238 | ) |
| 220 | 239 | ||
| 221 | return output, permuted_probs | 240 | return output, permuted_probs |
| 222 | - | 241 | + |
| 223 | 242 | ||
| 224 | class _moe_chunk_sort(torch.autograd.Function): | 243 | class _moe_chunk_sort(torch.autograd.Function): |
| 225 | """functional MoE chunk permute""" | 244 | """functional MoE chunk permute""" |
| @@ -247,7 +266,7 @@ class _moe_chunk_sort(torch.autograd.Function): | |||
| 247 | raise AssertionError('This operation needs NPU') | 266 | raise AssertionError('This operation needs NPU') |
| 248 | if sorted_idxs.device.type != 'npu': | 267 | if sorted_idxs.device.type != 'npu': |
| 249 | raise AssertionError('This operation needs NPU') | 268 | raise AssertionError('This operation needs NPU') |
| 250 | - if probs is not None: | 269 | + if probs is not None: |
| 251 | if probs.device.type != 'npu': | 270 | if probs.device.type != 'npu': |
| 252 | raise AssertionError('This operation needs NPU') | 271 | raise AssertionError('This operation needs NPU') |
| 253 | 272 | ||
| @@ -345,4 +364,4 @@ def moe_sort_chunks_by_index_with_probs( | |||
| 345 | Chunk indices used to permute the chunks. | 364 | Chunk indices used to permute the chunks. |
| 346 | """ | 365 | """ |
| 347 | output, permuted_probs = _moe_chunk_sort.apply(inp, split_sizes, sorted_index, probs) | 366 | output, permuted_probs = _moe_chunk_sort.apply(inp, split_sizes, sorted_index, probs) |
| 348 | - return output, permuted_probs | 367 | + return output, permuted_probs |
| @@ -3,10 +3,21 @@ | |||
| 3 | 3 | ||
| 4 | from typing import Tuple | 4 | from typing import Tuple |
| 5 | import torch | 5 | import torch |
| 6 | +import functools | ||
| 6 | import triton | 7 | import triton |
| 7 | import triton.language as tl | 8 | import triton.language as tl |
| 8 | 9 | ||
| 9 | 10 | ||
| 11 | + | ||
| 12 | +def get_vector_num() -> int: | ||
| 13 | + from triton.runtime import driver | ||
| 14 | + import torch_npu | ||
| 15 | + | ||
| 16 | + device = torch_npu.npu.current_device() | ||
| 17 | + properties = driver.active.utils.get_device_properties(device) | ||
| 18 | + return properties["num_vectorcore"] | ||
| 19 | + | ||
| 20 | + | ||
| 10 | 21 | ||
| 11 | configs=[ | 22 | configs=[ |
| 12 | triton.Config({'SUB_BLOCK_SIZE': 8}), | 23 | triton.Config({'SUB_BLOCK_SIZE': 8}), |
| @@ -37,31 +48,37 @@ def _make_chunk_sort_map_kernel( | |||
| 37 | IDX_LOAD_WIDTH: tl.constexpr, | 48 | IDX_LOAD_WIDTH: tl.constexpr, |
| 38 | ): | 49 | ): |
| 39 | pid = tl.program_id(0) | 50 | pid = tl.program_id(0) |
| 51 | + pid_start = pid * BLOCK_SIZE | ||
| 40 | 52 | ||
| 41 | load_split_offset = tl.arange(0, IDX_LOAD_WIDTH) | 53 | load_split_offset = tl.arange(0, IDX_LOAD_WIDTH) |
| 42 | - input_split_sizes = tl.load( | 54 | + input_split_sizes = tl.load(split_sizes_ptr + load_split_offset, mask=load_split_offset < num_splits, other=0).to( |
| 43 | - split_sizes_ptr + load_split_offset, mask=load_split_offset < num_splits, other=0 | 55 | + tl.float32 |
| 44 | - ).to(tl.float32) | 56 | + ) |
| 45 | input_split_sizes_cumsum = tl.cumsum(input_split_sizes) | 57 | input_split_sizes_cumsum = tl.cumsum(input_split_sizes) |
| 46 | 58 | ||
| 47 | for off in range(0, BLOCK_SIZE, SUB_BLOCK_SIZE): | 59 | for off in range(0, BLOCK_SIZE, SUB_BLOCK_SIZE): |
| 48 | - token_offsets = pid * BLOCK_SIZE + off + tl.arange(0, SUB_BLOCK_SIZE) | 60 | + block_offsets = off + tl.arange(0, SUB_BLOCK_SIZE) |
| 61 | + token_offsets = pid_start + block_offsets | ||
| 49 | token_offsets_cmp = token_offsets.to(tl.float32) | 62 | token_offsets_cmp = token_offsets.to(tl.float32) |
| 50 | 63 | ||
| 51 | input_split_sizes_mask = tl.where(input_split_sizes_cumsum[None, :] <= token_offsets_cmp[:, None], 1, 0) | 64 | input_split_sizes_mask = tl.where(input_split_sizes_cumsum[None, :] <= token_offsets_cmp[:, None], 1, 0) |
| 52 | input_chunk_indices = tl.sum(input_split_sizes_mask, axis=-1) | 65 | input_chunk_indices = tl.sum(input_split_sizes_mask, axis=-1) |
| 53 | cumsum_mask = input_chunk_indices < (num_splits + 1) | 66 | cumsum_mask = input_chunk_indices < (num_splits + 1) |
| 54 | - input_split_sizes_presums = tl.load(cumsum_ptr + input_chunk_indices, mask=cumsum_mask, other=0) | 67 | + safe_indices = tl.where(cumsum_mask, input_chunk_indices, 0) |
| 68 | + input_split_sizes_presums = tl.load(cumsum_ptr + safe_indices, mask=cumsum_mask, other=0) | ||
| 55 | 69 | ||
| 56 | inv_mask = input_chunk_indices < num_splits | 70 | inv_mask = input_chunk_indices < num_splits |
| 57 | - output_chunk_indices = tl.load(inverse_sorted_indices_ptr + input_chunk_indices, mask=inv_mask, other=0) | 71 | + safe_inv_indices = tl.where(inv_mask, input_chunk_indices, 0) |
| 72 | + output_chunk_indices = tl.load(inverse_sorted_indices_ptr + safe_inv_indices, mask=inv_mask, other=0) | ||
| 58 | 73 | ||
| 59 | output_chunk_mask = output_chunk_indices < (num_splits + 1) | 74 | output_chunk_mask = output_chunk_indices < (num_splits + 1) |
| 60 | - output_presums = tl.load(output_cumsum_ptr + output_chunk_indices, mask=output_chunk_mask, other=0) | 75 | + safe_out_indices = tl.where(output_chunk_mask, output_chunk_indices, 0) |
| 76 | + output_presums = tl.load(output_cumsum_ptr + safe_out_indices, mask=output_chunk_mask, other=0) | ||
| 61 | 77 | ||
| 62 | dst_rows = output_presums + token_offsets_cmp - input_split_sizes_presums | 78 | dst_rows = output_presums + token_offsets_cmp - input_split_sizes_presums |
| 63 | - store_mask = token_offsets < num_tokens | 79 | + store_mask = (block_offsets < BLOCK_SIZE) & (token_offsets < num_tokens) |
| 64 | - tl.store(dst_rows_ptr + token_offsets, dst_rows, mask=store_mask) | 80 | + safe_token_offsets = tl.where(store_mask, token_offsets, 0) |
| 81 | + tl.store(dst_rows_ptr + safe_token_offsets, dst_rows, mask=store_mask) | ||
| 65 | 82 | ||
| 66 | 83 | ||
| 67 | def make_chunk_sort_map( | 84 | def make_chunk_sort_map( |
| @@ -85,7 +102,7 @@ def make_chunk_sort_map( | |||
| 85 | Number of splits of split_sizes and sorted_indices. | 102 | Number of splits of split_sizes and sorted_indices. |
| 86 | """ | 103 | """ |
| 87 | row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device="npu") | 104 | row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device="npu") |
| 88 | - num_blocks = min(48, num_tokens) | 105 | + num_blocks = min(get_vector_num(), num_tokens) |
| 89 | block_size = triton.cdiv(num_tokens, num_blocks) | 106 | block_size = triton.cdiv(num_tokens, num_blocks) |
| 90 | grid = (num_blocks, 1, 1) | 107 | grid = (num_blocks, 1, 1) |
| 91 | 108 | ||
| @@ -94,11 +111,13 @@ def make_chunk_sort_map( | |||
| 94 | cumsum[1:] = split_sizes.cumsum(dim=0, dtype=torch.int32) | 111 | cumsum[1:] = split_sizes.cumsum(dim=0, dtype=torch.int32) |
| 95 | 112 | ||
| 96 | inverse_sorted_indices = torch.empty(num_splits, dtype=torch.int32, device="npu") | 113 | inverse_sorted_indices = torch.empty(num_splits, dtype=torch.int32, device="npu") |
| 97 | - inverse_sorted_indices.scatter_(dim=0, index=sorted_indices, src=torch.arange(num_splits, dtype=torch.int32, device="npu")) | 114 | + inverse_sorted_indices.scatter_( |
| 115 | + dim=0, index=sorted_indices, src=torch.arange(num_splits, dtype=torch.int32, device="npu") | ||
| 116 | + ) | ||
| 98 | 117 | ||
| 99 | output_split_sizes = split_sizes[sorted_indices] | 118 | output_split_sizes = split_sizes[sorted_indices] |
| 100 | output_cumsum = torch.empty(split_sizes.size(0) + 1, dtype=torch.int32, device="npu") | 119 | output_cumsum = torch.empty(split_sizes.size(0) + 1, dtype=torch.int32, device="npu") |
| 101 | - output_cumsum[0] = 0 | 120 | + output_cumsum[0] = 0 |
| 102 | output_cumsum[1:] = output_split_sizes.cumsum(dim=0, dtype=torch.int32) | 121 | output_cumsum[1:] = output_split_sizes.cumsum(dim=0, dtype=torch.int32) |
| 103 | 122 | ||
| 104 | _make_chunk_sort_map_kernel[grid]( | 123 | _make_chunk_sort_map_kernel[grid]( |
| @@ -144,11 +163,11 @@ def _sort_chunks_by_map_kernel( | |||
| 144 | 163 | ||
| 145 | for i in range(pid_start, pid_end): | 164 | for i in range(pid_start, pid_end): |
| 146 | if FORWARD: | 165 | if FORWARD: |
| 147 | - src_row = i | 166 | + src_row = i.to(tl.int64) |
| 148 | dst_row = tl.load(row_id_map_ptr + i).to(tl.int64) | 167 | dst_row = tl.load(row_id_map_ptr + i).to(tl.int64) |
| 149 | else: | 168 | else: |
| 150 | src_row = tl.load(row_id_map_ptr + i).to(tl.int64) | 169 | src_row = tl.load(row_id_map_ptr + i).to(tl.int64) |
| 151 | - dst_row = i | 170 | + dst_row = i.to(tl.int64) |
| 152 | 171 | ||
| 153 | current_offset = tl.arange(0, hidden_size) | 172 | current_offset = tl.arange(0, hidden_size) |
| 154 | input_offsets = src_row * stride_input_token + current_offset * stride_input_hidden | 173 | input_offsets = src_row * stride_input_token + current_offset * stride_input_hidden |
| @@ -196,7 +215,7 @@ def sort_chunks_by_map( | |||
| 196 | else: | 215 | else: |
| 197 | permuted_probs = None | 216 | permuted_probs = None |
| 198 | 217 | ||
| 199 | - num_blocks = 48 | 218 | + num_blocks = min(get_vector_num(), num_tokens) |
| 200 | block_size = triton.cdiv(num_tokens, num_blocks) | 219 | block_size = triton.cdiv(num_tokens, num_blocks) |
| 201 | grid = (num_blocks, 1, 1) | 220 | grid = (num_blocks, 1, 1) |
| 202 | _sort_chunks_by_map_kernel[grid]( | 221 | _sort_chunks_by_map_kernel[grid]( |
| @@ -219,7 +238,7 @@ def sort_chunks_by_map( | |||
| 219 | ) | 238 | ) |
| 220 | 239 | ||
| 221 | return output, permuted_probs | 240 | return output, permuted_probs |
| 222 | - | 241 | + |
| 223 | 242 | ||
| 224 | class _moe_chunk_sort(torch.autograd.Function): | 243 | class _moe_chunk_sort(torch.autograd.Function): |
| 225 | """functional MoE chunk permute""" | 244 | """functional MoE chunk permute""" |
| @@ -247,7 +266,7 @@ class _moe_chunk_sort(torch.autograd.Function): | |||
| 247 | raise AssertionError('This operation needs NPU') | 266 | raise AssertionError('This operation needs NPU') |
| 248 | if sorted_idxs.device.type != 'npu': | 267 | if sorted_idxs.device.type != 'npu': |
| 249 | raise AssertionError('This operation needs NPU') | 268 | raise AssertionError('This operation needs NPU') |
| 250 | - if probs is not None: | 269 | + if probs is not None: |
| 251 | if probs.device.type != 'npu': | 270 | if probs.device.type != 'npu': |
| 252 | raise AssertionError('This operation needs NPU') | 271 | raise AssertionError('This operation needs NPU') |
| 253 | 272 | ||
| @@ -345,4 +364,4 @@ def moe_sort_chunks_by_index_with_probs( | |||
| 345 | Chunk indices used to permute the chunks. | 364 | Chunk indices used to permute the chunks. |
| 346 | """ | 365 | """ |
| 347 | output, permuted_probs = _moe_chunk_sort.apply(inp, split_sizes, sorted_index, probs) | 366 | output, permuted_probs = _moe_chunk_sort.apply(inp, split_sizes, sorted_index, probs) |
| 348 | - return output, permuted_probs | 367 | + return output, permuted_probs |
| @@ -3,241 +3,10 @@ | |||
| 3 | 3 | ||
| 4 | import pytest | 4 | import pytest |
| 5 | import torch | 5 | import torch |
| 6 | -import triton | 6 | +from mindspeed.lite.ops.triton.sort_chunks_by_idx import ( |
| 7 | -import triton.language as tl | 7 | + moe_sort_chunks_by_index_with_probs, |
| 8 | -from mindspeed.lite.ops.triton.sort_chunks_by_idx import make_chunk_sort_map, sort_chunks_by_map | 8 | +) |
| 9 | - | 9 | +from megatron.core.transformer.moe.moe_utils import sort_chunks_by_idxs |
| 10 | - | ||
| 11 | -def accuracy_comparison(y_cal, y_ref): | ||
| 12 | - """ | ||
| 13 | - 精度比对函数：根据数据类型选择合适的比对策略。 | ||
| 14 | - | ||
| 15 | - 不同数据类型的处理策略： | ||
| 16 | - - 浮点类型（float16/32, bfloat16）：使用 torch.testing.assert_close，设置相对/绝对误差容限 | ||
| 17 | - - 整数类型（int8/16/32/64）：要求完全相等（torch.equal） | ||
| 18 | - - 布尔类型（bool）：CPU 上严格比较（避免设备差异） | ||
| 19 | - """ | ||
| 20 | - # 检查输出数据类型是否一致 | ||
| 21 | - assert y_cal.dtype == y_ref.dtype, f"dtype mismatch: {y_cal.dtype} vs {y_ref.dtype}" | ||
| 22 | - tensor_dtype = y_cal.dtype | ||
| 23 | - | ||
| 24 | - # 将张量移动到 NPU（假设测试在 NPU 上进行） | ||
| 25 | - y_cal = y_cal.npu() | ||
| 26 | - y_ref = y_ref.npu() | ||
| 27 | - | ||
| 28 | - # 根据数据类型选择不同的比对方式 | ||
| 29 | - if tensor_dtype == torch.float16: | ||
| 30 | - # float16 精度较低，允许稍大误差 | ||
| 31 | - torch.testing.assert_close(y_ref, y_cal, rtol=1e-3, atol=1e-3, equal_nan=True) | ||
| 32 | - elif tensor_dtype == torch.bfloat16: | ||
| 33 | - # bfloat16 精度更低，建议转为 float32 再比较 | ||
| 34 | - torch.testing.assert_close( | ||
| 35 | - y_ref.to(torch.float32), | ||
| 36 | - y_cal.to(torch.float32), | ||
| 37 | - rtol=1e-3, | ||
| 38 | - atol=1e-3, | ||
| 39 | - equal_nan=True | ||
| 40 | - ) | ||
| 41 | - elif tensor_dtype == torch.float32: | ||
| 42 | - # float32 精度较高，使用更严格的容差 | ||
| 43 | - torch.testing.assert_close(y_ref, y_cal, rtol=1e-4, atol=1e-4, equal_nan=True) | ||
| 44 | - elif tensor_dtype in [torch.int64, torch.int32, torch.int16, torch.int8]: | ||
| 45 | - # 整数类型应完全相等 | ||
| 46 | - assert torch.equal(y_cal, y_ref), f"Integer tensors are not equal for dtype {tensor_dtype}" | ||
| 47 | - elif tensor_dtype == torch.bool: | ||
| 48 | - # 布尔类型建议在 CPU 上比较，避免设备间布尔表示差异 | ||
| 49 | - assert torch.equal(y_cal.cpu(), y_ref.cpu()), "Boolean tensors are not equal" | ||
| 50 | - else: | ||
| 51 | - raise ValueError(f'Invalid or unsupported tensor dtype: {tensor_dtype}') | ||
| 52 | - | ||
| 53 | - | ||
| 54 | - | ||
| 55 | -def _make_chunk_sort_map_kernel_gpu( | ||
| 56 | - # pointers | ||
| 57 | - split_sizes_ptr, | ||
| 58 | - sorted_indices_ptr, | ||
| 59 | - dst_rows_ptr, | ||
| 60 | - # sizes | ||
| 61 | - num_splits: tl.constexpr, | ||
| 62 | - # metas | ||
| 63 | - IDX_LOAD_WIDTH: tl.constexpr, | ||
| 64 | -): | ||
| 65 | - pid = tl.program_id(0) | ||
| 66 | - | ||
| 67 | - load_split_offset = tl.arange(0, IDX_LOAD_WIDTH) | ||
| 68 | - sorted_indices = tl.load( | ||
| 69 | - sorted_indices_ptr + load_split_offset, mask=load_split_offset < num_splits | ||
| 70 | - ) | ||
| 71 | - | ||
| 72 | - # get chunk idx of the current token in the input tensor | ||
| 73 | - input_split_sizes = tl.load( | ||
| 74 | - split_sizes_ptr + load_split_offset, mask=load_split_offset < num_splits, other=0 | ||
| 75 | - ).to(tl.int32) | ||
| 76 | - input_split_sizes_cumsum = tl.cumsum(input_split_sizes) | ||
| 77 | - input_split_sizes_mask = tl.where(input_split_sizes_cumsum <= pid, 1, 0) | ||
| 78 | - input_chunk_idx = tl.sum(input_split_sizes_mask) | ||
| 79 | - input_split_sizes_presum = tl.sum(input_split_sizes * input_split_sizes_mask) | ||
| 80 | - in_chunk_offset = pid - input_split_sizes_presum | ||
| 81 | - | ||
| 82 | - # get chunk idx of the current token in the output tensor | ||
| 83 | - output_chunk_mask = tl.where(sorted_indices == input_chunk_idx, 1, 0) | ||
| 84 | - output_chunk_idx = tl.argmax(output_chunk_mask, axis=-1) | ||
| 85 | - | ||
| 86 | - # make row_id_map | ||
| 87 | - output_split_sizes = tl.load( | ||
| 88 | - split_sizes_ptr + sorted_indices, mask=load_split_offset < num_splits | ||
| 89 | - ).to(tl.int32) | ||
| 90 | - output_pre_split_sizes = tl.where(load_split_offset < output_chunk_idx, output_split_sizes, 0) | ||
| 91 | - dst_row = tl.sum(output_pre_split_sizes) + in_chunk_offset | ||
| 92 | - tl.store(dst_rows_ptr + pid, dst_row) | ||
| 93 | - | ||
| 94 | - | ||
| 95 | - | ||
| 96 | -def _sort_chunks_by_map_kernel_gpu( | ||
| 97 | - # pointers | ||
| 98 | - input_ptr, | ||
| 99 | - output_ptr, | ||
| 100 | - row_id_map_ptr, | ||
| 101 | - probs_ptr, | ||
| 102 | - permuted_probs_ptr, | ||
| 103 | - # sizes | ||
| 104 | - hidden_size: tl.constexpr, | ||
| 105 | - # strides | ||
| 106 | - stride_input_token, | ||
| 107 | - stride_input_hidden, | ||
| 108 | - stride_output_token, | ||
| 109 | - stride_output_hidden, | ||
| 110 | - stride_probs_token, | ||
| 111 | - stride_permuted_probs_token, | ||
| 112 | - # metas | ||
| 113 | - PERMUTE_PROBS: tl.constexpr, | ||
| 114 | - BLOCK_SIZE: tl.constexpr, | ||
| 115 | - FORWARD: tl.constexpr, | ||
| 116 | -): | ||
| 117 | - pid_t = tl.program_id(0) | ||
| 118 | - pid_h = tl.program_id(1) | ||
| 119 | - if FORWARD: | ||
| 120 | - src_row = pid_t.to(tl.int64) | ||
| 121 | - dst_row = tl.load(row_id_map_ptr + pid_t).to(tl.int64) | ||
| 122 | - else: | ||
| 123 | - src_row = tl.load(row_id_map_ptr + pid_t).to(tl.int64) | ||
| 124 | - dst_row = pid_t.to(tl.int64) | ||
| 125 | - current_offset = pid_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) | ||
| 126 | - mask = current_offset < hidden_size | ||
| 127 | - input_offsets = src_row * stride_input_token + current_offset * stride_input_hidden | ||
| 128 | - output_offsets = dst_row * stride_output_token + current_offset * stride_output_hidden | ||
| 129 | - inp = tl.load(input_ptr + input_offsets, mask=mask) | ||
| 130 | - tl.store(output_ptr + output_offsets, inp, mask=mask) | ||
| 131 | - if PERMUTE_PROBS: | ||
| 132 | - if pid_h == 0: | ||
| 133 | - prob_off = src_row * stride_probs_token | ||
| 134 | - prob = tl.load(probs_ptr + prob_off) | ||
| 135 | - permuted_prob_off = dst_row * stride_permuted_probs_token | ||
| 136 | - tl.store(permuted_probs_ptr + permuted_prob_off, prob) | ||
| 137 | - | ||
| 138 | - | ||
| 139 | -try: | ||
| 140 | - _sort_chunks_by_map_kernel_gpu = triton.autotune( | ||
| 141 | - configs=[ | ||
| 142 | - triton.Config({"BLOCK_SIZE": 64}), | ||
| 143 | - triton.Config({"BLOCK_SIZE": 128}), | ||
| 144 | - triton.Config({"BLOCK_SIZE": 256}), | ||
| 145 | - triton.Config({"BLOCK_SIZE": 512}), | ||
| 146 | - triton.Config({"BLOCK_SIZE": 1024}), | ||
| 147 | - triton.Config({"BLOCK_SIZE": 2048}), | ||
| 148 | - triton.Config({"BLOCK_SIZE": 4096}), | ||
| 149 | - ], | ||
| 150 | - key=["hidden_size"], | ||
| 151 | - )(_sort_chunks_by_map_kernel_gpu) | ||
| 152 | -except RuntimeError: | ||
| 153 | - pass | ||
| 154 | - | ||
| 155 | - | ||
| 156 | -def make_chunk_sort_map_gpu( | ||
| 157 | - split_sizes: torch.Tensor, | ||
| 158 | - sorted_indices: torch.Tensor, | ||
| 159 | - num_tokens: int, | ||
| 160 | - num_splits: int, | ||
| 161 | -): | ||
| 162 | - """ | ||
| 163 | - Make a row_id_map for chunk sort. | ||
| 164 | - | ||
| 165 | - Parameters | ||
| 166 | - ---------- | ||
| 167 | - split_sizes: torch.Tensor | ||
| 168 | - The sizes of the chunks of shape `[num_splits,]`. | ||
| 169 | - sorted_indices: torch.Tensor | ||
| 170 | - The indices of the sorted chunks of shape `[num_splits,]`. | ||
| 171 | - num_tokens: int | ||
| 172 | - Number of tokens in the input tensor. | ||
| 173 | - num_splits: int | ||
| 174 | - Number of splits of split_sizes and sorted_indices. | ||
| 175 | - """ | ||
| 176 | - row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device='npu') | ||
| 177 | - grid = (num_tokens,) | ||
| 178 | - _make_chunk_sort_map_kernel_gpu[grid]( | ||
| 179 | - split_sizes, | ||
| 180 | - sorted_indices, | ||
| 181 | - row_id_map, | ||
| 182 | - num_splits, | ||
| 183 | - IDX_LOAD_WIDTH=triton.next_power_of_2(num_splits), | ||
| 184 | - ) | ||
| 185 | - return row_id_map | ||
| 186 | - | ||
| 187 | - | ||
| 188 | -def sort_chunks_by_map_gpu( | ||
| 189 | - inp: torch.Tensor, | ||
| 190 | - row_id_map: torch.Tensor, | ||
| 191 | - probs: torch.Tensor, | ||
| 192 | - num_tokens: int, | ||
| 193 | - hidden_size: int, | ||
| 194 | - is_forward: bool, | ||
| 195 | -): | ||
| 196 | - """ | ||
| 197 | - Sort chunks with row_id_map. | ||
| 198 | - | ||
| 199 | - Parameters | ||
| 200 | - ---------- | ||
| 201 | - inp: torch.Tensor | ||
| 202 | - Input tensor of shape `[num_tokens, hidden_size]`. | ||
| 203 | - row_id_map: torch.Tensor | ||
| 204 | - The token to expert mapping tensor of shape `[num_tokens,]`. | ||
| 205 | - probs: torch.Tensor | ||
| 206 | - The probabilities of the input tensor. If it is not None, it will be permuted. | ||
| 207 | - num_tokens: int | ||
| 208 | - Number of tokens in the input tensor. | ||
| 209 | - hidden_size: int | ||
| 210 | - Hidden size of the input tensor. | ||
| 211 | - is_forward: bool | ||
| 212 | - Whether the sort is for forward or backward. | ||
| 213 | - """ | ||
| 214 | - output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device='npu') | ||
| 215 | - if probs is not None: | ||
| 216 | - permuted_probs = torch.empty((num_tokens,), dtype=probs.dtype, device='npu') | ||
| 217 | - else: | ||
| 218 | - permuted_probs = None | ||
| 219 | - # pylint: disable=unnecessary-lambda-assignment | ||
| 220 | - | ||
| 221 | - def get_grid(META): | ||
| 222 | - return (num_tokens, triton.cdiv(hidden_size, META["BLOCK_SIZE"])) | ||
| 223 | - | ||
| 224 | - _sort_chunks_by_map_kernel_gpu[get_grid]( | ||
| 225 | - inp, | ||
| 226 | - output, | ||
| 227 | - row_id_map, | ||
| 228 | - probs, | ||
| 229 | - permuted_probs, | ||
| 230 | - hidden_size, | ||
| 231 | - inp.stride(0), | ||
| 232 | - inp.stride(1), | ||
| 233 | - output.stride(0), | ||
| 234 | - output.stride(1), | ||
| 235 | - probs.stride(0) if probs is not None else None, | ||
| 236 | - permuted_probs.stride(0) if permuted_probs is not None else None, | ||
| 237 | - PERMUTE_PROBS=probs is not None, | ||
| 238 | - FORWARD=is_forward, | ||
| 239 | - ) | ||
| 240 | - return output, permuted_probs | ||
| 241 | 10 | ||
| 242 | 11 | ||
| 243 | def gen_split_sizes(num_tokens, num_splits): | 12 | def gen_split_sizes(num_tokens, num_splits): |
| @@ -247,39 +16,32 @@ def gen_split_sizes(num_tokens, num_splits): | |||
| 247 | scaled_numbers[-1] += num_tokens - torch.sum(scaled_numbers) | 16 | scaled_numbers[-1] += num_tokens - torch.sum(scaled_numbers) |
| 248 | return scaled_numbers | 17 | return scaled_numbers |
| 249 | 18 | ||
| 250 | -TEST_CASES = [ | 19 | + |
| 251 | - (16, 2048, 256), | 20 | +TEST_CASES = [(16, 2048, 256), (32, 4096, 128), (1024, 600000, 7168)] |
| 252 | - (32, 4096, 128) | ||
| 253 | -] | ||
| 254 | 21 | ||
| 255 | 22 | ||
| 256 | 23 | ||
| 257 | "num_splits,num_tokens,hidden_size", | 24 | "num_splits,num_tokens,hidden_size", |
| 258 | - [pytest.param(*case, id=f"split{case[0]}-tokens{case[1]}-hid{case[2]}") for case in TEST_CASES] | 25 | + [pytest.param(*case, id=f"split{case[0]}-tokens{case[1]}-hid{case[2]}") for case in TEST_CASES], |
| 259 | ) | 26 | ) |
| 260 | - | ||
| 261 | def test_sort_chunks_by_idx(num_splits, num_tokens, hidden_size): | 27 | def test_sort_chunks_by_idx(num_splits, num_tokens, hidden_size): |
| 262 | split_sizes = gen_split_sizes(num_tokens, num_splits) | 28 | split_sizes = gen_split_sizes(num_tokens, num_splits) |
| 263 | sorted_indices = torch.randperm(num_splits, device='npu') | 29 | sorted_indices = torch.randperm(num_splits, device='npu') |
| 264 | 30 | ||
| 265 | - ref_row_id_map = make_chunk_sort_map_gpu(split_sizes, sorted_indices, num_tokens, num_splits) | ||
| 266 | - row_id_map = make_chunk_sort_map(split_sizes, sorted_indices, num_tokens, num_splits) | ||
| 267 | - | ||
| 268 | - accuracy_comparison(ref_row_id_map, row_id_map) | ||
| 269 | - | ||
| 270 | inp = torch.randn(num_tokens, hidden_size, dtype=torch.float32, device='npu') | 31 | inp = torch.randn(num_tokens, hidden_size, dtype=torch.float32, device='npu') |
| 271 | probs = torch.rand(num_tokens, dtype=torch.float32, device='npu') | 32 | probs = torch.rand(num_tokens, dtype=torch.float32, device='npu') |
| 272 | - | ||
| 273 | - # test forward | ||
| 274 | - ref_output, ref_permuted_probs = sort_chunks_by_map_gpu(inp, ref_row_id_map, probs, num_tokens, hidden_size, True) | ||
| 275 | - output, permuted_probs = sort_chunks_by_map(inp, row_id_map, probs, num_tokens, hidden_size, True) | ||
| 276 | 33 | ||
| 277 | - accuracy_comparison(ref_output, output) | 34 | + output, permuted_probs = moe_sort_chunks_by_index_with_probs(inp, probs, split_sizes, sorted_indices) |
| 278 | - accuracy_comparison(ref_permuted_probs, permuted_probs) | ||
| 279 | - | ||
| 280 | - # test backward | ||
| 281 | - ref_output, ref_permuted_probs = sort_chunks_by_map_gpu(inp, ref_row_id_map, probs, num_tokens, hidden_size, False) | ||
| 282 | - output, permuted_probs = sort_chunks_by_map(inp, row_id_map, probs, num_tokens, hidden_size, False) | ||
| 283 | 35 | ||
| 284 | - accuracy_comparison(ref_output, output) | 36 | + ref_output, ref_permuted_probs = sort_chunks_by_idxs(inp, split_sizes, sorted_indices, probs, fused=False) |
| 285 | - accuracy_comparison(ref_permuted_probs, permuted_probs) | 37 | + |
| 38 | + torch.testing.assert_close(output, ref_output, rtol=1e-4, atol=1e-4) | ||
| 39 | + torch.testing.assert_close(permuted_probs, ref_permuted_probs, rtol=1e-4, atol=1e-4) | ||
| 40 | + | ||
| 41 | + output.backward(torch.ones_like(output)) | ||
| 42 | + npu_grad = inp.grad.clone() | ||
| 43 | + | ||
| 44 | + ref_output.backward(torch.ones_like(ref_output)) | ||
| 45 | + ref_grad = inp.grad.clone() | ||
| 46 | + | ||
| 47 | + torch.testing.assert_close(npu_grad, ref_grad, rtol=1e-4, atol=1e-4) | ||