fused_mla_yarn_rope_apply Triton kernels read and write out of bounds when head_num % BLOCK_H != 0

**Describe the bug**

@NVIDIA/mcore-oncall

The bounds masks in the four Triton kernels in
`megatron/core/fusions/fused_mla_yarn_rope_apply.py` are computed in **block-local** head
coordinates but compared against the **global** head count. Each kernel advances its base
pointer by `pid_head * BLOCK_H * stride_nheads` and then builds its offset from
`tl.arange(0, BLOCK_H)` alone, never adding `pid_head` back:

```python
Q = Q + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads # block folded in
x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim # LOCAL only
mask = x_off < head_num * stride_x_nheads # GLOBAL bound
...
tl.store(Q + x_left_off, x_left, mask=mask)
tl.store(Q + x_right_off, x_right, mask=mask)
```

So `mask` tests `local_h < head_num` where correctness requires
`pid_head * BLOCK_H + local_h < head_num`. It is accidentally correct at `pid_head == 0`
and strictly too permissive above it: every lane of the last, partial head block passes.

`BLOCK_H` is autotuned over `{1, 2, 4, 8, 16, 32, 64, 128}`, so **any `head_num` that is
not a power of two** admits a tiling with `head_num % BLOCK_H != 0`. That is why this has
gone unnoticed — the existing unit tests all use `num_heads=32`.

Five mask sites are affected (line numbers at `b1fe7599e`):

| line | kernel | offset var | block index |
|---|---|---|---|
| 142 | `_mla_rope_fwd_inplace_kernel` | `x_off` | `pid_head` |
| 238 | `_mla_rope_bwd_inplace_kernel` | `x_off` | `pid_head` |
| 569 | `_mla_rope_fwd_kv_split_kernel` | `kv_off` | `pid_head` |
| 679 | `_mla_rope_bwd_kv_split_kernel` | `dkv_off` | `pid_head` |
| 698 | `_mla_rope_bwd_kv_split_kernel` | `x_off` | `i` (inside `tl.static_range`) |

**Why it is a write, not just a read**

The backward kernels `tl.store` through the same mask, so the phantom lanes are written,
not only loaded. `restore_value=["Q"]` restores the tensor's **own** storage between
autotune trials and cannot undo a write **past** it. And `Autotuner.run` benchmarks every
candidate whenever more than one config is registered, so the unsafe tiling is *executed*
on the first call on every process, regardless of which config finally wins.

**What it costs, concretely**

At `head_num=12`, `qk_head_dim=128`, `qk_pos_emb_head_dim=64` (per-head stride 192),
`BLOCK_H=8` gives `cdiv(12, 8) = 2` blocks. Block `pid_head=1` covers heads 8-15 while
only 8-11 exist, and all eight lanes pass the mask. A token's row is `12 * 192 = 2304`
elements, so the four phantom lanes land at `+128`, `+320`, `+512`, `+704` past the row:

- **interior tokens**: silently corrupts the *next* token's heads 0-3 rope slice, racing
the legitimate writer, block `(t+1, pid_head=0)`;
- **the last token**: writes 256 elements **past the end of the allocation**.

**Steps/Code to reproduce bug**

Standalone, no training stack, ~30 s. Faults outright when there is no allocation slack
behind `q`:

```python
import torch
import megatron.core.fusions.fused_mla_yarn_rope_apply as M

q = torch.randn(8192, 1, 12, 192, device="cuda", dtype=torch.bfloat16, requires_grad=True)
cos = torch.randn(8192, 64, device="cuda", dtype=torch.bfloat16).contiguous()
sin = torch.randn(8192, 64, device="cuda", dtype=torch.bfloat16).contiguous()
out = M.fused_mla_rope_inplace(q, cos, sin, 128, 64)
out.backward(torch.randn_like(out)) # dies during autotuning
```

Observed on an H100 under `CUDA_LAUNCH_BLOCKING=1`:

```
AcceleratorError: CUDA error: an illegal memory access was encountered
```

raised out of `_mla_rope_bwd_inplace_kernel` during the autotune benchmark pass
(`triton/runtime/autotuner.py` `_bench` -> `benchmark` -> `run`). The same script completes
cleanly with the mask corrected. With allocation slack behind `q`, the process survives and silently
corrupts whatever is behind `q` instead — which is the mode that reaches real training, as
a NaN or an illegal access hundreds of steps later with no precursor.

Deterministically, with `BLOCK_H` pinned to 8 and a sentinel guard region behind `q`, the
guard records exactly **256 clobbered elements** on the unpatched kernels and **0** with
the mask corrected; the corrected `BLOCK_H=8` output is bit-identical to `BLOCK_H=4`
(a tiling that divides 12).

**Expected behavior**

The partial head block writes only the heads that exist. `BLOCK_H` selects a tiling, not
an arithmetic, so the numerical result should be identical for every candidate.

**Environment overview**

- Megatron-LM `b1fe7599e` (also present in `core_r0.16.1`)
- H100 80GB HBM3, torch 2.10.0+cu129, triton 3.6.0, CUDA 12.9
- Reproduces on a single GPU with no distributed setup

**Proposed fix**

Add back the block offset the base pointer already consumed, at each of the five sites:

```diff
- mask = x_off < head_num * stride_x_nheads
+ mask = pid_head * BLOCK_H * stride_x_nheads + x_off < head_num * stride_x_nheads
```

using that kernel's stride and block index (`i`, not `pid_head`, at line 698).

PR to follow.

**Possibly related**

- #5317 — `apply_rope_fusion=True` NaN in a DSv4-Hybrid mock pretrain. Distinct root causes
are already discussed there (Q aliasing, `rotary_percent`); this mask defect is a third,
independent one in the same file and would not be caught by either.

Start in megatron/core/fusions/fused_mla_yarn_rope_apply.py and inspect the five mask sites in the four Triton kernels, then run the standalone CUDA reproducer with head_num=12 and BLOCK_H pinned to 8. Add coverage for a partial head block and verify that no guard-region elements are changed, the corrected output matches a dividing tiling, and the existing 32-head behavior remains intact.

Written by the indexing model from the issue text.