| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the MOE permute/unpermute kernel. |
| 4 | |
| 5 | Run `pytest tests/kernels/test_moe_permute_unpermute.py`. |
| 6 | """ |
| 7 | |
| 8 | import numpy as np |
| 9 | import pytest |
| 10 | import torch |
| 11 | |
| 12 | from vllm.model_executor.layers.fused_moe import fused_topk |
| 13 | from vllm.model_executor.layers.fused_moe.expert_map_manager import ( |
| 14 | determine_expert_map, |
| 15 | ) |
| 16 | from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( |
| 17 | MoEPermuteScratch, |
| 18 | get_moe_permute_scratch, |
| 19 | moe_permute, |
| 20 | moe_permute_unpermute_supported, |
| 21 | moe_prepare_scatter, |
| 22 | moe_unpermute, |
| 23 | ) |
| 24 | from vllm.platforms import current_platform |
| 25 | from vllm.utils.torch_utils import set_random_seed |
| 26 | from vllm.v1.worker.workspace import current_workspace_manager |
| 27 | |
| 28 | NUM_EXPERTS = [16, 64, 256] |
| 29 | TOP_KS = [2, 6, 8] |
| 30 | EP_SIZE = [1, 4, 16] |
| 31 | set_random_seed(0) |
| 32 | |
| 33 | if current_platform.is_rocm(): |
| 34 | pytest.skip( |
| 35 | "moe_permute_unpermute_supported is not defined for ROCm", |
| 36 | allow_module_level=True, |
| 37 | ) |
| 38 | |
| 39 | |
| 40 | def torch_permute( |
| 41 | hidden_states: torch.Tensor, |
| 42 | topk_ids: torch.Tensor, |
| 43 | # token_expert_indices: torch.Tensor, |
| 44 | topk: int, |
| 45 | n_expert: int, |
| 46 | n_local_expert: int, |
| 47 | start_expert: int, |
| 48 | expert_map: torch.Tensor | None = None, |
| 49 | ) -> list[torch.Tensor]: |
| 50 | n_token = hidden_states.shape[0] |
| 51 | if expert_map is not None: |
| 52 | is_local_expert = expert_map[topk_ids] != -1 |
| 53 | not_local_expert = expert_map[topk_ids] == -1 |
| 54 | topk_ids = is_local_expert * (topk_ids - start_expert) + not_local_expert * ( |
| 55 | topk_ids + n_expert |
| 56 | ) |
| 57 | token_expert_indices = torch.arange( |
| 58 | 0, n_token * topk, dtype=torch.int32, device=hidden_states.device |
| 59 | ).reshape((n_token, topk)) |
| 60 | |
| 61 | sorted_topk_ids, sorted_indices = torch.sort(topk_ids.flatten(), stable=True) |
| 62 | dst_row_id2src_row_id_map = token_expert_indices.flatten()[sorted_indices] |
| 63 | |
| 64 | expert_first_token_offset = torch.zeros( |
| 65 | n_local_expert + 1, dtype=torch.int64, device="cuda" |
| 66 | ) |
| 67 | idx = 0 |
| 68 | for i in range(0, n_local_expert): |
| 69 | cnt = 0 |
| 70 | while idx < sorted_topk_ids.numel() and sorted_topk_ids[idx] == i: |
| 71 | cnt += 1 |
| 72 | idx += 1 |
| 73 | expert_first_token_offset[i + 1] = expert_first_token_offset[i] + cnt |
| 74 | |
| 75 | _, src2dst_idx = torch.sort(dst_row_id2src_row_id_map) |
| 76 | valid_row_idx = [] |
| 77 | permuted_hidden_states = hidden_states[dst_row_id2src_row_id_map // topk, ...] |
| 78 | src_row_id2dst_row_id_map = torch.arange( |
| 79 | 0, n_token * topk, device="cuda", dtype=torch.int32 |
| 80 | )[src2dst_idx].reshape((n_token, topk)) |
| 81 | valid_row_idx += [i for i in range(expert_first_token_offset[-1])] |
| 82 | dst_row_id2src_row_id_map[expert_first_token_offset[-1] :] = n_token * topk |
| 83 | return [ |
| 84 | permuted_hidden_states, |
| 85 | expert_first_token_offset, |
| 86 | src_row_id2dst_row_id_map, |
| 87 | dst_row_id2src_row_id_map, |
| 88 | valid_row_idx, |
| 89 | ] |
| 90 | |
| 91 | |
| 92 | def torch_unpermute( |
| 93 | permuted_hidden_states: torch.Tensor, |
| 94 | topk_weights: torch.Tensor, |
| 95 | topk_ids: torch.Tensor, |
| 96 | token_expert_indices: torch.Tensor, |
| 97 | src_row_id2dst_row_id_map: torch.Tensor, |
| 98 | valid_row_idx: torch.Tensor, |
| 99 | topk: int, |
| 100 | n_expert: int, |
| 101 | ) -> torch.Tensor: |
| 102 | # ignore invalid row |
| 103 | n_hidden = permuted_hidden_states.shape[1] |
| 104 | mask = torch.zeros(permuted_hidden_states.shape[0], dtype=bool, device="cuda") |
| 105 | mask[valid_row_idx] = True |
| 106 | permuted_hidden_states[~mask] = 0 |
| 107 | |
| 108 | permuted_hidden_states = permuted_hidden_states[ |
| 109 | src_row_id2dst_row_id_map.flatten(), ... |
| 110 | ] |
| 111 | permuted_hidden_states = permuted_hidden_states.view(-1, topk, n_hidden) |
| 112 | output = ( |
| 113 | (permuted_hidden_states * topk_weights.unsqueeze(2)) |
| 114 | .sum(1) |
| 115 | .to(permuted_hidden_states.dtype) |
| 116 | ) |
| 117 | return output |
| 118 | |
| 119 | |
| 120 | @pytest.mark.parametrize("n_token", [1, 33, 1024, 5000]) |
| 121 | @pytest.mark.parametrize("n_hidden", [2048, 7168]) |
| 122 | @pytest.mark.parametrize("n_expert", NUM_EXPERTS) |
| 123 | @pytest.mark.parametrize("topk", TOP_KS) |
| 124 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 125 | @pytest.mark.parametrize("ep_size", EP_SIZE) |
| 126 | @pytest.mark.parametrize("use_scratch", [False, True]) |
| 127 | def test_moe_permute_unpermute( |
| 128 | n_token: int, |
| 129 | n_hidden: int, |
| 130 | topk: int, |
| 131 | n_expert: int, |
| 132 | ep_size: int, |
| 133 | dtype: torch.dtype, |
| 134 | use_scratch: bool, |
| 135 | workspace_init, |
| 136 | ): |
| 137 | if not moe_permute_unpermute_supported(): |
| 138 | pytest.skip("moe_permute_unpermute is not supported on this platform.") |
| 139 | ep_rank = np.random.randint(0, ep_size) |
| 140 | expert_map = None |
| 141 | n_local_expert = n_expert |
| 142 | if ep_size != 1: |
| 143 | n_local_expert, expert_map, _ = determine_expert_map(ep_size, ep_rank, n_expert) |
| 144 | expert_map = expert_map.cuda() |
| 145 | start_expert = n_local_expert * ep_rank |
| 146 | set_random_seed(0) |
| 147 | hidden_states = torch.randn((n_token, n_hidden), device="cuda").to(dtype) |
| 148 | gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype) |
| 149 | topk_weights, topk_ids, token_expert_indices = fused_topk( |
| 150 | hidden_states, gating_output, topk, False |
| 151 | ) |
| 152 | ( |
| 153 | gold_permuted_hidden_states, |
| 154 | gold_expert_first_token_offset, |
| 155 | gold_inv_permuted_idx, |
| 156 | gold_permuted_idx, |
| 157 | valid_row_idx, |
| 158 | ) = torch_permute( |
| 159 | hidden_states, |
| 160 | topk_ids, |
| 161 | # token_expert_indices, |
| 162 | topk, |
| 163 | n_expert, |
| 164 | n_local_expert, |
| 165 | start_expert, |
| 166 | expert_map=expert_map, |
| 167 | ) |
| 168 | |
| 169 | scratch = None |
| 170 | if use_scratch: |
| 171 | scratch = get_moe_permute_scratch( |
| 172 | max_num_tokens=n_token, |
| 173 | topk=topk, |
| 174 | num_experts=n_expert, |
| 175 | num_local_experts=n_local_expert, |
| 176 | device=hidden_states.device, |
| 177 | hidden_size=n_hidden, |
| 178 | hidden_dtype=dtype, |
| 179 | ) |
| 180 | |
| 181 | ( |
| 182 | permuted_hidden_states, |
| 183 | _, |
| 184 | expert_first_token_offset, |
| 185 | inv_permuted_idx, |
| 186 | _, |
| 187 | ) = moe_permute( |
| 188 | hidden_states=hidden_states, |
| 189 | a1q_scale=None, |
| 190 | topk_ids=topk_ids, |
| 191 | n_expert=n_expert, |
| 192 | n_local_expert=n_local_expert, |
| 193 | expert_map=expert_map, |
| 194 | scratch=scratch, |
| 195 | ) |
| 196 | |
| 197 | # check expert_first_token_offset |
| 198 | torch.testing.assert_close( |
| 199 | gold_expert_first_token_offset, expert_first_token_offset, atol=0, rtol=0 |
| 200 | ) |
| 201 | # check src_row_id2dst_row_id_map |
| 202 | torch.testing.assert_close( |
| 203 | gold_inv_permuted_idx.flatten(), inv_permuted_idx, atol=0, rtol=0 |
| 204 | ) |
| 205 | |
| 206 | # check permuted_hidden_states, only valid token |
| 207 | torch.testing.assert_close( |
| 208 | gold_permuted_hidden_states[valid_row_idx], |
| 209 | permuted_hidden_states[valid_row_idx], |
| 210 | atol=0, |
| 211 | rtol=0, |
| 212 | ) |
| 213 | # add a random tensor to simulate group gemm |
| 214 | result0 = 0.5 * permuted_hidden_states + torch.randn_like(permuted_hidden_states) |
| 215 | result4 = torch.empty_like(hidden_states) |
| 216 | moe_unpermute( |
| 217 | result4, result0, topk_weights, inv_permuted_idx, expert_first_token_offset |
| 218 | ) |
| 219 | |
| 220 | gold4 = torch_unpermute( |
| 221 | result0, |
| 222 | topk_weights, |
| 223 | topk_ids, |
| 224 | token_expert_indices, |
| 225 | inv_permuted_idx, |
| 226 | valid_row_idx, |
| 227 | topk, |
| 228 | n_local_expert, |
| 229 | ) |
| 230 | # check unpermuted hidden |
| 231 | torch.testing.assert_close(result4, gold4, atol=2e-2, rtol=0) |
| 232 | |
| 233 | |
| 234 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 235 | @pytest.mark.parametrize("n_token", [1, 33, 128]) |
| 236 | @pytest.mark.parametrize("topk", [1, 6, 8]) |
| 237 | def test_moe_permute_reuses_scratch_buffers( |
| 238 | dtype: torch.dtype, n_token: int, topk: int, workspace_init |
| 239 | ): |
| 240 | if not moe_permute_unpermute_supported(): |
| 241 | pytest.skip("moe_permute_unpermute is not supported on this platform.") |
| 242 | |
| 243 | n_hidden = 2048 |
| 244 | n_expert = 16 |
| 245 | hidden_states = torch.randn((n_token, n_hidden), device="cuda").to(dtype) |
| 246 | gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype) |
| 247 | _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False) |
| 248 | |
| 249 | scratch_config = dict( |
| 250 | max_num_tokens=n_token, |
| 251 | topk=topk, |
| 252 | num_experts=n_expert, |
| 253 | num_local_experts=n_expert, |
| 254 | device=hidden_states.device, |
| 255 | hidden_size=n_hidden, |
| 256 | hidden_dtype=hidden_states.dtype, |
| 257 | ) |
| 258 | scratch = get_moe_permute_scratch(**scratch_config) |
| 259 | |
| 260 | first = moe_permute( |
| 261 | hidden_states=hidden_states, |
| 262 | a1q_scale=None, |
| 263 | topk_ids=topk_ids, |
| 264 | n_expert=n_expert, |
| 265 | scratch=scratch, |
| 266 | ) |
| 267 | current_workspace_manager().lock() |
| 268 | assert get_moe_permute_scratch(**scratch_config) is scratch |
| 269 | graph = torch.cuda.CUDAGraph() |
| 270 | with torch.cuda.graph(graph): |
| 271 | second = moe_permute( |
| 272 | hidden_states=hidden_states, |
| 273 | a1q_scale=None, |
| 274 | topk_ids=topk_ids, |
| 275 | n_expert=n_expert, |
| 276 | scratch=get_moe_permute_scratch(**scratch_config), |
| 277 | ) |
| 278 | |
| 279 | for _ in range(2): |
| 280 | hidden_states.add_(1) |
| 281 | topk_ids.copy_(topk_ids.roll(1, 0)) |
| 282 | graph.replay() |
| 283 | expected = moe_permute(hidden_states, None, topk_ids, n_expert) |
| 284 | for actual, reference in zip(second, expected): |
| 285 | if actual is not None: |
| 286 | torch.testing.assert_close(actual, reference) |
| 287 | |
| 288 | ( |
| 289 | permuted_hidden_states_1, |
| 290 | _, |
| 291 | expert_first_token_offset_1, |
| 292 | inv_permuted_idx_1, |
| 293 | permuted_idx_1, |
| 294 | ) = first |
| 295 | ( |
| 296 | permuted_hidden_states_2, |
| 297 | _, |
| 298 | expert_first_token_offset_2, |
| 299 | inv_permuted_idx_2, |
| 300 | permuted_idx_2, |
| 301 | ) = second |
| 302 | |
| 303 | torch.testing.assert_close(permuted_hidden_states_1, permuted_hidden_states_2) |
| 304 | torch.testing.assert_close(expert_first_token_offset_1, expert_first_token_offset_2) |
| 305 | torch.testing.assert_close(inv_permuted_idx_1, inv_permuted_idx_2) |
| 306 | torch.testing.assert_close(permuted_idx_1, permuted_idx_2) |
| 307 | |
| 308 | assert ( |
| 309 | permuted_hidden_states_1.untyped_storage().data_ptr() |
| 310 | == permuted_hidden_states_2.untyped_storage().data_ptr() |
| 311 | ) |
| 312 | assert ( |
| 313 | expert_first_token_offset_1.untyped_storage().data_ptr() |
| 314 | == expert_first_token_offset_2.untyped_storage().data_ptr() |
| 315 | ) |
| 316 | assert ( |
| 317 | inv_permuted_idx_1.untyped_storage().data_ptr() |
| 318 | == scratch.inv_permuted_idx.untyped_storage().data_ptr() |
| 319 | ) |
| 320 | assert ( |
| 321 | permuted_idx_1.untyped_storage().data_ptr() |
| 322 | == scratch.permuted_idx.untyped_storage().data_ptr() |
| 323 | ) |
| 324 | |
| 325 | |
| 326 | def test_moe_permute_scratch_reused_across_graph_sizes(workspace_init) -> None: |
| 327 | """Switching graph sizes must not expose metadata left by a larger batch.""" |
| 328 | if not moe_permute_unpermute_supported(): |
| 329 | pytest.skip("moe_permute_unpermute is not supported on this platform.") |
| 330 | |
| 331 | device = torch.device("cuda:0") |
| 332 | scratch = get_moe_permute_scratch( |
| 333 | max_num_tokens=128, |
| 334 | topk=6, |
| 335 | num_experts=16, |
| 336 | num_local_experts=16, |
| 337 | device=device, |
| 338 | hidden_size=2048, |
| 339 | hidden_dtype=torch.bfloat16, |
| 340 | ) |
| 341 | current_workspace_manager().lock() |
| 342 | runs = [] |
| 343 | for n_token in (128, 1, 33): |
| 344 | hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device) |
| 345 | topk_ids = torch.randint(16, (n_token, 6), device=device) |
| 346 | moe_permute(hidden, None, topk_ids, 16, scratch=scratch) |
| 347 | graph = torch.cuda.CUDAGraph() |
| 348 | with torch.cuda.graph(graph): |
| 349 | output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch) |
| 350 | runs.append((graph, hidden, topk_ids, output)) |
| 351 | |
| 352 | for graph, hidden, topk_ids, output in runs + runs[::-1]: |
| 353 | hidden.add_(1) |
| 354 | topk_ids.copy_((topk_ids + 1) % 16) |
| 355 | graph.replay() |
| 356 | expected = moe_permute(hidden, None, topk_ids, 16) |
| 357 | for actual, reference in zip(output, expected): |
| 358 | if actual is not None: |
| 359 | torch.testing.assert_close(actual, reference) |
| 360 | |
| 361 | |
| 362 | def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> None: |
| 363 | """Concurrent graph replays in different ubatches/lanes cannot share scratch.""" |
| 364 | if not moe_permute_unpermute_supported(): |
| 365 | pytest.skip("moe_permute_unpermute is not supported on this platform.") |
| 366 | |
| 367 | from vllm.v1.worker import workspace |
| 368 | |
| 369 | device = torch.device("cuda:0") |
| 370 | manager = workspace.WorkspaceManager(device, num_ubatches=2, num_lanes=2) |
| 371 | monkeypatch.setattr(workspace, "_manager", manager) |
| 372 | ubatch = 0 |
| 373 | monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: ubatch) |
| 374 | scratch_config = dict( |
| 375 | max_num_tokens=64, |
| 376 | topk=6, |
| 377 | num_experts=16, |
| 378 | num_local_experts=16, |
| 379 | device=device, |
| 380 | hidden_size=2048, |
| 381 | hidden_dtype=torch.bfloat16, |
| 382 | ) |
| 383 | runs = [] |
| 384 | for i, n_token in enumerate((1, 7, 33, 64)): |
| 385 | ubatch = i // 2 |
| 386 | stream = torch.cuda.Stream() |
| 387 | stream.wait_stream(torch.cuda.current_stream()) |
| 388 | with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream): |
| 389 | hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device) |
| 390 | topk_ids = torch.randint(16, (n_token, 6), device=device) |
| 391 | scratch = get_moe_permute_scratch(**scratch_config) |
| 392 | moe_permute(hidden, None, topk_ids, 16, scratch=scratch) |
| 393 | graph = torch.cuda.CUDAGraph() |
| 394 | with torch.cuda.graph(graph, stream=stream): |
| 395 | output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch) |
| 396 | runs.append((stream, graph, hidden, topk_ids, output, scratch)) |
| 397 | torch.cuda.current_stream().wait_stream(stream) |
| 398 | |
| 399 | assert len({run[-1].permuted_hidden_states.data_ptr() for run in runs}) == 4 |
| 400 | manager.lock() |
| 401 | for _ in range(3): |
| 402 | for i, (stream, graph, hidden, topk_ids, _, scratch) in enumerate(runs): |
| 403 | ubatch = i // 2 |
| 404 | with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream): |
| 405 | assert get_moe_permute_scratch(**scratch_config) is scratch |
| 406 | hidden.add_(1) |
| 407 | topk_ids.copy_((topk_ids + 1) % 16) |
| 408 | graph.replay() |
| 409 | for stream, _, hidden, topk_ids, output, _ in runs: |
| 410 | torch.cuda.current_stream().wait_stream(stream) |
| 411 | expected = moe_permute(hidden, None, topk_ids, 16) |
| 412 | for actual, reference in zip(output, expected): |
| 413 | if actual is not None: |
| 414 | torch.testing.assert_close(actual, reference) |
| 415 | |
| 416 | |
| 417 | def test_moe_permute_scratch_without_manager(monkeypatch) -> None: |
| 418 | """Standalone calls get independent scratch with initialized row indices.""" |
| 419 | if not moe_permute_unpermute_supported(): |
| 420 | pytest.skip("moe_permute_unpermute is not supported on this platform.") |
| 421 | |
| 422 | from vllm.v1.worker import workspace |
| 423 | |
| 424 | monkeypatch.setattr(workspace, "_manager", None) |
| 425 | config = dict( |
| 426 | max_num_tokens=4, |
| 427 | topk=2, |
| 428 | num_experts=4, |
| 429 | num_local_experts=4, |
| 430 | device=torch.device("cuda"), |
| 431 | ) |
| 432 | first = get_moe_permute_scratch(**config) |
| 433 | second = get_moe_permute_scratch(**config) |
| 434 | |
| 435 | assert ( |
| 436 | first.token_expert_indices.data_ptr() != second.token_expert_indices.data_ptr() |
| 437 | ) |
| 438 | torch.testing.assert_close( |
| 439 | first.token_expert_indices, torch.arange(8, dtype=torch.int32, device="cuda") |
| 440 | ) |
| 441 | |
| 442 | |
| 443 | def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: |
| 444 | if not moe_permute_unpermute_supported(): |
| 445 | pytest.skip("moe_permute_unpermute is not supported on this platform.") |
| 446 | |
| 447 | hidden_states = torch.arange(5 * 16, dtype=torch.bfloat16, device="cuda").view( |
| 448 | 5, 16 |
| 449 | ) |
| 450 | topk_ids = torch.tensor([[0], [-1], [1], [4], [2]], device="cuda") |
| 451 | expert_map = torch.tensor([0, 1, -1, -1], dtype=torch.int32, device="cuda") |
| 452 | scratch = MoEPermuteScratch( |
| 453 | max_num_tokens=5, |
| 454 | topk=1, |
| 455 | num_experts=4, |
| 456 | num_local_experts=2, |
| 457 | device=hidden_states.device, |
| 458 | hidden_size=16, |
| 459 | hidden_dtype=hidden_states.dtype, |
| 460 | ) |
| 461 | |
| 462 | permuted, _, expert_offsets, inverse, _ = moe_permute( |
| 463 | hidden_states=hidden_states, |
| 464 | a1q_scale=None, |
| 465 | topk_ids=topk_ids, |
| 466 | n_expert=4, |
| 467 | n_local_expert=2, |
| 468 | expert_map=expert_map, |
| 469 | scratch=scratch, |
| 470 | ) |
| 471 | |
| 472 | torch.testing.assert_close( |
| 473 | expert_offsets, |
| 474 | torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"), |
| 475 | ) |
| 476 | torch.testing.assert_close(permuted[:2], hidden_states[[0, 2]]) |
| 477 | assert torch.all(inverse[[1, 3, 4]] >= expert_offsets[-1]) |
| 478 | expected = torch.zeros_like(hidden_states) |
| 479 | expected[[0, 2]] = hidden_states[[0, 2]] |
| 480 | output = torch.empty_like(hidden_states) |
| 481 | moe_unpermute( |
| 482 | output, permuted, torch.ones(5, 1, device="cuda"), inverse, expert_offsets |
| 483 | ) |
| 484 | torch.testing.assert_close(output, expected) |
| 485 | |
| 486 | expected_inverse = inverse.clone() |
| 487 | expected_offsets = expert_offsets.clone() |
| 488 | expert_offsets, indices = moe_prepare_scatter(topk_ids, expert_map, scratch) |
| 489 | torch.testing.assert_close(indices.flatten(), expected_inverse) |
| 490 | torch.testing.assert_close(expert_offsets, expected_offsets) |
| 491 | graph = torch.cuda.CUDAGraph() |
| 492 | with torch.cuda.graph(graph): |
| 493 | expert_offsets, indices = moe_prepare_scatter(topk_ids, expert_map, scratch) |
| 494 | topk_ids.zero_() |
| 495 | graph.replay() |
| 496 | |
| 497 | expected_inverse = torch.arange(5, dtype=torch.int32, device="cuda") |
| 498 | expected_offsets = torch.tensor([0, 5, 5], dtype=torch.int64, device="cuda") |
| 499 | torch.testing.assert_close(indices.flatten(), expected_inverse) |
| 500 | torch.testing.assert_close(expert_offsets, expected_offsets) |