ci: validate exact commit from #6890

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)

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"]