| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from types import SimpleNamespace |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | import torch.nn as nn |
| 8 | import torch.nn.functional as F |
| 9 | |
| 10 | import vllm.model_executor.kernels.mhc # noqa: F401 |
| 11 | import vllm.model_executor.layers.mhc as mhc_layers |
| 12 | from vllm.model_executor.kernels.mhc.tilelang import ( |
| 13 | _torch_hc_prenorm_gemm, |
| 14 | mhc_post_tilelang, |
| 15 | mhc_pre_delayed_tilelang, |
| 16 | ) |
| 17 | from vllm.model_executor.kernels.mhc.tilelang_kernels import ( |
| 18 | _HC_PRENORM_GEMM_TILELANG_KERNEL, |
| 19 | mhc_fused_post_pre_split_config, |
| 20 | ) |
| 21 | from vllm.model_executor.kernels.mhc.torch import ( |
| 22 | mhc_post_torch, |
| 23 | mhc_pre_delayed_torch, |
| 24 | ) |
| 25 | from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton |
| 26 | from vllm.model_executor.layers.mhc import ( |
| 27 | HAS_AITER_MHC, |
| 28 | HAS_AITER_MHC_FUSED, |
| 29 | HAS_AITER_MHC_FUSED_NORM, |
| 30 | HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM, |
| 31 | HAS_AITER_MHC_PRE_NORM, |
| 32 | HAS_TILELANG_MHC, |
| 33 | MHCFusedPostPreOp, |
| 34 | MHCPreDelayedOp, |
| 35 | MHCPreOp, |
| 36 | ) |
| 37 | from vllm.models.deepseek_v4.nvidia.model import ( |
| 38 | DeepseekV4DecoderLayer, |
| 39 | DeepseekV4Model, |
| 40 | ) |
| 41 | from vllm.models.deepseek_v41.nvidia.model import ( |
| 42 | DeepseekV4DecoderLayer as DeepseekV41DecoderLayer, |
| 43 | ) |
| 44 | from vllm.models.deepseek_v41.nvidia.ops.mega_mhc import ( |
| 45 | is_mega_mhc_supported, |
| 46 | mhc_shifted_post_pre_deep_gemm, |
| 47 | ) |
| 48 | from vllm.platforms import current_platform |
| 49 | from vllm.utils.torch_utils import set_random_seed |
| 50 | |
| 51 | DEVICE = current_platform.device_type |
| 52 | |
| 53 | |
| 54 | @pytest.mark.parametrize( |
| 55 | "tp,ep,hidden,hc,multicast,expected", |
| 56 | [ |
| 57 | (4, False, 5120, 4, 1, True), |
| 58 | (2, False, 5120, 4, 1, False), |
| 59 | (4, True, 5120, 4, 1, False), |
| 60 | (4, False, 4096, 4, 1, False), |
| 61 | (4, False, 5120, 2, 1, False), |
| 62 | (4, False, 5120, 4, 0, False), |
| 63 | ], |
| 64 | ) |
| 65 | def test_deepseek_v41_all_reduce_fusion_requires_kernel_support( |
| 66 | monkeypatch, tp, ep, hidden, hc, multicast, expected |
| 67 | ): |
| 68 | """Eligibility depends on kernel inputs, not serving or attention settings.""" |
| 69 | from vllm.models.deepseek_v41.nvidia.ops import mhc |
| 70 | |
| 71 | config = SimpleNamespace( |
| 72 | parallel_config=SimpleNamespace( |
| 73 | tensor_parallel_size=tp, enable_expert_parallel=ep |
| 74 | ), |
| 75 | model_config=SimpleNamespace( |
| 76 | hf_config=SimpleNamespace(hidden_size=hidden, hc_mult=hc) |
| 77 | ), |
| 78 | ) |
| 79 | group = SimpleNamespace( |
| 80 | device_communicator=SimpleNamespace( |
| 81 | ca_comm=SimpleNamespace(mnnvl_lamport_ag_multicast_ptr=multicast) |
| 82 | ) |
| 83 | ) |
| 84 | monkeypatch.setattr(mhc, "get_tp_group", lambda: group) |
| 85 | assert mhc.supports_mhc_all_reduce(config) is expected |
| 86 | |
| 87 | |
| 88 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 89 | @pytest.mark.parametrize( |
| 90 | "num_tokens,hc_mult,hidden_size", |
| 91 | [ |
| 92 | (0, 4, 5120), |
| 93 | (1, 4, 5120), |
| 94 | (33, 4, 5137), |
| 95 | (128, 4, 5120), |
| 96 | (1024, 4, 5120), |
| 97 | (2, 1, 513), |
| 98 | (7, 8, 5137), |
| 99 | (65537, 4, 65), |
| 100 | ], |
| 101 | ) |
| 102 | @pytest.mark.parametrize("strided", [False, True]) |
| 103 | def test_hc_collapse_preserves_weighted_residual_sum( |
| 104 | num_tokens, hc_mult, hidden_size, strided |
| 105 | ): |
| 106 | """Handle tile tails and strided streams without changing FP32 mixing.""" |
| 107 | set_random_seed(0) |
| 108 | x = torch.randn( |
| 109 | num_tokens, |
| 110 | hc_mult, |
| 111 | hidden_size * (2 if strided else 1), |
| 112 | dtype=torch.bfloat16, |
| 113 | device=DEVICE, |
| 114 | ) |
| 115 | pre = torch.rand(num_tokens, hc_mult, device=DEVICE) |
| 116 | if strided: |
| 117 | x = x[..., ::2] |
| 118 | pre = pre.t().contiguous().t() |
| 119 | if num_tokens: |
| 120 | pre[0].zero_() |
| 121 | pre[0, -1] = 1 |
| 122 | if num_tokens > 1 and hc_mult == 4: |
| 123 | pre[1].fill_(1) |
| 124 | x[1, 1] = -x[1, 0] |
| 125 | x[1, 3] = -x[1, 2] |
| 126 | expected = (pre.unsqueeze(-1) * x.float()).sum(dim=1).to(x.dtype) |
| 127 | actual = hc_collapse_triton(x, pre) |
| 128 | assert actual.is_contiguous() |
| 129 | # BF16 rounding only: ptxas <13.1 contracts the FP32 mix into FFMA despite |
| 130 | # enable_fp_fusion=False, so compare within a ULP and pin the exact cases below. |
| 131 | torch.testing.assert_close(actual, expected, atol=1.6e-2, rtol=1e-2) |
| 132 | if num_tokens: |
| 133 | torch.testing.assert_close(actual[0], x[0, -1], atol=0, rtol=0) |
| 134 | if num_tokens > 1 and hc_mult == 4: |
| 135 | torch.testing.assert_close( |
| 136 | actual[1], torch.zeros_like(actual[1]), atol=0, rtol=0 |
| 137 | ) |
| 138 | |
| 139 | |
| 140 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 141 | def test_hc_collapse_custom_op_supports_compile(): |
| 142 | x = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE) |
| 143 | pre = torch.rand(2, 4, device=DEVICE) |
| 144 | torch.library.opcheck(torch.ops.vllm.hc_collapse_triton.default, (x, pre)) |
| 145 | |
| 146 | |
| 147 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 148 | @pytest.mark.parametrize("num_tokens", [1, 7]) |
| 149 | def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch): |
| 150 | """Match reference DSparkBlock.forward_head's hc_pre before its RMSNorm.""" |
| 151 | from vllm.models.deepseek_v41.nvidia import dspark |
| 152 | |
| 153 | set_random_seed(0) |
| 154 | hidden_size, hc_mult = 5120, 4 |
| 155 | streams = torch.randn( |
| 156 | num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE |
| 157 | ) |
| 158 | mixes = [torch.rand(num_tokens, hc_mult, device=DEVICE) for _ in range(2)] |
| 159 | carried_mixes = [] |
| 160 | |
| 161 | def make_layer(mix): |
| 162 | def forward(hidden, positions, ids, pre, post, res, residual, **kwargs): |
| 163 | carried_mixes.append(pre) |
| 164 | return hidden, None, None, None, mix, None |
| 165 | |
| 166 | return forward |
| 167 | |
| 168 | monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams) |
| 169 | draft = SimpleNamespace( |
| 170 | use_mega_moe=False, |
| 171 | use_sequence_parallel=False, |
| 172 | hc_mult=hc_mult, |
| 173 | layers=[make_layer(mix) for mix in mixes], |
| 174 | ) |
| 175 | actual = dspark.DSparkDeepseekV4Model.forward( |
| 176 | draft, |
| 177 | input_ids=torch.zeros(num_tokens, dtype=torch.long, device=DEVICE), |
| 178 | positions=torch.arange(num_tokens, device=DEVICE), |
| 179 | inputs_embeds=torch.zeros( |
| 180 | num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE |
| 181 | ), |
| 182 | ) |
| 183 | expected = (streams.float() * mixes[-1].unsqueeze(-1)).sum(dim=1) |
| 184 | assert carried_mixes[0] is None |
| 185 | assert carried_mixes[1] is mixes[0] |
| 186 | torch.testing.assert_close( |
| 187 | actual, expected.to(streams.dtype), atol=1.6e-2, rtol=1e-2 |
| 188 | ) |
| 189 | |
| 190 | |
| 191 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 192 | @pytest.mark.parametrize("num_tokens", [1, 7, 32]) |
| 193 | @pytest.mark.parametrize("tp_size", [4, 8]) |
| 194 | def test_hc_collapse_before_gather_is_bitwise_equal(num_tokens, tp_size): |
| 195 | """Simulate SP on one GPU, including padding and ranks without valid tokens.""" |
| 196 | set_random_seed(0) |
| 197 | padded_tokens = num_tokens + (-num_tokens) % tp_size |
| 198 | streams = torch.randn(padded_tokens, 4, 5120, dtype=torch.bfloat16, device=DEVICE) |
| 199 | pre_mix = torch.rand(padded_tokens, 4, device=DEVICE) |
| 200 | stream_shards = streams.chunk(tp_size) |
| 201 | mix_shards = pre_mix.chunk(tp_size) |
| 202 | before = hc_collapse_triton( |
| 203 | torch.cat(stream_shards)[:num_tokens], torch.cat(mix_shards)[:num_tokens] |
| 204 | ) |
| 205 | after = torch.cat( |
| 206 | [hc_collapse_triton(x, pre) for x, pre in zip(stream_shards, mix_shards)] |
| 207 | )[:num_tokens] |
| 208 | assert after.shape == (num_tokens, 5120) |
| 209 | torch.testing.assert_close(after, before, atol=0, rtol=0) |
| 210 | |
| 211 | |
| 212 | @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") |
| 213 | @pytest.mark.parametrize( |
| 214 | "num_tokens,sinkhorn_repeat,norm_eps", |
| 215 | [ |
| 216 | (0, 20, 1e-6), |
| 217 | (1, 1, 1e-6), |
| 218 | (1, 20, 1e-6), |
| 219 | (33, 20, 1e-20), |
| 220 | (128, 20, 1e-6), |
| 221 | (1024, 20, 1e-6), |
| 222 | ], |
| 223 | ) |
| 224 | @pytest.mark.parametrize("entry", ["broadcast", "identity", "carried"]) |
| 225 | @pytest.mark.parametrize("use_deep_gemm", [False, True]) |
| 226 | @pytest.mark.parametrize("fused_norm", [False, True]) |
| 227 | def test_deepseek_v41_mhc_pre_delayed( |
| 228 | num_tokens, |
| 229 | sinkhorn_repeat, |
| 230 | norm_eps, |
| 231 | entry, |
| 232 | use_deep_gemm, |
| 233 | fused_norm, |
| 234 | monkeypatch, |
| 235 | ): |
| 236 | """Collapse with the carried pre-mix and save the new one for the next sublayer.""" |
| 237 | from vllm.utils.deep_gemm import is_deep_gemm_supported |
| 238 | |
| 239 | if use_deep_gemm and not is_deep_gemm_supported(): |
| 240 | pytest.skip("DeepGEMM not available") |
| 241 | monkeypatch.setattr( |
| 242 | "vllm.utils.deep_gemm.is_deep_gemm_supported", lambda: use_deep_gemm |
| 243 | ) |
| 244 | set_random_seed(0) |
| 245 | hc_mult, hidden_size = 4, 5120 |
| 246 | residual = torch.randn( |
| 247 | num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE |
| 248 | ) |
| 249 | if num_tokens > 1: |
| 250 | residual[0].zero_() |
| 251 | fn = torch.randn(24, hc_mult, hidden_size, device=DEVICE) * 0.02 |
| 252 | x = None |
| 253 | pre_mix = None |
| 254 | if entry == "broadcast": |
| 255 | x = residual[:, 0].contiguous() |
| 256 | residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() |
| 257 | fn = fn.sum(1) |
| 258 | else: |
| 259 | fn = fn.flatten(1) |
| 260 | if entry == "carried": |
| 261 | pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) |
| 262 | if num_tokens > 1: |
| 263 | # A constant collapse just above 1 rounds to BF16 1 before RMSNorm. |
| 264 | residual[1].fill_(1) |
| 265 | pre_mix[1].zero_() |
| 266 | pre_mix[1, 0] = 1.003 |
| 267 | scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE) |
| 268 | base = torch.randn(24, device=DEVICE) |
| 269 | args = (residual, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, sinkhorn_repeat) |
| 270 | expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix, x=x) |
| 271 | norm_kwargs = {} |
| 272 | if fused_norm: |
| 273 | weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE) |
| 274 | weight.uniform_(0.5, 1.5) |
| 275 | norm_kwargs = dict(norm_weight=weight, norm_eps=norm_eps) |
| 276 | # Compare to the original collapse + CUDA RMSNorm, including BF16 rounding. |
| 277 | unfused = mhc_pre_delayed_tilelang(*args, pre_mix=pre_mix, x=x) |
| 278 | normalized = torch.empty_like(unfused[2]) |
| 279 | if num_tokens: |
| 280 | torch.ops._C.rms_norm(normalized, unfused[2], weight, norm_eps) |
| 281 | expected = (*unfused[:2], normalized, unfused[3]) |
| 282 | actual = mhc_pre_delayed_tilelang(*args, pre_mix=pre_mix, x=x, **norm_kwargs) |
| 283 | if fused_norm and num_tokens: |
| 284 | graph = torch.cuda.CUDAGraph() |
| 285 | with torch.cuda.graph(graph): |
| 286 | captured = mhc_pre_delayed_tilelang( |
| 287 | *args, pre_mix=pre_mix, x=x, **norm_kwargs |
| 288 | ) |
| 289 | graph.replay() |
| 290 | for eager, replayed in zip(actual, captured, strict=True): |
| 291 | torch.testing.assert_close(eager, replayed, atol=0, rtol=0) |
| 292 | if fused_norm and entry == "carried" and num_tokens > 1: |
| 293 | torch.testing.assert_close(actual[2][1], expected[2][1], atol=0, rtol=0) |
| 294 | # DeepGEMM uses TF32 for the projection; the TileLang fallback uses FP32. |
| 295 | atol, rtol = (1e-3, 1e-3) if use_deep_gemm else (1e-5, 1e-4) |
| 296 | for i in (0, 1, 3): |
| 297 | torch.testing.assert_close(actual[i], expected[i], atol=atol, rtol=rtol) |
| 298 | if pre_mix is None and not fused_norm: |
| 299 | torch.testing.assert_close(actual[2], expected[2], atol=0, rtol=0) |
| 300 | else: |
| 301 | torch.testing.assert_close(actual[2], expected[2], atol=1.6e-2, rtol=1e-2) |
| 302 | |
| 303 | |
| 304 | @pytest.mark.skipif( |
| 305 | not HAS_TILELANG_MHC or not current_platform.is_device_capability_family(100), |
| 306 | reason="SM100 TileLang mHC required", |
| 307 | ) |
| 308 | @pytest.mark.parametrize("num_tokens", [0, 1, 8, 9, 32, 128, 256]) |
| 309 | @pytest.mark.parametrize("entry", ["broadcast", "identity", "carried"]) |
| 310 | @pytest.mark.parametrize("hidden_size,hc_mult", [(4096, 2), (5120, 4), (7168, 4)]) |
| 311 | def test_deepseek_v41_mhc_overlap_preserves_outputs_on_replay( |
| 312 | num_tokens, entry, hidden_size, hc_mult |
| 313 | ): |
| 314 | """The input is usable before joining; coefficients stay exact on graph replay.""" |
| 315 | from vllm.models.deepseek_v41.nvidia.ops.mhc import mhc_pre_delayed_overlap |
| 316 | from vllm.utils.deep_gemm import is_deep_gemm_supported |
| 317 | |
| 318 | if not is_deep_gemm_supported(): |
| 319 | pytest.skip("DeepGEMM required") |
| 320 | set_random_seed(42) |
| 321 | mix_size = hc_mult * (hc_mult + 2) |
| 322 | residual = torch.randn( |
| 323 | num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE |
| 324 | ) |
| 325 | fn = torch.randn(mix_size, hc_mult, hidden_size, device=DEVICE) * 0.02 |
| 326 | x = None |
| 327 | pre_mix = None |
| 328 | if entry == "broadcast": |
| 329 | x = residual[:, 0].contiguous() |
| 330 | residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() |
| 331 | fn = fn.sum(1) |
| 332 | else: |
| 333 | fn = fn.flatten(1) |
| 334 | if entry == "carried": |
| 335 | pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) |
| 336 | scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE) |
| 337 | base = torch.randn(mix_size, device=DEVICE) |
| 338 | weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE).uniform_( |
| 339 | 0.5, 1.5 |
| 340 | ) |
| 341 | args = (residual, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20) |
| 342 | kwargs = dict(pre_mix=pre_mix, x=x, norm_weight=weight, norm_eps=1e-20) |
| 343 | side = torch.cuda.Stream() |
| 344 | |
| 345 | def run(): |
| 346 | outputs = mhc_pre_delayed_overlap(*args, **kwargs, stream=side) |
| 347 | consumed_input = outputs[2].clone() |
| 348 | torch.cuda.current_stream().wait_stream(side) |
| 349 | return (*outputs, consumed_input) |
| 350 | |
| 351 | def check(actual): |
| 352 | expected = mhc_pre_delayed_tilelang(*args, **kwargs) |
| 353 | for a, b in zip(actual[:4], expected, strict=True): |
| 354 | torch.testing.assert_close(a, b, atol=0, rtol=0) |
| 355 | torch.testing.assert_close(actual[4], expected[2], atol=0, rtol=0) |
| 356 | |
| 357 | check(run()) |
| 358 | if num_tokens == 0: |
| 359 | return |
| 360 | graph = torch.cuda.CUDAGraph() |
| 361 | with torch.cuda.graph(graph): |
| 362 | captured = run() |
| 363 | for _ in range(5): |
| 364 | if x is not None: |
| 365 | x.normal_() |
| 366 | residual.copy_(x.unsqueeze(1).expand_as(residual)) |
| 367 | else: |
| 368 | residual.normal_() |
| 369 | if pre_mix is not None: |
| 370 | pre_mix.uniform_() |
| 371 | graph.replay() |
| 372 | check(captured) |
| 373 | |
| 374 | |
| 375 | def mhc_fused_post_pre_delayed_ref( |
| 376 | x: torch.Tensor, |
| 377 | residual: torch.Tensor, |
| 378 | post_layer_mix: torch.Tensor, |
| 379 | comb_res_mix: torch.Tensor, |
| 380 | *mix_args, |
| 381 | pre_mix: torch.Tensor | None, |
| 382 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 383 | """Project the post result before it is rounded to BF16. |
| 384 | |
| 385 | The fused kernel keeps the updated residual streams in registers, so its |
| 386 | projection sees the FP32 post result while the streams it stores are the |
| 387 | BF16 rounding of it. |
| 388 | """ |
| 389 | post = torch.einsum("tij,tih->tjh", comb_res_mix, residual.float()) |
| 390 | post = post + post_layer_mix * x.float().unsqueeze(-2) |
| 391 | return mhc_pre_delayed_torch( |
| 392 | post.bfloat16(), *mix_args, pre_mix=pre_mix, x=post.flatten(1) |
| 393 | ) |
| 394 | |
| 395 | |
| 396 | @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") |
| 397 | @pytest.mark.parametrize("num_tokens", [0, 1, 8, 128]) |
| 398 | @pytest.mark.parametrize("hidden_size", [4096, 7168]) |
| 399 | @pytest.mark.parametrize("carried", [False, True]) |
| 400 | def test_deepseek_v41_mhc_fused_post_pre_delayed(num_tokens, hidden_size, carried): |
| 401 | """Fold the post block into the delayed pre without changing its outputs.""" |
| 402 | set_random_seed(0) |
| 403 | hc_mult = 4 |
| 404 | mix_size = hc_mult * (hc_mult + 2) |
| 405 | x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) |
| 406 | residual = torch.randn( |
| 407 | num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE |
| 408 | ) |
| 409 | post_layer_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE) |
| 410 | comb_res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE) |
| 411 | fn = torch.randn(mix_size, hc_mult * hidden_size, device=DEVICE) * 0.02 |
| 412 | scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE) |
| 413 | base = torch.randn(mix_size, device=DEVICE) |
| 414 | pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) if carried else None |
| 415 | weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE) |
| 416 | weight.uniform_(0.5, 1.5) |
| 417 | mix_args = (fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20) |
| 418 | |
| 419 | residual_cur, post, comb, layer_input, next_pre, aux = ( |
| 420 | torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( |
| 421 | x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, weight, 1e-6 |
| 422 | ) |
| 423 | ) |
| 424 | assert aux.shape == (0, hidden_size) |
| 425 | assert [ |
| 426 | residual_cur.shape, |
| 427 | post.shape, |
| 428 | comb.shape, |
| 429 | layer_input.shape, |
| 430 | next_pre.shape, |
| 431 | ] == [ |
| 432 | (num_tokens, hc_mult, hidden_size), |
| 433 | (num_tokens, hc_mult, 1), |
| 434 | (num_tokens, hc_mult, hc_mult), |
| 435 | (num_tokens, hidden_size), |
| 436 | (num_tokens, hc_mult), |
| 437 | ] |
| 438 | if num_tokens == 0: |
| 439 | return |
| 440 | |
| 441 | # The post mapping and the collapse are what the unfused kernels produce. |
| 442 | residual_ref = torch.ops.vllm.mhc_post_tilelang( |
| 443 | x, residual, post_layer_mix, comb_res_mix |
| 444 | ) |
| 445 | layer_input_ref = mhc_pre_delayed_tilelang( |
| 446 | residual_ref, *mix_args, pre_mix=pre_mix, norm_weight=weight, norm_eps=1e-6 |
| 447 | )[2] |
| 448 | torch.testing.assert_close(residual_cur, residual_ref, atol=0, rtol=0) |
| 449 | torch.testing.assert_close(layer_input, layer_input_ref, atol=0, rtol=0) |
| 450 | |
| 451 | post_ref, comb_ref, _, next_pre_ref = mhc_fused_post_pre_delayed_ref( |
| 452 | x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix=pre_mix |
| 453 | ) |
| 454 | # Above the fused kernel's token count the projection falls back to a |
| 455 | # split-k TF32 GEMM over the stored BF16 streams. |
| 456 | tol = ( |
| 457 | 1e-5 |
| 458 | if mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) is not None |
| 459 | else 5e-3 |
| 460 | ) |
| 461 | for actual, expected in ( |
| 462 | (post, post_ref), |
| 463 | (comb, comb_ref), |
| 464 | (next_pre, next_pre_ref), |
| 465 | ): |
| 466 | torch.testing.assert_close(actual, expected, atol=tol, rtol=tol) |
| 467 | |
| 468 | # Folding the stream mean into the collapse must not change it, or anything |
| 469 | # else the op returns. |
| 470 | captured = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( |
| 471 | x, |
| 472 | residual, |
| 473 | post_layer_mix, |
| 474 | comb_res_mix, |
| 475 | *mix_args, |
| 476 | pre_mix, |
| 477 | weight, |
| 478 | 1e-6, |
| 479 | True, |
| 480 | ) |
| 481 | torch.testing.assert_close(captured[5], residual_cur.mean(dim=1), atol=0, rtol=0) |
| 482 | unchanged = (residual_cur, post, comb, layer_input, next_pre) |
| 483 | for with_aux, without in zip(captured[:5], unchanged, strict=True): |
| 484 | torch.testing.assert_close(with_aux, without, atol=0, rtol=0) |
| 485 | |
| 486 | # The unnormalized epilogue has its own fused aux path. |
| 487 | bare = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( |
| 488 | x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, None, 1e-6, True |
| 489 | ) |
| 490 | torch.testing.assert_close(bare[5], bare[0].mean(dim=1), atol=0, rtol=0) |
| 491 | |
| 492 | # Decode replays this op from a captured graph. |
| 493 | graph = torch.cuda.CUDAGraph() |
| 494 | with torch.cuda.graph(graph): |
| 495 | captured = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( |
| 496 | x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, weight, 1e-6 |
| 497 | ) |
| 498 | graph.replay() |
| 499 | eager = (residual_cur, post, comb, layer_input, next_pre, aux) |
| 500 | for replayed, expected in zip(captured, eager, strict=True): |
| 501 | torch.testing.assert_close(replayed, expected, atol=0, rtol=0) |
| 502 | |
| 503 | |
| 504 | @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") |
| 505 | @pytest.mark.parametrize("carried", [False, True]) |
| 506 | @pytest.mark.parametrize("capture_aux", [False, True]) |
| 507 | def test_mhc_fused_post_pre_delayed_custom_op_supports_compile(carried, capture_aux): |
| 508 | set_random_seed(0) |
| 509 | x = torch.randn(2, 5120, dtype=torch.bfloat16, device=DEVICE) |
| 510 | residual = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE) |
| 511 | post_layer_mix = torch.rand(2, 4, 1, device=DEVICE) |
| 512 | comb_res_mix = torch.rand(2, 4, 4, device=DEVICE) |
| 513 | fn = torch.randn(24, 20480, device=DEVICE) * 0.02 |
| 514 | pre_mix = torch.rand(2, 4, device=DEVICE) if carried else None |
| 515 | scale = torch.ones(3, device=DEVICE) |
| 516 | base = torch.zeros(24, device=DEVICE) |
| 517 | torch.library.opcheck( |
| 518 | torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang.default, |
| 519 | ( |
| 520 | x, |
| 521 | residual, |
| 522 | post_layer_mix, |
| 523 | comb_res_mix, |
| 524 | fn, |
| 525 | scale, |
| 526 | base, |
| 527 | 1e-20, |
| 528 | 1e-6, |
| 529 | 1e-6, |
| 530 | 2.0, |
| 531 | 20, |
| 532 | pre_mix, |
| 533 | None, |
| 534 | 1e-6, |
| 535 | capture_aux, |
| 536 | ), |
| 537 | ) |
| 538 | |
| 539 | |
| 540 | @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") |
| 541 | @pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"]) |
| 542 | @pytest.mark.parametrize( |
| 543 | "mhc_mode", ["disabled", "eager", "overlap", "piecewise", "large_batch"] |
| 544 | ) |
| 545 | def test_deepseek_v41_decoder_mixes_match_torch( |
| 546 | entry, mhc_mode, monkeypatch, default_vllm_config |
| 547 | ): |
| 548 | """Preserve carried pre-mixes and Engram ordering at every decoder entry.""" |
| 549 | from vllm.models.deepseek_v41.nvidia.ops.mhc import MHC_OVERLAP_MAX_TOKENS |
| 550 | |
| 551 | num_tokens = MHC_OVERLAP_MAX_TOKENS + 1 if mhc_mode == "large_batch" else 3 |
| 552 | set_random_seed(0) |
| 553 | decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer) |
| 554 | nn.Module.__init__(decoder) |
| 555 | decoder.hc_mult = 4 |
| 556 | decoder.hc_sinkhorn_iters = 20 |
| 557 | decoder.hc_eps = 1e-6 |
| 558 | decoder.rms_norm_eps = 1e-20 |
| 559 | decoder.hc_post_alpha = 2.0 |
| 560 | decoder.use_sequence_parallel = False |
| 561 | decoder.fuse_mhc_all_reduce = False |
| 562 | decoder.mhc_stream = None |
| 563 | if mhc_mode != "disabled": |
| 564 | from vllm.utils.deep_gemm import is_deep_gemm_supported |
| 565 | |
| 566 | if ( |
| 567 | not current_platform.is_device_capability_family(100) |
| 568 | or not is_deep_gemm_supported() |
| 569 | ): |
| 570 | pytest.skip("SM100 DeepGEMM required for overlap") |
| 571 | decoder.mhc_stream = torch.cuda.Stream() |
| 572 | if mhc_mode in ("eager", "piecewise", "large_batch"): |
| 573 | |
| 574 | def unexpected_overlap(*args, **kwargs): |
| 575 | pytest.fail("unsupported execution must retain native mHC") |
| 576 | |
| 577 | monkeypatch.setattr( |
| 578 | "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_overlap", |
| 579 | unexpected_overlap, |
| 580 | ) |
| 581 | monkeypatch.setattr( |
| 582 | "vllm.models.deepseek_v41.nvidia.ops.mhc.mhc_pre_delayed_overlap", |
| 583 | unexpected_overlap, |
| 584 | ) |
| 585 | decoder.engram = None |
| 586 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 587 | |
| 588 | decoder.attn_norm = decoder.ffn_norm = RMSNorm(5120, 1e-6).to( |
| 589 | device=DEVICE, dtype=torch.bfloat16 |
| 590 | ) |
| 591 | decoder.attn = lambda positions, x, _: x * 0.5 |
| 592 | decoder.ffn = lambda x, input_ids, mega_gate_metadata=None: x * 0.25 |
| 593 | with torch.device(DEVICE): |
| 594 | decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02 |
| 595 | decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02 |
| 596 | decoder.hc_attn_fn_broadcast = decoder.hc_attn_fn.view(24, 4, 5120).sum(1) |
| 597 | decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3) |
| 598 | decoder.hc_attn_base = torch.randn(24) |
| 599 | decoder.hc_ffn_base = torch.randn(24) |
| 600 | x = torch.randn(num_tokens, 5120, dtype=torch.bfloat16) |
| 601 | positions = torch.arange(num_tokens) |
| 602 | kwargs = {} |
| 603 | if entry != "broadcast": |
| 604 | kwargs["pre_mix"] = torch.rand(num_tokens, 4) |
| 605 | if entry == "pipeline": |
| 606 | x = torch.randn(num_tokens, 4, 5120, dtype=torch.bfloat16) |
| 607 | if entry in ("residual", "engram"): |
| 608 | kwargs.update( |
| 609 | residual=torch.randn(num_tokens, 4, 5120, dtype=torch.bfloat16), |
| 610 | post_mix=torch.rand(num_tokens, 4, 1), |
| 611 | res_mix=torch.rand(num_tokens, 4, 4), |
| 612 | ) |
| 613 | if entry == "engram": |
| 614 | |
| 615 | class FakeEngram(nn.Module): |
| 616 | layer_hash_index = 0 |
| 617 | |
| 618 | def forward(self, residual, hashes, mask): |
| 619 | return residual + 0.125 |
| 620 | |
| 621 | decoder.engram = FakeEngram() |
| 622 | kwargs["engram_hashes"] = torch.zeros(num_tokens, 1, 1, dtype=torch.int32) |
| 623 | |
| 624 | from vllm.config import CUDAGraphMode |
| 625 | from vllm.forward_context import set_forward_context |
| 626 | |
| 627 | mode = CUDAGraphMode.PIECEWISE if mhc_mode == "piecewise" else CUDAGraphMode.NONE |
| 628 | overlap_calls = [] |
| 629 | if mhc_mode == "overlap": |
| 630 | from vllm.models.deepseek_v41.nvidia.ops.mhc import mhc_pre_delayed_overlap |
| 631 | |
| 632 | def checked_overlap(*args, **kwargs): |
| 633 | overlap_calls.append(torch.cuda.is_current_stream_capturing()) |
| 634 | return mhc_pre_delayed_overlap(*args, **kwargs) |
| 635 | |
| 636 | for module in ("model", "ops.mhc"): |
| 637 | monkeypatch.setattr( |
| 638 | f"vllm.models.deepseek_v41.nvidia.{module}.mhc_pre_delayed_overlap", |
| 639 | checked_overlap, |
| 640 | ) |
| 641 | with set_forward_context(None, default_vllm_config, cudagraph_runtime_mode=mode): |
| 642 | actual = decoder(x, positions, None, **kwargs) |
| 643 | if mhc_mode == "overlap": |
| 644 | graph = torch.cuda.CUDAGraph() |
| 645 | with torch.cuda.graph(graph): |
| 646 | actual = decoder(x, positions, None, **kwargs) |
| 647 | graph.replay() |
| 648 | assert overlap_calls and all(overlap_calls) |
| 649 | |
| 650 | def reference(*args, norm_weight, norm_eps, **kwargs): |
| 651 | post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs) |
| 652 | return post, res, decoder.attn_norm(collapsed), pre |
| 653 | |
| 654 | def fused_reference( |
| 655 | x, |
| 656 | residual, |
| 657 | post_mix, |
| 658 | res_mix, |
| 659 | *args, |
| 660 | capture_aux=False, |
| 661 | stream=None, |
| 662 | reduce_results=False, |
| 663 | **kw, |
| 664 | ): |
| 665 | residual = mhc_post_torch(x, residual, post_mix, res_mix) |
| 666 | aux = residual.mean(dim=1) if capture_aux else residual.new_empty(0) |
| 667 | return residual, *reference(residual, *args, **kw), aux |
| 668 | |
| 669 | monkeypatch.setattr( |
| 670 | "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_tilelang", |
| 671 | reference, |
| 672 | ) |
| 673 | decoder.mhc_stream = None |
| 674 | monkeypatch.setattr( |
| 675 | "vllm.models.deepseek_v41.nvidia.model.mhc_shifted_post_pre", |
| 676 | fused_reference, |
| 677 | ) |
| 678 | expected = decoder(x, positions, None, **kwargs) |
| 679 | assert actual[-1] is None and expected[-1] is None |
| 680 | for result, ref in zip(actual[:-1], expected[:-1], strict=True): |
| 681 | torch.testing.assert_close(result, ref, atol=2e-2, rtol=1e-2) |
| 682 | |
| 683 | |
| 684 | @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") |
| 685 | @pytest.mark.parametrize("entry", ["fused", "engram"]) |
| 686 | def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_config): |
| 687 | """Read the aux hidden state back out of the seam that already computed it. |
| 688 | |
| 689 | The draft model's input is the mean over hc streams of the previous |
| 690 | layer's post. That post now runs inside this layer's fused pre, so the |
| 691 | captured value must equal a standalone ``mhc_post_tilelang`` on the same |
| 692 | inputs -- and Engram must not be folded into it, since the injection |
| 693 | happens after the aux consumers read the stream. |
| 694 | """ |
| 695 | set_random_seed(0) |
| 696 | decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer) |
| 697 | nn.Module.__init__(decoder) |
| 698 | decoder.hc_mult = 4 |
| 699 | decoder.hc_sinkhorn_iters = 20 |
| 700 | decoder.hc_eps = 1e-6 |
| 701 | decoder.rms_norm_eps = 1e-6 |
| 702 | decoder.hc_post_alpha = 2.0 |
| 703 | decoder.use_sequence_parallel = False |
| 704 | decoder.fuse_mhc_all_reduce = False |
| 705 | decoder.mhc_stream = None |
| 706 | decoder.engram = None |
| 707 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 708 | |
| 709 | decoder.attn_norm = decoder.ffn_norm = RMSNorm(5120, 1e-6).to( |
| 710 | device=DEVICE, dtype=torch.bfloat16 |
| 711 | ) |
| 712 | decoder.attn = lambda positions, x, _: x * 0.5 |
| 713 | decoder.ffn = lambda x, input_ids, mega_gate_metadata=None: x * 0.25 |
| 714 | with torch.device(DEVICE): |
| 715 | decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02 |
| 716 | decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02 |
| 717 | decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3) |
| 718 | decoder.hc_attn_base = torch.randn(24) |
| 719 | decoder.hc_ffn_base = torch.randn(24) |
| 720 | x = torch.randn(3, 5120, dtype=torch.bfloat16) |
| 721 | kwargs = dict( |
| 722 | pre_mix=torch.rand(3, 4), |
| 723 | residual=torch.randn(3, 4, 5120, dtype=torch.bfloat16), |
| 724 | post_mix=torch.rand(3, 4, 1), |
| 725 | res_mix=torch.rand(3, 4, 4), |
| 726 | ) |
| 727 | if entry == "engram": |
| 728 | |
| 729 | class FakeEngram(nn.Module): |
| 730 | layer_hash_index = 0 |
| 731 | |
| 732 | def forward(self, residual, hashes, mask): |
| 733 | return residual + 0.125 |
| 734 | |
| 735 | decoder.engram = FakeEngram() |
| 736 | kwargs["engram_hashes"] = torch.zeros(3, 1, 1, dtype=torch.int32) |
| 737 | positions = torch.arange(3) |
| 738 | |
| 739 | expected = torch.ops.vllm.mhc_post_tilelang( |
| 740 | x, kwargs["residual"], kwargs["post_mix"], kwargs["res_mix"] |
| 741 | ).mean(dim=1) |
| 742 | previous_aux = decoder(x, positions, None, **kwargs, capture_previous_aux=True)[-1] |
| 743 | torch.testing.assert_close(previous_aux, expected, atol=0, rtol=0) |
| 744 | assert decoder(x, positions, None, **kwargs)[-1] is None |
| 745 | |
| 746 | |
| 747 | @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") |
| 748 | @pytest.mark.parametrize("carried", [False, True]) |
| 749 | def test_mhc_pre_delayed_custom_op_supports_compile(carried): |
| 750 | set_random_seed(0) |
| 751 | x = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE) |
| 752 | fn = torch.randn(24, 20480, device=DEVICE) * 0.02 |
| 753 | pre_mix = torch.rand(2, 4, device=DEVICE) if carried else None |
| 754 | scale = torch.ones(3, device=DEVICE) |
| 755 | base = torch.zeros(24, device=DEVICE) |
| 756 | torch.library.opcheck( |
| 757 | torch.ops.vllm.mhc_pre_delayed_tilelang.default, |
| 758 | (x, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20, pre_mix), |
| 759 | ) |
| 760 | |
| 761 | |
| 762 | @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") |
| 763 | @pytest.mark.skipif( |
| 764 | not current_platform.is_device_capability_family(100), |
| 765 | reason="DeepGEMM Mega mHC requires SM100-family CUDA", |
| 766 | ) |
| 767 | @pytest.mark.parametrize("num_tokens", [0, 1, 17, 128, 1024]) |
| 768 | @pytest.mark.parametrize("hidden_size", [5120, 7168]) |
| 769 | def test_deep_gemm_mega_mhc_correctness(num_tokens, hidden_size): |
| 770 | set_random_seed(0) |
| 771 | hc_mult = 4 |
| 772 | if not is_mega_mhc_supported(hidden_size, 4): |
| 773 | pytest.skip("DeepGEMM Mega mHC is not available") |
| 774 | |
| 775 | x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) |
| 776 | residual = torch.randn( |
| 777 | num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE |
| 778 | ) |
| 779 | previous_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) |
| 780 | post_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE) |
| 781 | res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE) |
| 782 | fn = torch.randn(24, hc_mult * hidden_size, device=DEVICE) * 0.01 |
| 783 | scale = torch.randn(3, device=DEVICE) * 0.1 |
| 784 | base = torch.randn(24, device=DEVICE) * 0.1 |
| 785 | norm_weight = torch.empty( |
| 786 | hidden_size, dtype=torch.bfloat16, device=DEVICE |
| 787 | ).uniform_(0.9, 1.1) |
| 788 | |
| 789 | expected_residual = ( |
| 790 | mhc_post_tilelang(x, residual, post_mix, res_mix) if num_tokens else residual |
| 791 | ) |
| 792 | ( |
| 793 | expected_post_mix, |
| 794 | expected_res_mix, |
| 795 | expected_y_bf16, |
| 796 | expected_previous_mix, |
| 797 | ) = mhc_pre_delayed_tilelang( |
| 798 | expected_residual, |
| 799 | fn, |
| 800 | scale, |
| 801 | base, |
| 802 | 2e-5, |
| 803 | 3e-4, |
| 804 | 2e-6, |
| 805 | 1.25, |
| 806 | 10, |
| 807 | pre_mix=previous_mix, |
| 808 | norm_weight=norm_weight, |
| 809 | norm_eps=7e-6, |
| 810 | ) |
| 811 | args = ( |
| 812 | x, |
| 813 | residual, |
| 814 | previous_mix, |
| 815 | post_mix, |
| 816 | res_mix, |
| 817 | fn, |
| 818 | scale, |
| 819 | base, |
| 820 | 2e-5, |
| 821 | 3e-4, |
| 822 | 1.25, |
| 823 | 2e-6, |
| 824 | 10, |
| 825 | norm_weight, |
| 826 | 7e-6, |
| 827 | ) |
| 828 | actual = mhc_shifted_post_pre_deep_gemm(*args) |
| 829 | expected = ( |
| 830 | expected_residual, |
| 831 | expected_post_mix, |
| 832 | expected_res_mix, |
| 833 | expected_y_bf16, |
| 834 | expected_previous_mix, |
| 835 | ) |
| 836 | for i, (result, ref) in enumerate(zip(actual, expected, strict=True)): |
| 837 | # Match the existing mHC BF16 tolerance: post rounding differences |
| 838 | # can be amplified by cancellation in the carried collapse. |
| 839 | atol, rtol = (1.6e-2, 1e-2) if i in (0, 3) else (1e-6, 1e-3) |
| 840 | torch.testing.assert_close(result, ref, atol=atol, rtol=rtol) |
| 841 | |
| 842 | if num_tokens: |
| 843 | stream = torch.cuda.Stream() |
| 844 | stream.wait_stream(torch.cuda.current_stream()) |
| 845 | with torch.cuda.stream(stream): |
| 846 | # DeepGEMM initializes barriers per stream before capture. |
| 847 | mhc_shifted_post_pre_deep_gemm(*args) |
| 848 | graph = torch.cuda.CUDAGraph() |
| 849 | with torch.cuda.graph(graph, stream=stream): |
| 850 | captured = mhc_shifted_post_pre_deep_gemm(*args) |
| 851 | for multiplier in (0.5, -1.0): |
| 852 | x.mul_(multiplier) |
| 853 | graph.replay() |
| 854 | eager = mhc_shifted_post_pre_deep_gemm(*args) |
| 855 | for result, ref in zip(captured, eager, strict=True): |
| 856 | torch.testing.assert_close(result, ref, atol=0, rtol=0) |
| 857 | torch.cuda.current_stream().wait_stream(stream) |
| 858 | |
| 859 | |
| 860 | def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor: |
| 861 | x = x.softmax(-1) + eps |
| 862 | x = x / (x.sum(-2, keepdim=True) + eps) |
| 863 | for _ in range(repeat - 1): |
| 864 | x = x / (x.sum(-1, keepdim=True) + eps) |
| 865 | x = x / (x.sum(-2, keepdim=True) + eps) |
| 866 | return x |
| 867 | |
| 868 | |
| 869 | def mhc_pre_ref( |
| 870 | residual: torch.Tensor, |
| 871 | fn: torch.Tensor, |
| 872 | hc_scale: torch.Tensor, |
| 873 | hc_base: torch.Tensor, |
| 874 | rms_eps: float, |
| 875 | hc_pre_eps: float, |
| 876 | hc_sinkhorn_eps: float, |
| 877 | hc_post_mult_value: float, |
| 878 | sinkhorn_repeat: int, |
| 879 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 880 | """MHC pre reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_pre.py#L303.""" |
| 881 | hc_mult = residual.shape[-2] |
| 882 | |
| 883 | residual_flat = residual.flatten(-2, -1).float() |
| 884 | sqrsum = residual_flat.square().sum(-1) |
| 885 | mixes = ( |
| 886 | residual_flat @ fn.T * (sqrsum.unsqueeze(-1) / fn.shape[-1] + rms_eps).rsqrt() |
| 887 | ) |
| 888 | |
| 889 | hc_scale = torch.cat( |
| 890 | [ |
| 891 | hc_scale[0].expand(hc_mult), |
| 892 | hc_scale[1].expand(hc_mult), |
| 893 | hc_scale[2].expand(hc_mult * hc_mult), |
| 894 | ], |
| 895 | ) |
| 896 | mixes = mixes * hc_scale + hc_base |
| 897 | |
| 898 | pre_mix = mixes[:, :hc_mult].sigmoid().unsqueeze(-1) + hc_pre_eps |
| 899 | post_mix = ( |
| 900 | mixes[:, hc_mult : 2 * hc_mult].sigmoid() * hc_post_mult_value |
| 901 | ).unsqueeze(-1) |
| 902 | res_mix = mixes[:, 2 * hc_mult :].view(-1, hc_mult, hc_mult) |
| 903 | |
| 904 | res_mix = sinkhorn_normalize_ref( |
| 905 | res_mix, repeat=sinkhorn_repeat, eps=hc_sinkhorn_eps |
| 906 | ) |
| 907 | |
| 908 | layer_input = (residual * pre_mix).sum(-2).bfloat16() |
| 909 | |
| 910 | return post_mix, res_mix, layer_input |
| 911 | |
| 912 | |
| 913 | def mhc_post_ref( |
| 914 | x: torch.Tensor, |
| 915 | residual: torch.Tensor, |
| 916 | post_layer_mix: torch.Tensor, |
| 917 | comb_res_mix: torch.Tensor, |
| 918 | ) -> torch.Tensor: |
| 919 | """MHC post reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_post.py#L68.""" |
| 920 | term2 = torch.bmm(comb_res_mix.mT, residual.float()) |
| 921 | return (x.float().unsqueeze(-2) * post_layer_mix + term2).bfloat16() |
| 922 | |
| 923 | |
| 924 | def hc_head_ref( |
| 925 | residual: torch.Tensor, |
| 926 | fn: torch.Tensor, |
| 927 | hc_scale: torch.Tensor, |
| 928 | hc_base: torch.Tensor, |
| 929 | rms_eps: float, |
| 930 | hc_eps: float, |
| 931 | ) -> torch.Tensor: |
| 932 | residual_flat = residual.flatten(-2).float() |
| 933 | residual_norm = residual_flat * torch.rsqrt( |
| 934 | residual_flat.square().mean(dim=-1, keepdim=True) + rms_eps |
| 935 | ) |
| 936 | pre_mix = torch.nn.functional.linear(residual_norm, fn) |
| 937 | pre_mix = torch.sigmoid(pre_mix * hc_scale + hc_base) + hc_eps |
| 938 | return torch.sum(pre_mix.unsqueeze(-1) * residual.float(), dim=-2).bfloat16() |
| 939 | |
| 940 | |
| 941 | @pytest.mark.skipif( |
| 942 | not HAS_TILELANG_MHC, |
| 943 | reason="TileLang MHC support required", |
| 944 | ) |
| 945 | @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) |
| 946 | @pytest.mark.parametrize("hidden_size", [4096, 7168]) |
| 947 | @pytest.mark.parametrize("hc_mult", [4]) |
| 948 | @pytest.mark.parametrize("fused_norm", [False, True]) |
| 949 | def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult, fused_norm): |
| 950 | torch.set_default_device(DEVICE) |
| 951 | set_random_seed(0) |
| 952 | |
| 953 | residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) |
| 954 | hc_mult2 = hc_mult * hc_mult |
| 955 | hc_mult3 = 2 * hc_mult + hc_mult2 |
| 956 | fn = ( |
| 957 | torch.randn((hc_mult3, hc_mult, hidden_size), dtype=torch.float) |
| 958 | * 1e-4 |
| 959 | * (1 + torch.arange(hc_mult).mul(0.01).view(1, -1, 1)) |
| 960 | ).flatten(1, 2) |
| 961 | hc_scale = torch.randn((3,), dtype=torch.float) * 0.1 |
| 962 | hc_base = torch.randn((hc_mult3,), dtype=torch.float) * 0.1 |
| 963 | |
| 964 | hc_sinkhorn_eps = hc_pre_eps = rms_eps = 1e-6 |
| 965 | sinkhorn_repeat = 20 |
| 966 | hc_post_alpha = 1.0 |
| 967 | |
| 968 | ref = mhc_pre_ref( |
| 969 | residual, |
| 970 | fn, |
| 971 | hc_scale, |
| 972 | hc_base, |
| 973 | rms_eps, |
| 974 | hc_pre_eps, |
| 975 | hc_sinkhorn_eps, |
| 976 | hc_post_alpha, |
| 977 | sinkhorn_repeat, |
| 978 | ) |
| 979 | norm_kwargs = {} |
| 980 | if fused_norm: |
| 981 | # Non-shifted mHC must still use its freshly computed pre-mix. |
| 982 | weight = torch.ones(hidden_size, dtype=torch.bfloat16) |
| 983 | norm_kwargs = dict(norm_weight=weight, norm_eps=rms_eps) |
| 984 | ref = (*ref[:2], F.rms_norm(ref[2], (hidden_size,), weight, rms_eps)) |
| 985 | out = torch.ops.vllm.mhc_pre_tilelang( |
| 986 | residual, |
| 987 | fn, |
| 988 | hc_scale, |
| 989 | hc_base, |
| 990 | rms_eps, |
| 991 | hc_pre_eps, |
| 992 | hc_sinkhorn_eps, |
| 993 | hc_post_alpha, |
| 994 | sinkhorn_repeat, |
| 995 | **norm_kwargs, |
| 996 | ) |
| 997 | |
| 998 | for actual, expected in zip(out, ref, strict=True): |
| 999 | torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2) |
| 1000 | |
| 1001 | |
| 1002 | @pytest.mark.skipif( |
| 1003 | not HAS_TILELANG_MHC, |
| 1004 | reason="TileLang MHC support required", |
| 1005 | ) |
| 1006 | @pytest.mark.parametrize( |
| 1007 | ("num_tokens", "hidden_size"), |
| 1008 | [ |
| 1009 | (1, 1280), |
| 1010 | (512, 1280), |
| 1011 | (2048, 1280), |
| 1012 | (1, 4096), |
| 1013 | (64, 4096), |
| 1014 | (512, 4096), |
| 1015 | (2048, 4096), |
| 1016 | (1, 7168), |
| 1017 | (64, 7168), |
| 1018 | (512, 7168), |
| 1019 | (2048, 7168), |
| 1020 | ], |
| 1021 | ) |
| 1022 | def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size): |
| 1023 | torch.set_default_device(DEVICE) |
| 1024 | set_random_seed(0) |
| 1025 | |
| 1026 | hc_mult = 4 |
| 1027 | hc_mult3 = 2 * hc_mult + hc_mult * hc_mult |
| 1028 | x = torch.randn((num_tokens, hc_mult * hidden_size), dtype=torch.bfloat16) |
| 1029 | fn = torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4 |
| 1030 | out_ref = torch.empty((1, num_tokens, hc_mult3), dtype=torch.float32) |
| 1031 | sqrsum_ref = torch.empty((1, num_tokens), dtype=torch.float32) |
| 1032 | out = torch.empty_like(out_ref) |
| 1033 | sqrsum = torch.empty_like(sqrsum_ref) |
| 1034 | |
| 1035 | _torch_hc_prenorm_gemm(x, fn, out_ref, sqrsum_ref) |
| 1036 | _HC_PRENORM_GEMM_TILELANG_KERNEL( |
| 1037 | x, |
| 1038 | fn, |
| 1039 | out, |
| 1040 | sqrsum, |
| 1041 | hidden_size, |
| 1042 | hc_mult, |
| 1043 | ) |
| 1044 | |
| 1045 | torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4) |
| 1046 | torch.testing.assert_close(sqrsum, sqrsum_ref, atol=1e-2, rtol=1e-6) |
| 1047 | |
| 1048 | |
| 1049 | @pytest.mark.skipif( |
| 1050 | not HAS_TILELANG_MHC, |
| 1051 | reason="TileLang MHC support required", |
| 1052 | ) |
| 1053 | @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) |
| 1054 | @pytest.mark.parametrize("hidden_size", [4096, 7168]) |
| 1055 | @pytest.mark.parametrize("hc_mult", [4]) |
| 1056 | def test_mhc_post_tilelang(num_tokens, hidden_size, hc_mult): |
| 1057 | torch.set_default_device(DEVICE) |
| 1058 | set_random_seed(0) |
| 1059 | |
| 1060 | x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16) |
| 1061 | residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) |
| 1062 | post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32) |
| 1063 | comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32) |
| 1064 | |
| 1065 | ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix) |
| 1066 | out = torch.ops.vllm.mhc_post_tilelang( |
| 1067 | x, |
| 1068 | residual, |
| 1069 | post_layer_mix, |
| 1070 | comb_res_mix, |
| 1071 | ) |
| 1072 | |
| 1073 | torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2) |
| 1074 | |
| 1075 | |
| 1076 | @pytest.mark.skipif( |
| 1077 | not HAS_TILELANG_MHC, |
| 1078 | reason="TileLang MHC support required", |
| 1079 | ) |
| 1080 | @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) |
| 1081 | @pytest.mark.parametrize("hidden_size", [4096, 7168]) |
| 1082 | @pytest.mark.parametrize("hc_mult", [4]) |
| 1083 | def test_mhc_fused_post_pre(num_tokens, hidden_size, hc_mult): |
| 1084 | torch.set_default_device(DEVICE) |
| 1085 | set_random_seed(0) |
| 1086 | |
| 1087 | x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16) |
| 1088 | residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) |
| 1089 | post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32) |
| 1090 | comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32) |
| 1091 | |
| 1092 | hc_mult2 = hc_mult * hc_mult |
| 1093 | hc_mult3 = hc_mult * 2 + hc_mult2 |
| 1094 | fn = ( |
| 1095 | torch.randn((hc_mult3, hc_mult, hidden_size), dtype=torch.float) |
| 1096 | * 1e-4 |
| 1097 | * (1 + torch.arange(hc_mult).mul(0.01).view(1, -1, 1)) |
| 1098 | ).flatten(1, 2) |
| 1099 | hc_scale = torch.randn((3,), dtype=torch.float) * 0.1 |
| 1100 | hc_base = torch.randn((hc_mult3,), dtype=torch.float) * 0.1 |
| 1101 | |
| 1102 | hc_sinkhorn_eps = hc_pre_eps = rms_eps = 1e-6 |
| 1103 | sinkhorn_repeat = 20 |
| 1104 | hc_post_alpha = 1.0 |
| 1105 | |
| 1106 | def run_ref(): |
| 1107 | residual_ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix) |
| 1108 | post_mix_ref, res_mix_ref, layer_input_ref = mhc_pre_ref( |
| 1109 | residual_ref, |
| 1110 | fn, |
| 1111 | hc_scale, |
| 1112 | hc_base, |
| 1113 | rms_eps, |
| 1114 | hc_pre_eps, |
| 1115 | hc_sinkhorn_eps, |
| 1116 | hc_post_alpha, |
| 1117 | sinkhorn_repeat, |
| 1118 | ) |
| 1119 | return residual_ref, post_mix_ref, res_mix_ref, layer_input_ref |
| 1120 | |
| 1121 | residual_ref, post_mix_ref, res_mix_ref, layer_input_ref = run_ref() |
| 1122 | |
| 1123 | residual, post_mix, res_mix, x = torch.ops.vllm.mhc_fused_post_pre_tilelang( |
| 1124 | x, |
| 1125 | residual, |
| 1126 | post_layer_mix, |
| 1127 | comb_res_mix, |
| 1128 | fn, |
| 1129 | hc_scale, |
| 1130 | hc_base, |
| 1131 | rms_eps, |
| 1132 | hc_pre_eps, |
| 1133 | hc_sinkhorn_eps, |
| 1134 | hc_post_alpha, |
| 1135 | sinkhorn_repeat, |
| 1136 | ) |
| 1137 | |
| 1138 | torch.testing.assert_close(residual, residual_ref, atol=1e-2, rtol=1e-2) |
| 1139 | torch.testing.assert_close(post_mix, post_mix_ref, atol=1e-2, rtol=1e-2) |
| 1140 | torch.testing.assert_close(res_mix, res_mix_ref, atol=1e-2, rtol=1e-2) |
| 1141 | torch.testing.assert_close(x, layer_input_ref, atol=1e-2, rtol=1e-2) |
| 1142 | |
| 1143 | |
| 1144 | def _rocm_mhc_inputs(num_tokens=2, hidden_size=256, hc_mult=4): |
| 1145 | residual = torch.randn( |
| 1146 | (num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16, device=DEVICE |
| 1147 | ) |
| 1148 | hc_mult3 = 2 * hc_mult + hc_mult * hc_mult |
| 1149 | fn = ( |
| 1150 | torch.randn( |
| 1151 | (hc_mult3, hc_mult * hidden_size), dtype=torch.float32, device=DEVICE |
| 1152 | ) |
| 1153 | * 1e-4 |
| 1154 | ) |
| 1155 | hc_scale = torch.randn((3,), dtype=torch.float32, device=DEVICE) * 0.1 |
| 1156 | hc_base = torch.randn((hc_mult3,), dtype=torch.float32, device=DEVICE) * 0.1 |
| 1157 | norm_weight = torch.randn(hidden_size, dtype=torch.bfloat16, device=DEVICE) |
| 1158 | return residual, fn, hc_scale, hc_base, norm_weight |
| 1159 | |
| 1160 | |
| 1161 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required") |
| 1162 | def test_mhc_pre_rocm_fallback_applies_norm(monkeypatch): |
| 1163 | set_random_seed(0) |
| 1164 | residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs() |
| 1165 | rms_eps = hc_pre_eps = hc_sinkhorn_eps = norm_eps = 1e-6 |
| 1166 | sinkhorn_repeat = 20 |
| 1167 | hc_post_alpha = 1.0 |
| 1168 | ref = mhc_pre_ref( |
| 1169 | residual, |
| 1170 | fn, |
| 1171 | hc_scale, |
| 1172 | hc_base, |
| 1173 | rms_eps, |
| 1174 | hc_pre_eps, |
| 1175 | hc_sinkhorn_eps, |
| 1176 | hc_post_alpha, |
| 1177 | sinkhorn_repeat, |
| 1178 | ) |
| 1179 | expected_layer_input = F.rms_norm( |
| 1180 | ref[2], (ref[2].shape[-1],), norm_weight, norm_eps |
| 1181 | ) |
| 1182 | monkeypatch.setattr(mhc_layers, "HAS_AITER_MHC", True) |
| 1183 | monkeypatch.setattr(mhc_layers, "HAS_AITER_MHC_PRE_NORM", False) |
| 1184 | monkeypatch.setattr(mhc_layers, "HAS_TILELANG_MHC", False) |
| 1185 | |
| 1186 | out = object.__new__(MHCPreOp).forward_hip( |
| 1187 | residual, |
| 1188 | fn, |
| 1189 | hc_scale, |
| 1190 | hc_base, |
| 1191 | rms_eps, |
| 1192 | hc_pre_eps, |
| 1193 | hc_sinkhorn_eps, |
| 1194 | hc_post_alpha, |
| 1195 | sinkhorn_repeat, |
| 1196 | norm_weight=norm_weight, |
| 1197 | norm_eps=norm_eps, |
| 1198 | ) |
| 1199 | |
| 1200 | torch.testing.assert_close(out[0], ref[0]) |
| 1201 | torch.testing.assert_close(out[1], ref[1]) |
| 1202 | torch.testing.assert_close(out[2], expected_layer_input) |
| 1203 | |
| 1204 | |
| 1205 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required") |
| 1206 | def test_mhc_fused_rocm_fallback_applies_norm(monkeypatch): |
| 1207 | set_random_seed(0) |
| 1208 | residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs() |
| 1209 | x = torch.randn((2, 256), dtype=torch.bfloat16, device=DEVICE) |
| 1210 | post_layer_mix = torch.randn((2, 4, 1), dtype=torch.float32, device=DEVICE) |
| 1211 | comb_res_mix = torch.randn((2, 4, 4), dtype=torch.float32, device=DEVICE) |
| 1212 | rms_eps = hc_pre_eps = hc_sinkhorn_eps = norm_eps = 1e-6 |
| 1213 | sinkhorn_repeat = 20 |
| 1214 | hc_post_alpha = 1.0 |
| 1215 | residual_ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix) |
| 1216 | pre_ref = mhc_pre_ref( |
| 1217 | residual_ref, |
| 1218 | fn, |
| 1219 | hc_scale, |
| 1220 | hc_base, |
| 1221 | rms_eps, |
| 1222 | hc_pre_eps, |
| 1223 | hc_sinkhorn_eps, |
| 1224 | hc_post_alpha, |
| 1225 | sinkhorn_repeat, |
| 1226 | ) |
| 1227 | expected_layer_input = F.rms_norm( |
| 1228 | pre_ref[2], (pre_ref[2].shape[-1],), norm_weight, norm_eps |
| 1229 | ) |
| 1230 | monkeypatch.setattr(mhc_layers, "HAS_AITER_MHC_FUSED", False) |
| 1231 | monkeypatch.setattr(mhc_layers, "HAS_TILELANG_MHC", False) |
| 1232 | |
| 1233 | out = object.__new__(MHCFusedPostPreOp).forward_hip( |
| 1234 | x, |
| 1235 | residual, |
| 1236 | post_layer_mix, |
| 1237 | comb_res_mix, |
| 1238 | fn, |
| 1239 | hc_scale, |
| 1240 | hc_base, |
| 1241 | rms_eps, |
| 1242 | hc_pre_eps, |
| 1243 | hc_sinkhorn_eps, |
| 1244 | hc_post_alpha, |
| 1245 | sinkhorn_repeat, |
| 1246 | norm_weight=norm_weight, |
| 1247 | norm_eps=norm_eps, |
| 1248 | ) |
| 1249 | |
| 1250 | torch.testing.assert_close(out[0], residual_ref) |
| 1251 | torch.testing.assert_close(out[1], pre_ref[0]) |
| 1252 | torch.testing.assert_close(out[2], pre_ref[1]) |
| 1253 | torch.testing.assert_close(out[3], expected_layer_input) |
| 1254 | |
| 1255 | |
| 1256 | @pytest.mark.skipif( |
| 1257 | not (current_platform.is_rocm() and HAS_AITER_MHC), |
| 1258 | reason="AITER mHC required", |
| 1259 | ) |
| 1260 | @pytest.mark.parametrize("num_tokens", [1, 2, 7, 128, 1024]) |
| 1261 | @pytest.mark.parametrize("carried", [False, True]) |
| 1262 | def test_mhc_pre_delayed_rocm_aiter(num_tokens, carried): |
| 1263 | """AITER must reproduce the delayed reference on both seam variants. |
| 1264 | |
| 1265 | ``num_tokens`` spans the split-k choices AITER makes for the projection, |
| 1266 | since the pre-mix is recovered from that unreduced output. |
| 1267 | """ |
| 1268 | set_random_seed(0) |
| 1269 | hc_mult, hidden_size = 4, 5120 |
| 1270 | residual, fn, hc_scale, hc_base, _ = _rocm_mhc_inputs( |
| 1271 | num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult |
| 1272 | ) |
| 1273 | pre_mix = ( |
| 1274 | torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5 |
| 1275 | if carried |
| 1276 | else None |
| 1277 | ) |
| 1278 | rms_eps = hc_pre_eps = hc_sinkhorn_eps = 1e-6 |
| 1279 | args = ( |
| 1280 | residual, |
| 1281 | fn, |
| 1282 | hc_scale, |
| 1283 | hc_base, |
| 1284 | rms_eps, |
| 1285 | hc_pre_eps, |
| 1286 | hc_sinkhorn_eps, |
| 1287 | 1.0, |
| 1288 | 20, |
| 1289 | ) |
| 1290 | |
| 1291 | expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix) |
| 1292 | residual_out, *actual = object.__new__(MHCPreDelayedOp).forward_hip( |
| 1293 | *args, pre_mix=pre_mix |
| 1294 | ) |
| 1295 | |
| 1296 | # No post was requested, so the residual comes straight back. |
| 1297 | assert residual_out is residual |
| 1298 | for i in (0, 1, 3): |
| 1299 | torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3) |
| 1300 | # The collapse is the same FP32 multiply-and-sum in both paths. |
| 1301 | torch.testing.assert_close(actual[2], expected[2], atol=0, rtol=0) |
| 1302 | |
| 1303 | |
| 1304 | @pytest.mark.skipif( |
| 1305 | not (current_platform.is_rocm() and HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM), |
| 1306 | reason="AITER fused delayed-seam kernel required (gfx950)", |
| 1307 | ) |
| 1308 | @pytest.mark.parametrize("num_tokens", [1, 2, 7, 128, 1024]) |
| 1309 | @pytest.mark.parametrize("seam", ["post", "no_post", "identity_pre"]) |
| 1310 | def test_mhc_pre_delayed_rocm_aiter_fused_rms_norm(num_tokens, seam): |
| 1311 | """One fused kernel must match the whole reference seam: the post block |
| 1312 | (the previous seam's post/comb gates fold ``sublayer_out`` into the |
| 1313 | streams), the delayed pre block (projection + Sinkhorn producing the next |
| 1314 | seam's post/comb/pre gates, plus the collapse with the carried pre-mix), |
| 1315 | and the RMSNorm folded onto the collapse. |
| 1316 | |
| 1317 | ``num_tokens`` spans the kernel's split-K config buckets. ``seam`` covers |
| 1318 | the three call shapes the model makes: a regular seam (post folded in), an |
| 1319 | Engram seam (post already applied), and the draft entry seam |
| 1320 | (``pre_mix=None`` = identity pre-mix). |
| 1321 | """ |
| 1322 | set_random_seed(0) |
| 1323 | hc_mult, hidden_size = 4, 5120 |
| 1324 | residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs( |
| 1325 | num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult |
| 1326 | ) |
| 1327 | rms_eps = hc_pre_eps = hc_sinkhorn_eps = norm_eps = 1e-6 |
| 1328 | args = ( |
| 1329 | residual, |
| 1330 | fn, |
| 1331 | hc_scale, |
| 1332 | hc_base, |
| 1333 | rms_eps, |
| 1334 | hc_pre_eps, |
| 1335 | hc_sinkhorn_eps, |
| 1336 | 1.0, |
| 1337 | 20, |
| 1338 | ) |
| 1339 | |
| 1340 | pre_mix = ( |
| 1341 | None |
| 1342 | if seam == "identity_pre" |
| 1343 | else torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5 |
| 1344 | ) |
| 1345 | post_kwargs = {} |
| 1346 | residual_ref = residual |
| 1347 | if seam == "post": |
| 1348 | sublayer_out = torch.randn( |
| 1349 | num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE |
| 1350 | ) |
| 1351 | post_layer_mix = 2 * torch.rand( |
| 1352 | num_tokens, hc_mult, 1, dtype=torch.float32, device=DEVICE |
| 1353 | ) |
| 1354 | comb_res_mix = torch.softmax( |
| 1355 | torch.randn(num_tokens, hc_mult, hc_mult, device=DEVICE), dim=-1 |
| 1356 | ) |
| 1357 | post_kwargs = dict( |
| 1358 | sublayer_out=sublayer_out, |
| 1359 | post_layer_mix=post_layer_mix, |
| 1360 | comb_res_mix=comb_res_mix, |
| 1361 | ) |
| 1362 | residual_ref = mhc_post_torch( |
| 1363 | sublayer_out, residual, post_layer_mix, comb_res_mix |
| 1364 | ) |
| 1365 | |
| 1366 | expected = mhc_pre_delayed_torch(residual_ref, *args[1:], pre_mix=pre_mix) |
| 1367 | expected_layer_input = F.rms_norm( |
| 1368 | expected[2], (hidden_size,), norm_weight, norm_eps |
| 1369 | ) |
| 1370 | |
| 1371 | residual_out, *actual = object.__new__(MHCPreDelayedOp).forward_hip( |
| 1372 | *args, |
| 1373 | pre_mix=pre_mix, |
| 1374 | norm_weight=norm_weight, |
| 1375 | norm_eps=norm_eps, |
| 1376 | **post_kwargs, |
| 1377 | ) |
| 1378 | |
| 1379 | if seam == "post": |
| 1380 | torch.testing.assert_close(residual_out, residual_ref) |
| 1381 | else: |
| 1382 | # No post was requested, so the residual comes straight back. |
| 1383 | assert residual_out is residual |
| 1384 | # Gates: post_mix, comb_mix, next_pre_mix (bf16 MFMA vs the fp32 reference). |
| 1385 | for i in (0, 1, 3): |
| 1386 | torch.testing.assert_close(actual[i], expected[i], atol=5e-4, rtol=1e-3) |
| 1387 | # The collapse is bf16 and the kernel projects the once-rounded R', so |
| 1388 | # near-zero values need an absolute tolerance (same bound as aiter's UT). |
| 1389 | torch.testing.assert_close(actual[2], expected_layer_input, atol=2e-2, rtol=1e-2) |
| 1390 | |
| 1391 | |
| 1392 | @pytest.mark.skipif( |
| 1393 | not (current_platform.is_rocm() and HAS_AITER_MHC), |
| 1394 | reason="AITER mHC required", |
| 1395 | ) |
| 1396 | def test_mhc_pre_delayed_rocm_aiter_declines_unsupported(monkeypatch): |
| 1397 | """The broadcast seam must not take the AITER path; a fused norm must not |
| 1398 | take the unfused AITER op. |
| 1399 | |
| 1400 | The broadcast projects a narrower ``x``, which no AITER pre kernel takes; |
| 1401 | ``mhc_pre_gemm_sqrsum`` folds no RMSNorm, so a fused norm routes to the |
| 1402 | fused seam kernel where it exists (gfx950) and away from AITER otherwise. |
| 1403 | What is asserted here is the routing decision, which is what this gate |
| 1404 | owns; the numerics of the paths are covered by |
| 1405 | ``test_deepseek_v41_mhc_pre_delayed`` and |
| 1406 | ``test_mhc_pre_delayed_rocm_aiter_fused_rms_norm``. |
| 1407 | """ |
| 1408 | set_random_seed(0) |
| 1409 | hc_mult, hidden_size = 4, 5120 |
| 1410 | residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs( |
| 1411 | num_tokens=4, hidden_size=hidden_size, hc_mult=hc_mult |
| 1412 | ) |
| 1413 | args = (residual, fn, hc_scale, hc_base, 1e-6, 1e-6, 1e-6, 1.0, 20) |
| 1414 | op = object.__new__(MHCPreDelayedOp) |
| 1415 | |
| 1416 | aiter_op = torch.ops.vllm.mhc_pre_delayed_aiter |
| 1417 | took_aiter = False |
| 1418 | |
| 1419 | def spy(*spy_args, **spy_kwargs): |
| 1420 | nonlocal took_aiter |
| 1421 | took_aiter = True |
| 1422 | return aiter_op(*spy_args, **spy_kwargs) |
| 1423 | |
| 1424 | monkeypatch.setattr(torch.ops.vllm, "mhc_pre_delayed_aiter", spy) |
| 1425 | |
| 1426 | fused_op = torch.ops.vllm.mhc_fused_post_pre_delayed_rms_norm_aiter |
| 1427 | took_fused = False |
| 1428 | |
| 1429 | def fused_spy(*spy_args, **spy_kwargs): |
| 1430 | nonlocal took_fused |
| 1431 | took_fused = True |
| 1432 | return fused_op(*spy_args, **spy_kwargs) |
| 1433 | |
| 1434 | monkeypatch.setattr( |
| 1435 | torch.ops.vllm, "mhc_fused_post_pre_delayed_rms_norm_aiter", fused_spy |
| 1436 | ) |
| 1437 | |
| 1438 | x = residual[:, 0].contiguous() |
| 1439 | broadcast_fn = fn.view(-1, hc_mult, hidden_size).sum(1) |
| 1440 | broadcast_residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() |
| 1441 | expected = mhc_pre_delayed_torch(broadcast_residual, broadcast_fn, *args[2:], x=x) |
| 1442 | # forward_hip leads with the residual; no post was requested here, so it is |
| 1443 | # the caller's own tensor and the reference's four outputs follow it. |
| 1444 | residual_out, *actual = op.forward_hip( |
| 1445 | broadcast_residual, broadcast_fn, *args[2:], x=x |
| 1446 | ) |
| 1447 | assert residual_out is broadcast_residual |
| 1448 | assert not took_aiter, "the broadcast seam must not reach AITER" |
| 1449 | assert not took_fused, "the broadcast seam must not reach AITER" |
| 1450 | for i in range(4): |
| 1451 | torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3) |
| 1452 | |
| 1453 | # A fused norm must not reach the unfused AITER op; where the fused seam |
| 1454 | # kernel exists (gfx950) it routes there instead. |
| 1455 | op.forward_hip(*args, norm_weight=norm_weight, norm_eps=1e-6) |
| 1456 | assert not took_aiter, "a fused norm must not reach the unfused AITER op" |
| 1457 | assert took_fused == mhc_layers.HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM |
| 1458 | |
| 1459 | # Positive control: the same inputs without a norm do take the AITER path, |
| 1460 | # so the two declines above are the gate discriminating rather than the |
| 1461 | # path being unavailable in this environment. |
| 1462 | op.forward_hip(*args) |
| 1463 | assert took_aiter |
| 1464 | |
| 1465 | |
| 1466 | @pytest.mark.skipif( |
| 1467 | not (current_platform.is_rocm() and HAS_AITER_MHC and HAS_AITER_MHC_PRE_NORM), |
| 1468 | reason="AITER mHC with fused RMSNorm required", |
| 1469 | ) |
| 1470 | def test_mhc_pre_rocm_aiter_fuses_norm(): |
| 1471 | set_random_seed(0) |
| 1472 | residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs( |
| 1473 | num_tokens=2, hidden_size=7168 |
| 1474 | ) |
| 1475 | rms_eps = hc_pre_eps = hc_sinkhorn_eps = norm_eps = 1e-6 |
| 1476 | sinkhorn_repeat = 20 |
| 1477 | hc_post_alpha = 1.0 |
| 1478 | ref = mhc_pre_ref( |
| 1479 | residual, |
| 1480 | fn, |
| 1481 | hc_scale, |
| 1482 | hc_base, |
| 1483 | rms_eps, |
| 1484 | hc_pre_eps, |
| 1485 | hc_sinkhorn_eps, |
| 1486 | hc_post_alpha, |
| 1487 | sinkhorn_repeat, |
| 1488 | ) |
| 1489 | expected_layer_input = F.rms_norm( |
| 1490 | ref[2], (ref[2].shape[-1],), norm_weight, norm_eps |
| 1491 | ) |
| 1492 | |
| 1493 | out = object.__new__(MHCPreOp).forward_hip( |
| 1494 | residual, |
| 1495 | fn, |
| 1496 | hc_scale, |
| 1497 | hc_base, |
| 1498 | rms_eps, |
| 1499 | hc_pre_eps, |
| 1500 | hc_sinkhorn_eps, |
| 1501 | hc_post_alpha, |
| 1502 | sinkhorn_repeat, |
| 1503 | norm_weight=norm_weight, |
| 1504 | norm_eps=norm_eps, |
| 1505 | ) |
| 1506 | |
| 1507 | torch.testing.assert_close(out[0], ref[0], atol=5e-2, rtol=1e-2) |
| 1508 | torch.testing.assert_close(out[1], ref[1], atol=5e-2, rtol=1e-2) |
| 1509 | torch.testing.assert_close(out[2], expected_layer_input, atol=5e-2, rtol=1e-2) |
| 1510 | |
| 1511 | |
| 1512 | @pytest.mark.skipif( |
| 1513 | not ( |
| 1514 | current_platform.is_rocm() and HAS_AITER_MHC_FUSED and HAS_AITER_MHC_FUSED_NORM |
| 1515 | ), |
| 1516 | reason="AITER fused mHC with RMSNorm required", |
| 1517 | ) |
| 1518 | def test_mhc_fused_rocm_aiter_fuses_norm(): |
| 1519 | set_random_seed(0) |
| 1520 | residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs( |
| 1521 | num_tokens=2, hidden_size=7168 |
| 1522 | ) |
| 1523 | x = torch.randn((2, 7168), dtype=torch.bfloat16, device=DEVICE) |
| 1524 | post_layer_mix = torch.randn((2, 4, 1), dtype=torch.float32, device=DEVICE) |
| 1525 | comb_res_mix = torch.randn((2, 4, 4), dtype=torch.float32, device=DEVICE) |
| 1526 | rms_eps = hc_pre_eps = hc_sinkhorn_eps = norm_eps = 1e-6 |
| 1527 | sinkhorn_repeat = 20 |
| 1528 | hc_post_alpha = 1.0 |
| 1529 | residual_ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix) |
| 1530 | pre_ref = mhc_pre_ref( |
| 1531 | residual_ref, |
| 1532 | fn, |
| 1533 | hc_scale, |
| 1534 | hc_base, |
| 1535 | rms_eps, |
| 1536 | hc_pre_eps, |
| 1537 | hc_sinkhorn_eps, |
| 1538 | hc_post_alpha, |
| 1539 | sinkhorn_repeat, |
| 1540 | ) |
| 1541 | expected_layer_input = F.rms_norm( |
| 1542 | pre_ref[2], (pre_ref[2].shape[-1],), norm_weight, norm_eps |
| 1543 | ) |
| 1544 | |
| 1545 | out = object.__new__(MHCFusedPostPreOp).forward_hip( |
| 1546 | x, |
| 1547 | residual, |
| 1548 | post_layer_mix, |
| 1549 | comb_res_mix, |
| 1550 | fn, |
| 1551 | hc_scale, |
| 1552 | hc_base, |
| 1553 | rms_eps, |
| 1554 | hc_pre_eps, |
| 1555 | hc_sinkhorn_eps, |
| 1556 | hc_post_alpha, |
| 1557 | sinkhorn_repeat, |
| 1558 | norm_weight=norm_weight, |
| 1559 | norm_eps=norm_eps, |
| 1560 | ) |
| 1561 | |
| 1562 | torch.testing.assert_close(out[0], residual_ref, atol=5e-2, rtol=1e-2) |
| 1563 | torch.testing.assert_close(out[1], pre_ref[0], atol=5e-2, rtol=1e-2) |
| 1564 | torch.testing.assert_close(out[2], pre_ref[1], atol=5e-2, rtol=1e-2) |
| 1565 | torch.testing.assert_close(out[3], expected_layer_input, atol=5e-2, rtol=1e-2) |
| 1566 | |
| 1567 | |
| 1568 | @pytest.mark.skipif( |
| 1569 | not current_platform.is_rocm(), |
| 1570 | reason="ROCm required", |
| 1571 | ) |
| 1572 | @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) |
| 1573 | @pytest.mark.parametrize("hidden_size", [4096, 7168]) |
| 1574 | @pytest.mark.parametrize("hc_mult", [4]) |
| 1575 | def test_hc_head_triton(num_tokens, hidden_size, hc_mult): |
| 1576 | torch.set_default_device(DEVICE) |
| 1577 | set_random_seed(0) |
| 1578 | |
| 1579 | residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) |
| 1580 | fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4 |
| 1581 | hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1 |
| 1582 | hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1 |
| 1583 | rms_eps = hc_eps = 1e-6 |
| 1584 | |
| 1585 | out = torch.empty((num_tokens, hidden_size), dtype=torch.bfloat16) |
| 1586 | out.fill_(float("nan")) |
| 1587 | |
| 1588 | result = torch.ops.vllm.hc_head_triton( |
| 1589 | residual, |
| 1590 | fn, |
| 1591 | hc_scale, |
| 1592 | hc_base, |
| 1593 | out, |
| 1594 | hidden_size, |
| 1595 | rms_eps, |
| 1596 | hc_eps, |
| 1597 | hc_mult, |
| 1598 | ) |
| 1599 | |
| 1600 | assert result is None |
| 1601 | assert not torch.isnan(out).any() |
| 1602 | |
| 1603 | out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps) |
| 1604 | torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2) |
| 1605 | |
| 1606 | |
| 1607 | @pytest.mark.skipif( |
| 1608 | not HAS_TILELANG_MHC, |
| 1609 | reason="TileLang MHC support required", |
| 1610 | ) |
| 1611 | @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) |
| 1612 | @pytest.mark.parametrize("hidden_size", [4096, 7168]) |
| 1613 | @pytest.mark.parametrize("hc_mult", [4]) |
| 1614 | def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult): |
| 1615 | torch.set_default_device(DEVICE) |
| 1616 | set_random_seed(0) |
| 1617 | |
| 1618 | residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) |
| 1619 | fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4 |
| 1620 | hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1 |
| 1621 | hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1 |
| 1622 | rms_eps = hc_eps = 1e-6 |
| 1623 | |
| 1624 | out = torch.ops.vllm.hc_head_fused_kernel_tilelang( |
| 1625 | residual, |
| 1626 | fn, |
| 1627 | hc_scale, |
| 1628 | hc_base, |
| 1629 | rms_eps, |
| 1630 | hc_eps, |
| 1631 | ) |
| 1632 | |
| 1633 | assert out.shape == (num_tokens, hidden_size) |
| 1634 | assert out.dtype == torch.bfloat16 |
| 1635 | assert not torch.isnan(out).any() |
| 1636 | |
| 1637 | out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps) |
| 1638 | torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2) |
| 1639 | |
| 1640 | |
| 1641 | def _make_mhc_decoder_layer(hc_mult: int, hidden_size: int) -> DeepseekV4DecoderLayer: |
| 1642 | layer = DeepseekV4DecoderLayer.__new__(DeepseekV4DecoderLayer) |
| 1643 | nn.Module.__init__(layer) |
| 1644 | layer.hc_mult = hc_mult |
| 1645 | layer.hidden_size = hidden_size |
| 1646 | mix_hc = (2 + hc_mult) * hc_mult |
| 1647 | layer.hc_attn_fn = nn.Parameter( |
| 1648 | torch.randn(mix_hc, hc_mult * hidden_size, dtype=torch.float32), |
| 1649 | requires_grad=False, |
| 1650 | ) |
| 1651 | layer.hc_attn_fn_broadcast = None |
| 1652 | return layer |
| 1653 | |
| 1654 | |
| 1655 | def _patch_first_rank_pp_group(monkeypatch): |
| 1656 | monkeypatch.setattr( |
| 1657 | "vllm.models.deepseek_v4.nvidia.model.get_pp_group", |
| 1658 | lambda: SimpleNamespace(is_first_rank=True), |
| 1659 | ) |
| 1660 | |
| 1661 | |
| 1662 | def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch): |
| 1663 | """First finalize (at the end of load_weights) allocates |
| 1664 | hc_attn_fn_broadcast as hc_attn_fn summed over hc streams.""" |
| 1665 | _patch_first_rank_pp_group(monkeypatch) |
| 1666 | layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) |
| 1667 | model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) |
| 1668 | |
| 1669 | DeepseekV4Model.finalize_mhc_broadcast_weights(model) |
| 1670 | |
| 1671 | assert layer.hc_attn_fn_broadcast is not None |
| 1672 | expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1) |
| 1673 | assert torch.equal(layer.hc_attn_fn_broadcast, expected) |
| 1674 | |
| 1675 | |
| 1676 | def test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place(monkeypatch): |
| 1677 | """Re-finalizing after a weight refit must copy into the existing |
| 1678 | broadcast tensor so its address stays stable for captured CUDA graphs, |
| 1679 | while picking up the new hc_attn_fn values.""" |
| 1680 | _patch_first_rank_pp_group(monkeypatch) |
| 1681 | layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) |
| 1682 | model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) |
| 1683 | |
| 1684 | DeepseekV4Model.finalize_mhc_broadcast_weights(model) |
| 1685 | buffer = layer.hc_attn_fn_broadcast |
| 1686 | |
| 1687 | layer.hc_attn_fn.add_(1.0) |
| 1688 | DeepseekV4Model.finalize_mhc_broadcast_weights(model) |
| 1689 | |
| 1690 | assert layer.hc_attn_fn_broadcast is buffer |
| 1691 | expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1) |
| 1692 | assert torch.equal(layer.hc_attn_fn_broadcast, expected) |
| 1693 | |
| 1694 | |
| 1695 | @pytest.mark.skipif( |
| 1696 | not (current_platform.is_rocm() and HAS_AITER_MHC), |
| 1697 | reason="AITER mHC required", |
| 1698 | ) |
| 1699 | @pytest.mark.parametrize("num_tokens", [1, 2, 7, 128]) |
| 1700 | @pytest.mark.parametrize("carried", [False, True]) |
| 1701 | def test_mhc_fused_post_pre_delayed_rocm_aiter( |
| 1702 | num_tokens, carried, default_vllm_config |
| 1703 | ): |
| 1704 | """Folding the post into the pre projection must match doing them apart. |
| 1705 | |
| 1706 | The residual is bfloat16, so the two accumulation orders are allowed to |
| 1707 | land a rounding step apart; the tolerance here is one bfloat16 ULP. |
| 1708 | """ |
| 1709 | set_random_seed(0) |
| 1710 | hc_mult, hidden_size = 4, 5120 |
| 1711 | residual, fn, hc_scale, hc_base, _ = _rocm_mhc_inputs( |
| 1712 | num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult |
| 1713 | ) |
| 1714 | x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16, device=DEVICE) |
| 1715 | post_layer_mix = ( |
| 1716 | torch.rand(num_tokens, hc_mult, 1, dtype=torch.float32, device=DEVICE) + 0.5 |
| 1717 | ) |
| 1718 | comb_res_mix = ( |
| 1719 | torch.rand(num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=DEVICE) |
| 1720 | + 0.2 |
| 1721 | ) |
| 1722 | pre_mix = ( |
| 1723 | torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5 |
| 1724 | if carried |
| 1725 | else None |
| 1726 | ) |
| 1727 | rms_eps = hc_pre_eps = hc_sinkhorn_eps = 1e-6 |
| 1728 | pre_args = (fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, 1.0, 20) |
| 1729 | |
| 1730 | next_residual = mhc_post_torch(x, residual, post_layer_mix, comb_res_mix) |
| 1731 | expected = ( |
| 1732 | next_residual, |
| 1733 | *mhc_pre_delayed_torch(next_residual, *pre_args, pre_mix=pre_mix), |
| 1734 | ) |
| 1735 | actual = MHCPreDelayedOp().forward_hip( |
| 1736 | residual, |
| 1737 | *pre_args, |
| 1738 | pre_mix=pre_mix, |
| 1739 | sublayer_out=x, |
| 1740 | post_layer_mix=post_layer_mix, |
| 1741 | comb_res_mix=comb_res_mix, |
| 1742 | ) |
| 1743 | |
| 1744 | for got, want in zip(actual, expected, strict=True): |
| 1745 | torch.testing.assert_close(got, want.reshape(got.shape), atol=2e-2, rtol=8e-3) |
| 1746 | |
| 1747 | |
| 1748 | @pytest.mark.skipif( |
| 1749 | not (current_platform.is_rocm() and HAS_AITER_MHC), |
| 1750 | reason="AITER mHC required", |
| 1751 | ) |
| 1752 | def test_mhc_fused_post_pre_delayed_falls_back_for_large_batches(): |
| 1753 | """AITER's heuristic hands large batches back to the unfused path.""" |
| 1754 | from vllm._aiter_ops import rocm_aiter_ops |
| 1755 | |
| 1756 | assert not rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1) |
| 1757 | assert rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1 << 20) |