| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | import torch.nn.functional as F |
| 8 | |
| 9 | import vllm._custom_ops as ops |
| 10 | from vllm.model_executor.layers.fused_moe.config import ( |
| 11 | RoutingMethodType, |
| 12 | get_routing_method_type, |
| 13 | ) |
| 14 | from vllm.model_executor.layers.fused_moe.router.dsv4_topk import dsv4_topk |
| 15 | from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( |
| 16 | fused_topk_bias, |
| 17 | ) |
| 18 | from vllm.platforms import current_platform |
| 19 | |
| 20 | |
| 21 | def _torch_topk_softplus_sqrt( |
| 22 | gating_output: torch.Tensor, |
| 23 | topk: int, |
| 24 | renormalize: bool, |
| 25 | routed_scaling_factor: float, |
| 26 | e_score_correction_bias: torch.Tensor | None = None, |
| 27 | input_ids: torch.Tensor | None = None, |
| 28 | hash_indices_table: torch.Tensor | None = None, |
| 29 | bias_vl: torch.Tensor | None = None, |
| 30 | image_sentinel_lo: int = 0, |
| 31 | ): |
| 32 | scores = F.softplus(gating_output.float()).sqrt() |
| 33 | original_scores = scores |
| 34 | |
| 35 | if hash_indices_table is not None: |
| 36 | assert input_ids is not None |
| 37 | if bias_vl is None or image_sentinel_lo == 0: |
| 38 | topk_ids = hash_indices_table[input_ids.long()] |
| 39 | else: |
| 40 | image_mask = (input_ids.long() >= image_sentinel_lo) & ( |
| 41 | input_ids.long() < image_sentinel_lo + 5 |
| 42 | ) |
| 43 | topk_ids = hash_indices_table[input_ids.long()] |
| 44 | vl_ids = (scores + bias_vl).topk(topk, dim=-1).indices |
| 45 | topk_ids = torch.where( |
| 46 | image_mask.unsqueeze(-1), vl_ids.to(topk_ids.dtype), topk_ids |
| 47 | ) |
| 48 | else: |
| 49 | if bias_vl is not None and image_sentinel_lo > 0: |
| 50 | # Image tokens carry five consecutive in-vocab sentinel ids |
| 51 | # starting at image_sentinel_lo and select experts with bias_vl |
| 52 | # instead of the regular bias. |
| 53 | assert input_ids is not None |
| 54 | assert e_score_correction_bias is not None |
| 55 | image_mask = ( |
| 56 | (input_ids.long() >= image_sentinel_lo) |
| 57 | & (input_ids.long() < image_sentinel_lo + 5) |
| 58 | ).unsqueeze(-1) |
| 59 | row_bias = torch.where( |
| 60 | image_mask, |
| 61 | bias_vl.unsqueeze(0), |
| 62 | e_score_correction_bias.unsqueeze(0), |
| 63 | ) |
| 64 | scores_for_choice = scores + row_bias |
| 65 | elif e_score_correction_bias is not None: |
| 66 | scores_for_choice = scores + e_score_correction_bias.unsqueeze(0) |
| 67 | else: |
| 68 | scores_for_choice = scores |
| 69 | # Match the fused kernel's deterministic tie-break: lower expert ids |
| 70 | # win when scores are equal. torch.topk does not guarantee which tied |
| 71 | # index it selects at the k-th boundary. |
| 72 | topk_ids = torch.argsort( |
| 73 | scores_for_choice, dim=-1, descending=True, stable=True |
| 74 | )[:, :topk] |
| 75 | |
| 76 | topk_weights = original_scores.gather(1, topk_ids.long()) |
| 77 | if renormalize: |
| 78 | topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) |
| 79 | if routed_scaling_factor != 1.0: |
| 80 | topk_weights = topk_weights * routed_scaling_factor |
| 81 | return topk_weights.to(torch.float32), topk_ids.to(torch.int32) |
| 82 | |
| 83 | |
| 84 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend selection") |
| 85 | @pytest.mark.parametrize("has_hash_routing", [False, True]) |
| 86 | @pytest.mark.parametrize( |
| 87 | "backend,expected_experts", |
| 88 | [ |
| 89 | ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"), |
| 90 | ("aiter", "AiterExperts"), |
| 91 | ("triton_unfused", "UnfusedOAITritonExperts"), |
| 92 | ], |
| 93 | ) |
| 94 | def test_hash_routing_backend_selection( |
| 95 | dist_init, |
| 96 | default_vllm_config, |
| 97 | monkeypatch, |
| 98 | has_hash_routing, |
| 99 | backend, |
| 100 | expected_experts, |
| 101 | ): |
| 102 | """Hash tables must reach backend selection; modular AITER remains eligible.""" |
| 103 | from dataclasses import replace |
| 104 | |
| 105 | from vllm._aiter_ops import rocm_aiter_ops |
| 106 | from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory |
| 107 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 108 | select_deepseek_v4_mxfp4_moe_backend, |
| 109 | ) |
| 110 | from vllm.platforms import rocm |
| 111 | |
| 112 | # Exercise both native MX and monolithic AITER selection on any ROCm device. |
| 113 | monkeypatch.setattr(rocm, "on_gfx950", lambda: True) |
| 114 | monkeypatch.setattr(rocm, "on_gfx1250", lambda: False) |
| 115 | monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) |
| 116 | monkeypatch.setattr(rocm_aiter_ops, "is_fused_moe_enabled", lambda: True) |
| 117 | default_vllm_config.kernel_config.moe_backend = "triton" |
| 118 | layer = FusedMoEFactory( |
| 119 | num_experts=8, |
| 120 | top_k=2, |
| 121 | hidden_size=256, |
| 122 | intermediate_size=256, |
| 123 | params_dtype=torch.bfloat16, |
| 124 | scoring_func="sqrtsoftplus", |
| 125 | e_score_correction_bias=torch.zeros(8), |
| 126 | hash_indices_table=( |
| 127 | torch.tensor([[0, 1], [2, 3]], dtype=torch.int32) |
| 128 | if has_hash_routing |
| 129 | else None |
| 130 | ), |
| 131 | prefix="hash_routing_selection", |
| 132 | ) |
| 133 | config = replace(layer.routed_experts.moe_config, moe_backend=backend) |
| 134 | |
| 135 | if has_hash_routing and backend == "aiter_triton_mxfp4_bf16": |
| 136 | with pytest.raises(ValueError, match="hash routing"): |
| 137 | select_deepseek_v4_mxfp4_moe_backend(config) |
| 138 | else: |
| 139 | _, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config) |
| 140 | assert experts_cls.__name__ == expected_experts |
| 141 | |
| 142 | |
| 143 | def test_torch_topk_softplus_sqrt_breaks_ties_by_expert_id(): |
| 144 | gating_output = torch.tensor([[2.0, 1.0, 1.0, 0.0]]) |
| 145 | |
| 146 | _, topk_ids = _torch_topk_softplus_sqrt( |
| 147 | gating_output, |
| 148 | topk=2, |
| 149 | renormalize=False, |
| 150 | routed_scaling_factor=1.0, |
| 151 | ) |
| 152 | |
| 153 | assert topk_ids.tolist() == [[0, 1]] |
| 154 | |
| 155 | |
| 156 | def test_sqrtsoftplus_bias_uses_deepseek_v4_routing_method(): |
| 157 | assert ( |
| 158 | get_routing_method_type( |
| 159 | scoring_func="sqrtsoftplus", |
| 160 | top_k=8, |
| 161 | renormalize=True, |
| 162 | num_expert_group=None, |
| 163 | has_e_score_bias=True, |
| 164 | ) |
| 165 | == RoutingMethodType.DeepseekV4 |
| 166 | ) |
| 167 | assert ( |
| 168 | get_routing_method_type( |
| 169 | scoring_func="sqrtsoftplus", |
| 170 | top_k=8, |
| 171 | renormalize=False, |
| 172 | num_expert_group=None, |
| 173 | has_e_score_bias=True, |
| 174 | ) |
| 175 | == RoutingMethodType.Unspecified |
| 176 | ) |
| 177 | |
| 178 | |
| 179 | @pytest.mark.skipif( |
| 180 | not current_platform.is_cuda_alike(), |
| 181 | reason="This test is skipped on non-CUDA platform.", |
| 182 | ) |
| 183 | @pytest.mark.parametrize("num_tokens", [1, 33, 128]) |
| 184 | @pytest.mark.parametrize("hidden_size", [1024, 2048]) |
| 185 | @pytest.mark.parametrize("num_experts", [128, 256, 384, 512]) |
| 186 | @pytest.mark.parametrize("topk", [6, 8, 16]) |
| 187 | @pytest.mark.parametrize("renormalize", [True, False]) |
| 188 | @pytest.mark.parametrize("routed_scaling_factor", [1.0, 1.5]) |
| 189 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32]) |
| 190 | def test_fused_topk_softplus_sqrt( |
| 191 | num_tokens: int, |
| 192 | hidden_size: int, |
| 193 | num_experts: int, |
| 194 | topk: int, |
| 195 | renormalize: bool, |
| 196 | routed_scaling_factor: float, |
| 197 | dtype: torch.dtype, |
| 198 | ): |
| 199 | torch.manual_seed(0) |
| 200 | hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda") |
| 201 | gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") |
| 202 | e_score_correction_bias = torch.randn( |
| 203 | (num_experts,), dtype=torch.float32, device="cuda" |
| 204 | ) |
| 205 | |
| 206 | topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( |
| 207 | gating_output=gating_output, |
| 208 | topk=topk, |
| 209 | renormalize=renormalize, |
| 210 | routed_scaling_factor=routed_scaling_factor, |
| 211 | e_score_correction_bias=e_score_correction_bias, |
| 212 | ) |
| 213 | |
| 214 | topk_weights, topk_ids = fused_topk_bias( |
| 215 | hidden_states=hidden_states, |
| 216 | gating_output=gating_output, |
| 217 | scoring_func="sqrtsoftplus", |
| 218 | e_score_correction_bias=e_score_correction_bias, |
| 219 | topk=topk, |
| 220 | renormalize=renormalize, |
| 221 | routed_scaling_factor=routed_scaling_factor, |
| 222 | ) |
| 223 | |
| 224 | # Different kernels may return the topk experts in different orders when |
| 225 | # scores tie; sort by expert id before comparing. |
| 226 | sorted_ref_ids, idx_ref = topk_ids_ref.sort(dim=-1) |
| 227 | sorted_ids, idx_ops = topk_ids.sort(dim=-1) |
| 228 | torch.testing.assert_close(sorted_ref_ids, sorted_ids, atol=0, rtol=0) |
| 229 | |
| 230 | sorted_w_ref = topk_weights_ref.gather(1, idx_ref) |
| 231 | sorted_w = topk_weights.gather(1, idx_ops) |
| 232 | torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2) |
| 233 | |
| 234 | |
| 235 | @pytest.mark.skipif( |
| 236 | not current_platform.is_cuda_alike(), |
| 237 | reason="This test is skipped on non-CUDA platform.", |
| 238 | ) |
| 239 | @pytest.mark.parametrize("num_tokens", [1, 33, 128]) |
| 240 | @pytest.mark.parametrize("hidden_size", [1024, 2048]) |
| 241 | @pytest.mark.parametrize("num_experts", [256, 384, 512]) |
| 242 | @pytest.mark.parametrize("topk", [6, 8, 16]) |
| 243 | @pytest.mark.parametrize("renormalize", [True, False]) |
| 244 | @pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5]) |
| 245 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32]) |
| 246 | def test_fused_topk_softplus_sqrt_hash( |
| 247 | num_tokens: int, |
| 248 | hidden_size: int, |
| 249 | num_experts: int, |
| 250 | topk: int, |
| 251 | renormalize: bool, |
| 252 | routed_scaling_factor: float, |
| 253 | dtype: torch.dtype, |
| 254 | ): |
| 255 | torch.manual_seed(0) |
| 256 | vocab_size = 1024 |
| 257 | hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda") |
| 258 | gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") |
| 259 | # Per-token fixed expert selection: for each vocab id pick `topk` distinct |
| 260 | # experts. |
| 261 | hash_indices_table = torch.stack( |
| 262 | [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)] |
| 263 | ).to(device="cuda", dtype=torch.long) |
| 264 | input_ids = torch.randint( |
| 265 | 0, vocab_size, (num_tokens,), dtype=torch.long, device="cuda" |
| 266 | ) |
| 267 | |
| 268 | topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( |
| 269 | gating_output=gating_output, |
| 270 | topk=topk, |
| 271 | renormalize=renormalize, |
| 272 | routed_scaling_factor=routed_scaling_factor, |
| 273 | input_ids=input_ids, |
| 274 | hash_indices_table=hash_indices_table, |
| 275 | ) |
| 276 | |
| 277 | topk_weights, topk_ids = fused_topk_bias( |
| 278 | hidden_states=hidden_states, |
| 279 | gating_output=gating_output, |
| 280 | scoring_func="sqrtsoftplus", |
| 281 | e_score_correction_bias=None, |
| 282 | topk=topk, |
| 283 | renormalize=renormalize, |
| 284 | input_tokens=input_ids, |
| 285 | hash_indices_table=hash_indices_table, |
| 286 | routed_scaling_factor=routed_scaling_factor, |
| 287 | ) |
| 288 | |
| 289 | sorted_ref_ids, idx_ref = topk_ids_ref.sort(dim=-1) |
| 290 | sorted_ids, idx_ops = topk_ids.sort(dim=-1) |
| 291 | torch.testing.assert_close(sorted_ref_ids, sorted_ids, atol=0, rtol=0) |
| 292 | |
| 293 | sorted_w_ref = topk_weights_ref.gather(1, idx_ref) |
| 294 | sorted_w = topk_weights.gather(1, idx_ops) |
| 295 | torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2) |
| 296 | |
| 297 | |
| 298 | @pytest.mark.skipif( |
| 299 | not current_platform.is_cuda(), |
| 300 | reason="The DeepSeek V4 fast path is CUDA-only.", |
| 301 | ) |
| 302 | @pytest.mark.parametrize( |
| 303 | ("num_tokens", "num_experts", "indices_type"), |
| 304 | [ |
| 305 | (0, 256, torch.uint32), |
| 306 | (17, 256, torch.uint32), |
| 307 | (17, 384, torch.int64), |
| 308 | ], |
| 309 | ) |
| 310 | def test_dsv4_fast_topk( |
| 311 | num_tokens: int, |
| 312 | num_experts: int, |
| 313 | indices_type: torch.dtype, |
| 314 | ): |
| 315 | torch.manual_seed(0) |
| 316 | gating_output = torch.randn( |
| 317 | (num_tokens, num_experts), dtype=torch.float32, device="cuda" |
| 318 | ) |
| 319 | correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda") |
| 320 | |
| 321 | topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( |
| 322 | gating_output=gating_output, |
| 323 | topk=6, |
| 324 | renormalize=True, |
| 325 | routed_scaling_factor=1.5, |
| 326 | e_score_correction_bias=correction_bias, |
| 327 | ) |
| 328 | topk_weights, topk_ids = dsv4_topk( |
| 329 | gating_output, correction_bias, indices_type, 1.5 |
| 330 | ) |
| 331 | |
| 332 | assert topk_ids.dtype == indices_type |
| 333 | torch.testing.assert_close(topk_ids_ref.to(indices_type), topk_ids, atol=0, rtol=0) |
| 334 | torch.testing.assert_close( |
| 335 | topk_weights_ref, |
| 336 | topk_weights, |
| 337 | atol=2e-5, |
| 338 | rtol=2e-5, |
| 339 | ) |
| 340 | |
| 341 | |
| 342 | @pytest.mark.skipif( |
| 343 | not current_platform.is_cuda_alike(), |
| 344 | reason="This test is skipped on non-CUDA platform.", |
| 345 | ) |
| 346 | @pytest.mark.parametrize("use_hash", [False, True]) |
| 347 | @pytest.mark.parametrize("use_bias", [False, True]) |
| 348 | @pytest.mark.parametrize("use_padding_mask", [False, True]) |
| 349 | @pytest.mark.parametrize("pad_with_nan", [False, True]) |
| 350 | @pytest.mark.parametrize("num_experts", [128, 256, 384]) |
| 351 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32]) |
| 352 | def test_fused_topk_softplus_sqrt_padding( |
| 353 | use_hash: bool, |
| 354 | use_bias: bool, |
| 355 | use_padding_mask: bool, |
| 356 | pad_with_nan: bool, |
| 357 | num_experts: int, |
| 358 | dtype: torch.dtype, |
| 359 | ): |
| 360 | """Verify explicit padding and NaN-padded rows do not affect real rows.""" |
| 361 | torch.manual_seed(0) |
| 362 | num_tokens = 8 |
| 363 | topk = 6 |
| 364 | indices_dtype = torch.int32 |
| 365 | |
| 366 | gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") |
| 367 | # Exercise top-k cutoff ties without relying on random low-precision collisions. |
| 368 | gating_output[0] = 0 |
| 369 | |
| 370 | padding_rows = torch.zeros(num_tokens, dtype=torch.bool, device="cuda") |
| 371 | padding_rows[1::2] = True |
| 372 | if pad_with_nan: |
| 373 | gating_output[padding_rows] = float("nan") |
| 374 | is_padding = padding_rows if use_padding_mask else None |
| 375 | |
| 376 | # A negative correction bias makes explicit pad rows look selectable unless |
| 377 | # the kernel uses the is_padding guard. |
| 378 | e_score_correction_bias = None |
| 379 | if use_bias: |
| 380 | e_score_correction_bias = ( |
| 381 | -torch.rand((num_experts,), dtype=torch.float32, device="cuda") - 1.0 |
| 382 | ) |
| 383 | |
| 384 | input_ids = None |
| 385 | hash_indices_table = None |
| 386 | if use_hash: |
| 387 | vocab_size = 64 |
| 388 | hash_indices_table = torch.stack( |
| 389 | [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)] |
| 390 | ).to(device="cuda", dtype=indices_dtype) |
| 391 | input_ids = torch.randint( |
| 392 | 0, vocab_size, (num_tokens,), dtype=indices_dtype, device="cuda" |
| 393 | ) |
| 394 | |
| 395 | topk_weights = torch.empty(num_tokens, topk, dtype=torch.float32, device="cuda") |
| 396 | topk_ids = torch.empty(num_tokens, topk, dtype=indices_dtype, device="cuda") |
| 397 | token_expert_indices = torch.empty( |
| 398 | num_tokens, topk, dtype=torch.int32, device="cuda" |
| 399 | ) |
| 400 | |
| 401 | ops.topk_hash_softplus_sqrt( |
| 402 | topk_weights, |
| 403 | topk_ids, |
| 404 | token_expert_indices, |
| 405 | gating_output, |
| 406 | renormalize=True, |
| 407 | routed_scaling_factor=1.0, |
| 408 | e_score_correction_bias=e_score_correction_bias, |
| 409 | input_tokens=input_ids, |
| 410 | hash_indices_table=hash_indices_table, |
| 411 | is_padding=is_padding, |
| 412 | ) |
| 413 | |
| 414 | if use_padding_mask: |
| 415 | pad_ids = topk_ids[padding_rows] |
| 416 | pad_weights = topk_weights[padding_rows] |
| 417 | assert torch.equal(pad_ids, torch.full_like(pad_ids, -1)), ( |
| 418 | f"Explicit pad rows should contain only -1 ids, got {pad_ids.tolist()}" |
| 419 | ) |
| 420 | assert (pad_weights == 0).all(), ( |
| 421 | "Explicit pad rows should have all-zero weights, " |
| 422 | f"got {pad_weights.tolist()}" |
| 423 | ) |
| 424 | |
| 425 | if pad_with_nan: |
| 426 | nan_pad_weights = topk_weights[padding_rows] |
| 427 | assert torch.isfinite(nan_pad_weights).all(), ( |
| 428 | f"NaN-padded rows have non-finite weights, got {nan_pad_weights.tolist()}" |
| 429 | ) |
| 430 | assert (nan_pad_weights == 0).all(), ( |
| 431 | "NaN-padded rows should have all-zero weights, " |
| 432 | f"got {nan_pad_weights.tolist()}" |
| 433 | ) |
| 434 | |
| 435 | topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( |
| 436 | gating_output=gating_output, |
| 437 | topk=topk, |
| 438 | renormalize=True, |
| 439 | routed_scaling_factor=1.0, |
| 440 | e_score_correction_bias=e_score_correction_bias, |
| 441 | input_ids=input_ids, |
| 442 | hash_indices_table=hash_indices_table, |
| 443 | ) |
| 444 | |
| 445 | rows_to_compare = torch.ones(num_tokens, dtype=torch.bool, device="cuda") |
| 446 | if use_padding_mask or pad_with_nan: |
| 447 | rows_to_compare = ~padding_rows |
| 448 | |
| 449 | sorted_ref_ids, idx_ref = topk_ids_ref[rows_to_compare].sort(dim=-1) |
| 450 | sorted_ids, idx_ops = topk_ids[rows_to_compare].sort(dim=-1) |
| 451 | torch.testing.assert_close( |
| 452 | sorted_ref_ids, sorted_ids.to(sorted_ref_ids.dtype), atol=0, rtol=0 |
| 453 | ) |
| 454 | |
| 455 | sorted_w_ref = topk_weights_ref[rows_to_compare].gather(1, idx_ref) |
| 456 | sorted_w = topk_weights[rows_to_compare].gather(1, idx_ops) |
| 457 | torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2) |
| 458 | |
| 459 | |
| 460 | def _make_mixed_input_ids( |
| 461 | num_tokens: int, image_sentinel_lo: int, dtype: torch.dtype = torch.long |
| 462 | ) -> torch.Tensor: |
| 463 | """Mix of text, image-sentinel, and above-sentinel ids. |
| 464 | |
| 465 | Every third token is an image sentinel id (cycling through the five slots |
| 466 | [image_sentinel_lo, image_sentinel_lo + 5)); every sixth token carries an |
| 467 | id just above the sentinel block, simulating the named special tokens |
| 468 | that follow it — those must route as regular text tokens. |
| 469 | """ |
| 470 | input_ids = torch.randint( |
| 471 | 0, image_sentinel_lo, (num_tokens,), dtype=torch.long, device="cuda" |
| 472 | ) |
| 473 | pos = torch.arange(num_tokens, device="cuda") |
| 474 | image_rows = pos % 3 == 0 |
| 475 | input_ids[image_rows] = image_sentinel_lo + (pos[image_rows] // 3) % 5 |
| 476 | above_rows = pos % 6 == 3 |
| 477 | input_ids[above_rows] = image_sentinel_lo + 5 |
| 478 | return input_ids.to(dtype) |
| 479 | |
| 480 | |
| 481 | def _assert_topk_matches( |
| 482 | topk_weights: torch.Tensor, |
| 483 | topk_ids: torch.Tensor, |
| 484 | topk_weights_ref: torch.Tensor, |
| 485 | topk_ids_ref: torch.Tensor, |
| 486 | ) -> None: |
| 487 | # Different kernels may return the topk experts in different orders when |
| 488 | # scores tie; sort by expert id before comparing. |
| 489 | sorted_ref_ids, idx_ref = topk_ids_ref.sort(dim=-1) |
| 490 | sorted_ids, idx_ops = topk_ids.sort(dim=-1) |
| 491 | torch.testing.assert_close( |
| 492 | sorted_ref_ids, sorted_ids.to(sorted_ref_ids.dtype), atol=0, rtol=0 |
| 493 | ) |
| 494 | |
| 495 | sorted_w_ref = topk_weights_ref.gather(1, idx_ref) |
| 496 | sorted_w = topk_weights.gather(1, idx_ops.to(torch.long)) |
| 497 | torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2) |
| 498 | |
| 499 | |
| 500 | @pytest.mark.skipif( |
| 501 | not current_platform.is_cuda_alike(), |
| 502 | reason="This test is skipped on non-CUDA platform.", |
| 503 | ) |
| 504 | @pytest.mark.parametrize( |
| 505 | ("num_experts", "topk", "renormalize", "dtype"), |
| 506 | [ |
| 507 | (256, 6, True, torch.float32), # Triton dsv4 fast path |
| 508 | (256, 8, True, torch.bfloat16), # CUDA topk kernel |
| 509 | (384, 6, False, torch.half), |
| 510 | (512, 16, True, torch.float32), |
| 511 | ], |
| 512 | ) |
| 513 | def test_fused_topk_softplus_sqrt_bias_vl( |
| 514 | num_experts: int, |
| 515 | topk: int, |
| 516 | renormalize: bool, |
| 517 | dtype: torch.dtype, |
| 518 | ): |
| 519 | """Image sentinel tokens must select experts with bias_vl. |
| 520 | |
| 521 | Rows with ids just above the sentinel block (image_sentinel_lo + 5) |
| 522 | simulate the named special tokens that follow it and must route through |
| 523 | the regular bias path; the reference comparison covers them. |
| 524 | """ |
| 525 | torch.manual_seed(0) |
| 526 | num_tokens = 64 |
| 527 | hidden_size = 1024 |
| 528 | vocab_size = 128 |
| 529 | image_sentinel_lo = vocab_size - 8 |
| 530 | hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda") |
| 531 | gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") |
| 532 | e_score_correction_bias = torch.randn( |
| 533 | (num_experts,), dtype=torch.float32, device="cuda" |
| 534 | ) |
| 535 | bias_vl = torch.randn((num_experts,), dtype=torch.float32, device="cuda") |
| 536 | input_ids = _make_mixed_input_ids(num_tokens, image_sentinel_lo) |
| 537 | assert (input_ids == image_sentinel_lo + 5).any() |
| 538 | |
| 539 | topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( |
| 540 | gating_output=gating_output, |
| 541 | topk=topk, |
| 542 | renormalize=renormalize, |
| 543 | routed_scaling_factor=1.5, |
| 544 | e_score_correction_bias=e_score_correction_bias, |
| 545 | input_ids=input_ids, |
| 546 | bias_vl=bias_vl, |
| 547 | image_sentinel_lo=image_sentinel_lo, |
| 548 | ) |
| 549 | |
| 550 | topk_weights, topk_ids = fused_topk_bias( |
| 551 | hidden_states=hidden_states, |
| 552 | gating_output=gating_output, |
| 553 | scoring_func="sqrtsoftplus", |
| 554 | e_score_correction_bias=e_score_correction_bias, |
| 555 | topk=topk, |
| 556 | renormalize=renormalize, |
| 557 | routed_scaling_factor=1.5, |
| 558 | input_tokens=input_ids, |
| 559 | bias_vl=bias_vl, |
| 560 | image_sentinel_lo=image_sentinel_lo, |
| 561 | ) |
| 562 | |
| 563 | _assert_topk_matches(topk_weights, topk_ids, topk_weights_ref, topk_ids_ref) |
| 564 | |
| 565 | |
| 566 | @pytest.mark.skipif( |
| 567 | not current_platform.is_cuda_alike(), |
| 568 | reason="This test is skipped on non-CUDA platform.", |
| 569 | ) |
| 570 | @pytest.mark.parametrize( |
| 571 | ("num_experts", "topk", "renormalize", "dtype", "indices_dtype"), |
| 572 | [ |
| 573 | (256, 6, True, torch.float32, torch.long), # specialized hash kernel |
| 574 | (384, 6, True, torch.float32, torch.int32), # specialized hash kernel |
| 575 | (384, 8, False, torch.bfloat16, torch.long), # generic hash kernel |
| 576 | (512, 6, True, torch.float32, torch.long), # generic hash kernel |
| 577 | ], |
| 578 | ) |
| 579 | def test_fused_topk_softplus_sqrt_hash_bias_vl( |
| 580 | num_experts: int, |
| 581 | topk: int, |
| 582 | renormalize: bool, |
| 583 | dtype: torch.dtype, |
| 584 | indices_dtype: torch.dtype, |
| 585 | ): |
| 586 | """Hash MoE: text rows use tid2eid, image rows use topk(score + bias_vl). |
| 587 | |
| 588 | Sentinel ids are in-vocab; rows with ids above the sentinel block |
| 589 | (image_sentinel_lo + 5) must still come straight from the hash table. |
| 590 | """ |
| 591 | torch.manual_seed(0) |
| 592 | num_tokens = 64 |
| 593 | hidden_size = 1024 |
| 594 | vocab_size = 64 |
| 595 | image_sentinel_lo = vocab_size - 8 |
| 596 | hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda") |
| 597 | gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") |
| 598 | hash_indices_table = torch.stack( |
| 599 | [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)] |
| 600 | ).to(device="cuda", dtype=indices_dtype) |
| 601 | bias_vl = torch.randn((num_experts,), dtype=torch.float32, device="cuda") |
| 602 | input_ids = _make_mixed_input_ids(num_tokens, image_sentinel_lo) |
| 603 | image_rows = (input_ids >= image_sentinel_lo) & (input_ids < image_sentinel_lo + 5) |
| 604 | assert image_rows.any() and (~image_rows).any() |
| 605 | assert (input_ids == image_sentinel_lo + 5).any() |
| 606 | |
| 607 | topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( |
| 608 | gating_output=gating_output, |
| 609 | topk=topk, |
| 610 | renormalize=renormalize, |
| 611 | routed_scaling_factor=2.5, |
| 612 | input_ids=input_ids, |
| 613 | hash_indices_table=hash_indices_table, |
| 614 | bias_vl=bias_vl, |
| 615 | image_sentinel_lo=image_sentinel_lo, |
| 616 | ) |
| 617 | |
| 618 | topk_weights, topk_ids = fused_topk_bias( |
| 619 | hidden_states=hidden_states, |
| 620 | gating_output=gating_output, |
| 621 | scoring_func="sqrtsoftplus", |
| 622 | e_score_correction_bias=None, |
| 623 | topk=topk, |
| 624 | renormalize=renormalize, |
| 625 | input_tokens=input_ids, |
| 626 | hash_indices_table=hash_indices_table, |
| 627 | routed_scaling_factor=2.5, |
| 628 | bias_vl=bias_vl, |
| 629 | image_sentinel_lo=image_sentinel_lo, |
| 630 | ) |
| 631 | |
| 632 | _assert_topk_matches(topk_weights, topk_ids, topk_weights_ref, topk_ids_ref) |
| 633 | |
| 634 | # Text rows must come straight from the hash table, image rows from |
| 635 | # topk(score + bias_vl). |
| 636 | text_ids = hash_indices_table[input_ids[~image_rows].long()] |
| 637 | torch.testing.assert_close( |
| 638 | topk_ids[~image_rows].to(text_ids.dtype), text_ids, atol=0, rtol=0 |
| 639 | ) |
| 640 | |
| 641 | |
| 642 | @pytest.mark.skipif( |
| 643 | not current_platform.is_cuda(), |
| 644 | reason="The DeepSeek V4 fast path is CUDA-only.", |
| 645 | ) |
| 646 | def test_dsv4_fast_topk_bias_vl(): |
| 647 | torch.manual_seed(0) |
| 648 | num_tokens = 33 |
| 649 | num_experts = 256 |
| 650 | vocab_size = 128 |
| 651 | image_sentinel_lo = vocab_size - 8 |
| 652 | gating_output = torch.randn( |
| 653 | (num_tokens, num_experts), dtype=torch.float32, device="cuda" |
| 654 | ) |
| 655 | correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda") |
| 656 | bias_vl = torch.randn(num_experts, dtype=torch.float32, device="cuda") |
| 657 | input_ids = _make_mixed_input_ids(num_tokens, image_sentinel_lo) |
| 658 | |
| 659 | topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( |
| 660 | gating_output=gating_output, |
| 661 | topk=6, |
| 662 | renormalize=True, |
| 663 | routed_scaling_factor=1.5, |
| 664 | e_score_correction_bias=correction_bias, |
| 665 | input_ids=input_ids, |
| 666 | bias_vl=bias_vl, |
| 667 | image_sentinel_lo=image_sentinel_lo, |
| 668 | ) |
| 669 | topk_weights, topk_ids = dsv4_topk( |
| 670 | gating_output, |
| 671 | correction_bias, |
| 672 | torch.int64, |
| 673 | 1.5, |
| 674 | input_ids=input_ids, |
| 675 | bias_vl=bias_vl, |
| 676 | image_sentinel_lo=image_sentinel_lo, |
| 677 | ) |
| 678 | |
| 679 | assert topk_ids.dtype == torch.int64 |
| 680 | torch.testing.assert_close(topk_ids_ref.to(torch.int64), topk_ids, atol=0, rtol=0) |
| 681 | torch.testing.assert_close(topk_weights_ref, topk_weights, atol=2e-5, rtol=2e-5) |