diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index ade9faf08ac..5ca376d7e4f 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -125,7 +125,7 @@ def _mla_rope_fwd_inplace_kernel( Q = Q + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim - mask = x_off < head_num * stride_x_nheads + mask = (pid_head * BLOCK_H + tl.arange(0, BLOCK_H))[:, None] < head_num # x1 = t[..., 0::2], x2 = t[..., 1::2] x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 x_2_off = x_1_off + 1 @@ -226,7 +226,7 @@ def _mla_rope_bwd_inplace_kernel( DO = DO + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim - mask = x_off < head_num * stride_x_nheads + mask = (pid_head * BLOCK_H + tl.arange(0, BLOCK_H))[:, None] < head_num if REMOVE_INTERLEAVING: x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 x_2_off = x_1_off + 1 @@ -577,7 +577,7 @@ def _mla_rope_fwd_kv_split_kernel( KV_ptr = KV + pid_m * stride_kv_seq + pid_head * BLOCK_H * stride_kv_nheads kv_off = tl.arange(0, BLOCK_H)[:, None] * stride_kv_nheads - mask = kv_off < head_num * stride_kv_nheads + mask = (pid_head * BLOCK_H + tl.arange(0, BLOCK_H))[:, None] < head_num k_in_off = kv_off + tl.arange(0, k_dim)[None, :] v_in_off = kv_off + k_dim + tl.arange(0, v_dim)[None, :] k = tl.load(KV_ptr + k_in_off, mask=mask) @@ -687,7 +687,7 @@ def _mla_rope_bwd_kv_split_kernel( dKV_ptr = dKV + pid_m * stride_dkv_seq + pid_head * BLOCK_H * stride_dkv_nheads dkv_off = tl.arange(0, BLOCK_H)[:, None] * stride_dkv_nheads - mask = dkv_off < head_num * stride_dkv_nheads + mask = (pid_head * BLOCK_H + tl.arange(0, BLOCK_H))[:, None] < head_num dk_out_off = dkv_off + tl.arange(0, k_dim)[None, :] dv_out_off = dkv_off + k_dim + tl.arange(0, v_dim)[None, :] @@ -706,17 +706,21 @@ def _mla_rope_bwd_kv_split_kernel( for i in tl.static_range(triton.cdiv(head_num, BLOCK_H)): dK_ptr = dK + pid_m * stride_dk_seq + i * BLOCK_H * stride_dk_nheads x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim - mask = x_off < head_num * stride_dk_nheads + mask = (i * BLOCK_H + tl.arange(0, BLOCK_H))[:, None] < head_num + # ``other=0`` is required, not cosmetic: a masked-out lane is undefined without it, + # and these values are added unconditionally into the accumulators below and then + # reduced with ``tl.sum``. Every other masked load in this file feeds a masked store, + # which discards the invalid lanes; a reduction cannot. if REMOVE_INTERLEAVING: x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 x_2_off = x_1_off + 1 - x_left = tl.load(dK_ptr + x_1_off, mask=mask) - x_right = tl.load(dK_ptr + x_2_off, mask=mask) + x_left = tl.load(dK_ptr + x_1_off, mask=mask, other=0.0) + x_right = tl.load(dK_ptr + x_2_off, mask=mask, other=0.0) else: x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] x_right_off = x_left_off + emb_dim // 2 - x_left = tl.load(dK_ptr + x_left_off, mask=mask) - x_right = tl.load(dK_ptr + x_right_off, mask=mask) + x_left = tl.load(dK_ptr + x_left_off, mask=mask, other=0.0) + x_right = tl.load(dK_ptr + x_right_off, mask=mask, other=0.0) x_left_accum += x_left x_right_accum += x_right x_left_accum = tl.sum(x_left_accum, axis=0) diff --git a/tests/unit_tests/determinism/kernels/test_fused_triton_kernels.py b/tests/unit_tests/determinism/kernels/test_fused_triton_kernels.py index bdd62a872b5..7b981976c8a 100644 --- a/tests/unit_tests/determinism/kernels/test_fused_triton_kernels.py +++ b/tests/unit_tests/determinism/kernels/test_fused_triton_kernels.py @@ -5,12 +5,17 @@ * ``fused_pad_routing_map`` / ``fused_indices_to_multihot``: integer routing bookkeeping with unique stores -- replay and agree with the torch reference. * MLA YaRN RoPE (``fused_mla_yarn_rope_apply``): elementwise rotations under a timing-based - ``triton.autotune``; forward and backward replay in sbhd and thd layouts. + ``triton.autotune``; forward and backward replay in sbhd and thd layouts, at a head count the + autotuned ``BLOCK_H`` divides and at one it does not, plus a check that the result does not + depend on ``BLOCK_H`` at all -- a replay cannot see that, because Triton caches the config the + first call chose. * mHC hyper-connection kernels (``fused_mhc_kernels``): Sinkhorn, h-aggregate and h-post-BDA contain real ``tl.sum`` reductions under autotune, on the Triton, native (``torch.compile``) and, when available, cuTile backends. """ +import contextlib + import pytest import torch @@ -32,12 +37,14 @@ HAVE_TRITON = False try: + from megatron.core.fusions import fused_mla_yarn_rope_apply as fused_mla_rope_module from megatron.core.fusions.fused_mla_yarn_rope_apply import ( fused_mla_rope_inplace, fused_mla_rope_kv_split, fused_mla_rope_out_of_place, ) except ImportError: + fused_mla_rope_module = None fused_mla_rope_inplace = fused_mla_rope_kv_split = fused_mla_rope_out_of_place = None pytestmark = pytest.mark.skipif( @@ -159,12 +166,43 @@ def _yarn_cos_sin(emb_dim, seqlen, dtype): return (torch.cos(freqs) * mscale).to(dtype), (torch.sin(freqs) * mscale).to(dtype) +@contextlib.contextmanager +def _pinned_block_h(block_h): + """Leave every autotuned MLA RoPE kernel one ``BLOCK_H`` (head-block size) to choose from. + + ``BLOCK_H`` is chosen by timing, so it is not an input the caller controls. The replay tests + below cannot vary it -- Triton caches the tuned config, so every replay reuses whatever the + first call picked -- which is exactly why a result that moves with the tiling can look stable + under replay and still move from run to run in a real job. + """ + kernels = [ + fused_mla_rope_module._autotuned_mla_rope_fwd_inplace_kernel, + fused_mla_rope_module._autotuned_mla_rope_bwd_inplace_kernel, + fused_mla_rope_module._mla_rope_fwd_kv_split_kernel, + fused_mla_rope_module._mla_rope_bwd_kv_split_kernel, + ] + saved = [(kernel, kernel.configs, kernel.cache) for kernel in kernels] + try: + for kernel in kernels: + kernel.configs = [triton.Config({"BLOCK_H": block_h})] + kernel.cache = {} + yield + finally: + for kernel, configs, cache in saved: + kernel.configs = configs + kernel.cache = cache + + +# Every kernel here is launched over ``cdiv(head_num, BLOCK_H)`` head programs. 12 heads at the +# autotuned ``BLOCK_H=8`` leaves a final program covering four head rows that do not exist; 32 -- +# the only count these tests used -- never leaves one while more than one program is launched. @pytest.mark.skipif(fused_mla_rope_out_of_place is None, reason="fused MLA RoPE unavailable") @pytest.mark.parametrize("layout", ["sbhd", "thd"]) @pytest.mark.parametrize("variant", ["out_of_place", "inplace"]) -def test_fused_mla_rope_q_replays_fwd_bwd(layout, variant): +@pytest.mark.parametrize("heads", [32, 12]) +def test_fused_mla_rope_q_replays_fwd_bwd(layout, variant, heads): seeded() - heads, nope_dim, emb_dim = 32, 128, 64 + nope_dim, emb_dim = 128, 64 dtype = torch.bfloat16 if layout == "sbhd": seqlen, batch = 2048, 2 @@ -185,11 +223,15 @@ def run(t): assert_replays_bit_exact(run, (t,), replays=3, what=f"fused_mla_rope_{variant}[{layout}]") +# Every kernel here is launched over ``cdiv(head_num, BLOCK_H)`` head programs. 12 heads at the +# autotuned ``BLOCK_H=8`` leaves a final program covering four head rows that do not exist; 32 -- +# the only count these tests used -- never leaves one while more than one program is launched. @pytest.mark.skipif(fused_mla_rope_kv_split is None, reason="fused MLA RoPE unavailable") @pytest.mark.parametrize("layout", ["sbhd", "thd"]) -def test_fused_mla_rope_kv_split_replays_fwd_bwd(layout): +@pytest.mark.parametrize("heads", [32, 12]) +def test_fused_mla_rope_kv_split_replays_fwd_bwd(layout, heads): seeded() - heads, k_dim, v_dim, emb_dim = 32, 128, 128, 64 + k_dim, v_dim, emb_dim = 128, 128, 64 dtype = torch.bfloat16 if layout == "sbhd": seqlen, batch = 2048, 2 @@ -216,6 +258,76 @@ def run(kv, k_pos_emb): ) +@pytest.mark.skipif(fused_mla_rope_kv_split is None, reason="fused MLA RoPE unavailable") +@pytest.mark.parametrize("path", ["q", "kv_split"]) +def test_fused_mla_rope_is_independent_of_block_h(path): + """One fixed input must give the same bits at both tilings. + + ``BLOCK_H`` is a tiling choice the autotuner makes on timing, so a result that depends on it is + a result that depends on which config happened to win -- non-determinism reached through the + tuner rather than through a reduction order, and invisible to a replay that reuses one config. + """ + seeded() + # 12 % 8 == 4: the second of two head programs covers four head rows that do not exist. + # 12 % 4 == 0: every program is full. Same arithmetic, so the two must agree bit for bit. + heads = 12 + block_heads_partial = 8 + block_heads_exact = 4 + nope_dim, k_dim, v_dim, emb_dim = 128, 128, 128, 64 + dtype = torch.bfloat16 + bounds = [0, 1000, 1500, 2048, 4096] + seqlen = max(b - a for a, b in zip(bounds, bounds[1:])) + cu_seqlens = torch.tensor(bounds, dtype=torch.int32, device="cuda") + cos, sin = _yarn_cos_sin(emb_dim, seqlen, dtype) + + if path == "q": + base = torch.randn(bounds[-1], heads, nope_dim + emb_dim, device="cuda", dtype=dtype) + grad = torch.randn_like(base) + + def run(): + t = base.detach().clone().requires_grad_(True) + out = fused_mla_rope_out_of_place( + t, cos, sin, nope_dim, emb_dim, cu_seqlens_q=cu_seqlens + ) + out.backward(grad.clone()) + return [out.detach(), t.grad] + + else: + base_kv = torch.randn(bounds[-1], heads, k_dim + v_dim, device="cuda", dtype=dtype) + base_emb = torch.randn(bounds[-1], 1, emb_dim, device="cuda", dtype=dtype) + grad_k = torch.randn(bounds[-1], heads, k_dim + emb_dim, device="cuda", dtype=dtype) + grad_v = torch.randn(bounds[-1], heads, v_dim, device="cuda", dtype=dtype) + + def run(): + kv = base_kv.detach().clone().requires_grad_(True) + k_pos_emb = base_emb.detach().clone().requires_grad_(True) + k_out, v_out = fused_mla_rope_kv_split( + kv, k_pos_emb, cos, sin, emb_dim, k_dim, v_dim, cu_seqlens_kv=cu_seqlens + ) + torch.autograd.backward((k_out, v_out), (grad_k.clone(), grad_v.clone())) + # k_pos_emb.grad is the dEMB reduction -- the one output here built by summing over + # head rows, so the only one a head row the mask should have dropped can reach. + return [k_out.detach(), v_out.detach(), kv.grad, k_pos_emb.grad] + + results = [] + for block_h in (block_heads_partial, block_heads_exact): + with _pinned_block_h(block_h): + results.append(run()) + + first, second = results + for i, (a, b) in enumerate(zip(first, second)): + torch.testing.assert_close( + a, + b, + rtol=0, + atol=0, + msg=lambda m, i=i: ( + f"fused_mla_rope_{path} output {i} differs between " + f"BLOCK_H={block_heads_partial} and BLOCK_H={block_heads_exact}: {m}" + ), + ) + + # --- mHC (hyper-connection) kernels ------------------------------------------------------- _MHC_BACKENDS = ["native", "triton"] diff --git a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py index 665957184b9..c4955b41f8d 100644 --- a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py +++ b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py @@ -1,5 +1,6 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +import contextlib import warnings from unittest.mock import MagicMock, patch @@ -17,6 +18,7 @@ try: import triton + from megatron.core.fusions import fused_mla_yarn_rope_apply as fused_mla_rope_module from megatron.core.fusions.fused_mla_yarn_rope_apply import ( HAVE_TRITON, _mla_rope_bwd_inplace_kernel, @@ -26,9 +28,10 @@ fused_mla_rope_kv_split, fused_mla_rope_out_of_place, ) -except Exception: +except ImportError: HAVE_TRITON = False triton = None + fused_mla_rope_module = None _mla_rope_bwd_inplace_kernel = None _mla_rope_fwd_inplace_kernel = None fused_apply_mla_rope_for_q = None @@ -197,9 +200,10 @@ def backward(ctx, _grad_output): return saved_output -def _test_fused_mla_rope_inplace(input_format, inverse=False, remove_interleaving=False): +def _test_fused_mla_rope_inplace( + input_format, inverse=False, remove_interleaving=False, num_heads=32 +): assert fused_mla_rope_inplace is not None - num_heads = 32 q_dim = 128 emb_dim = 64 dtype = torch.bfloat16 @@ -293,9 +297,8 @@ def _test_fused_mla_rope_inplace(input_format, inverse=False, remove_interleavin ) -def _test_fused_mla_rope_kv_split(input_format, remove_interleaving=False): +def _test_fused_mla_rope_kv_split(input_format, remove_interleaving=False, num_heads=32): assert fused_mla_rope_kv_split is not None - num_heads = 32 k_dim = 128 v_dim = 128 emb_dim = 64 @@ -424,6 +427,89 @@ def _test_fused_mla_rope_kv_split(input_format, remove_interleaving=False): ) +# ------------------------------------------------------------------------------------------------- +# Partial head block coverage. +# +# Every kernel here is launched over cdiv(head_num, BLOCK_H) head programs, so a head count that +# BLOCK_H does not divide leaves the final program covering head rows that do not exist. BLOCK_H is +# autotuned over {1, 2, ..., 128} and the tuner may pick a divisor, so an awkward head count is not +# enough on its own -- a run can pass without the partial block ever being built. The tests below +# pin BLOCK_H instead, and each one spells out the head count and the head-block size it needs. +# ------------------------------------------------------------------------------------------------- + + +@contextlib.contextmanager +def _pinned_block_h(block_h): + """Leave every autotuned kernel one BLOCK_H (the head-block size) to choose from.""" + kernels = [ + fused_mla_rope_module._autotuned_mla_rope_fwd_inplace_kernel, + fused_mla_rope_module._autotuned_mla_rope_bwd_inplace_kernel, + fused_mla_rope_module._mla_rope_fwd_kv_split_kernel, + fused_mla_rope_module._mla_rope_bwd_kv_split_kernel, + ] + saved = [(kernel, kernel.configs, kernel.cache) for kernel in kernels] + try: + for kernel in kernels: + kernel.configs = [triton.Config({"BLOCK_H": block_h})] + kernel.cache = {} + yield + finally: + for kernel, configs, cache in saved: + kernel.configs = configs + kernel.cache = cache + + +def _thd_rope_tables(cu_seqlens_list, emb_dim, dtype): + """cos/sin for a THD batch, plus the cu_seqlens tensor the kernels want.""" + max_seqlen = max(b - a for a, b in zip(cu_seqlens_list, cu_seqlens_list[1:])) + yarn_rope = YarnRotaryEmbedding(emb_dim, original_max_position_embeddings=max_seqlen) + freqs, mscale = yarn_rope(max_seqlen, 0) + cos = (torch.cos(freqs) * mscale).to(dtype) + sin = (torch.sin(freqs) * mscale).to(dtype) + cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int32, device="cuda") + return cos, sin, cu_seqlens + + +def _run_kv_split_once(block_h, remove_interleaving, num_heads, seed=1234): + """One seeded forward and backward of the KV-split path at a pinned BLOCK_H. + + Returns the four tensors whose values must not depend on the tiling: both forward outputs, the + kv gradient, and the k_pos_emb gradient. The last one is the ``dEMB`` reduction, the one site in + this file where an out-of-range lane survives into the result instead of being dropped by a + masked store. + """ + k_dim = v_dim = 128 + emb_dim = 64 + dtype = torch.bfloat16 + cu_seqlens_list = [0, 27, 54, 99, 128] + total_seqlen = cu_seqlens_list[-1] + cos, sin, cu_seqlens = _thd_rope_tables(cu_seqlens_list, emb_dim, dtype) + + torch.manual_seed(seed) + kv = torch.randn((total_seqlen, num_heads, k_dim + v_dim), dtype=dtype, device="cuda") + emb = torch.randn((total_seqlen, 1, emb_dim), dtype=dtype, device="cuda") + dk = torch.randn((total_seqlen, num_heads, k_dim + emb_dim), dtype=dtype, device="cuda") + dv = torch.randn((total_seqlen, num_heads, v_dim), dtype=dtype, device="cuda") + kv.requires_grad_(True) + emb.requires_grad_(True) + + with _pinned_block_h(block_h): + k_out, v_out = fused_mla_rope_kv_split( + kv, + emb, + cos, + sin, + emb_dim, + k_dim, + v_dim, + cu_seqlens_kv=cu_seqlens, + remove_interleaving=remove_interleaving, + ) + torch.autograd.backward((k_out, v_out), (dk, dv)) + + return k_out.detach(), v_out.detach(), kv.grad.clone(), emb.grad.clone() + + def _q_inplace_test_data(input_format, emb_dim): num_heads = 32 q_dim = 128 @@ -488,6 +574,97 @@ def _q_inplace_torch_reference(source, cos, sin, token_idx, q_dim, direction): @pytest.mark.experimental @pytest.mark.internal @pytest.mark.skipif(not is_torch_min_version("2.5.0"), reason="Requires PyTorch >= 2.5.0") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("remove_interleaving", [False, True]) +class TestFusedMLARopePartialHeadBlock: + @pytest.mark.parametrize("input_format", ["sbhd", "thd"]) + @pytest.mark.parametrize("inverse", [False, True]) + def test_inplace_matches_unfused_reference(self, input_format, inverse, remove_interleaving): + num_heads = 12 # Deliberately not a power of two. + block_heads_partial = 8 # 12 % 8 == 4: the second of two programs is partial. + with _pinned_block_h(block_heads_partial): + _test_fused_mla_rope_inplace( + input_format, + inverse=inverse, + remove_interleaving=remove_interleaving, + num_heads=num_heads, + ) + + @pytest.mark.parametrize("input_format", ["sbhd", "thd"]) + def test_kv_split_matches_unfused_reference(self, input_format, remove_interleaving): + """Covers the forward outputs, the kv gradient and the k_pos_emb (``dEMB``) gradient.""" + num_heads = 12 # Deliberately not a power of two. + block_heads_partial = 8 # 12 % 8 == 4: the second of two programs is partial. + with _pinned_block_h(block_heads_partial): + _test_fused_mla_rope_kv_split( + input_format, remove_interleaving=remove_interleaving, num_heads=num_heads + ) + + def test_kv_split_result_does_not_depend_on_block_size(self, remove_interleaving): + """BLOCK_H is a tiling choice, so it must not change one bit of the result. + + This needs no reference implementation, which is what makes it a sharper check than the + tolerance-based comparisons above: a lane that a mask should have excluded shows up as a + difference between two block sizes whatever value the hardware happened to give it. + """ + num_heads = 12 # Deliberately not a power of two. + block_heads_partial = 8 # 12 % 8 == 4: the second of two programs is partial. + block_heads_exact = 4 # 12 % 4 == 0: every program is full. Same arithmetic, no partial. + partial = _run_kv_split_once(block_heads_partial, remove_interleaving, num_heads) + exact = _run_kv_split_once(block_heads_exact, remove_interleaving, num_heads) + for name, from_partial, from_exact in zip(("k", "v", "d_kv", "d_emb"), partial, exact): + torch.testing.assert_close( + from_partial, + from_exact, + rtol=0, + atol=0, + msg=lambda msg, name=name: f"{name} depends on BLOCK_H: {msg}", + ) + + def test_inplace_does_not_write_past_the_tensor(self, remove_interleaving): + """The rotated tensor is a view of the head of a larger buffer whose tail is poisoned. + + ``stride_x_seq == head_num * stride_x_nheads``, so the head rows the final program covers + but does not own are the next token's leading heads -- and, for the last token, memory past + the end of the tensor. The poisoned tail is where an unmasked store lands. + """ + num_heads = 12 # Deliberately not a power of two. + block_heads_partial = 8 # 12 % 8 == 4: the second of two programs is partial. + nope_dim = 128 + emb_dim = 64 + dtype = torch.bfloat16 + cu_seqlens_list = [0, 27, 54, 99, 128] + total_seqlen = cu_seqlens_list[-1] + cos, sin, cu_seqlens = _thd_rope_tables(cu_seqlens_list, emb_dim, dtype) + + # One guard token row already exceeds the four head rows the partial block over-covers. + guard_rows = 2 + buffer = torch.randn( + (total_seqlen + guard_rows, num_heads, nope_dim + emb_dim), dtype=dtype, device="cuda" + ) + guard_before = buffer[total_seqlen:].clone() + rotated = buffer[:total_seqlen] + assert rotated.stride() == buffer.stride() + + with _pinned_block_h(block_heads_partial), torch.no_grad(): + fused_mla_rope_inplace( + rotated, + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + remove_interleaving=remove_interleaving, + ) + + overflow = int((buffer[total_seqlen:] != guard_before).sum()) + assert overflow == 0, ( + f"{overflow} element(s) written past the end of the tensor: the partial final head " + f"block (head_num={num_heads}, BLOCK_H={block_heads_partial}) stored to head rows it " + f"does not own" + ) + + @pytest.mark.skipif(not HAVE_TRITON, reason="Triton not available") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("input_format", ["sbhd", "thd"])