[huthvincent]

What

All five per-head bounds masks in megatron/core/fusions/fused_mla_yarn_rope_apply.py are built from the
intra-block head index, while the base pointer has already been advanced to the program's head block:

Q = Q + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads   # advanced to this block

x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim   # a LOCAL index
mask = x_off < head_num * stride_x_nheads                             # bounds the LOCAL index

The comparison reduces to local_head < head_num, which is true for every lane of every block. The launch
grid is cdiv(nheads, BLOCK_H) and BLOCK_H is autotuned over {1,2,...,128}, so whenever the chosen block
size does not divide the head count the final program runs unmasked. It then rotates
cdiv(head_num, BLOCK_H) * BLOCK_H - head_num head rows that do not exist — and because
stride_x_seq == head_num * stride_x_nheads, those rows are the next token's leading heads, or past the end
of the tensor for the last token.

So the output depends on the autotuned block size, which is meant to be a pure tiling choice.

This PR bounds the absolute head index instead. That is the idiom this file already uses: #6343 added new
kernels here with head_idx = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) / head_mask = head_idx < head_num
while leaving these five legacy sites as unchanged context.

The fifth site (the dEMB reduction in _mla_rope_bwd_kv_split_kernel) sits inside if pid_head == 0: and
iterates for i in tl.static_range(cdiv(head_num, BLOCK_H)), so it is offset by i, not pid_head — a
mechanical rewrite there would be a no-op and the k_pos_emb gradient would stay over-accumulated. Going
stride-free also fixes the reuse of a single mask across dKV/dK/dV, whose head strides differ.

We did not measure any performance impact of this change

Nothing in this PR is a performance claim, and no timing, throughput or memory figure is asserted anywhere in
it. The measurements below are correctness assertions only, taken on one NVIDIA L4 (PCIe, no NVLink), torch
2.12.0+cu130, triton 3.7.0 — hardware unlike anything this kernel targets. We have no H100 or GB200 access and
have deliberately not compared any number of ours against any number of yours.

How we checked it

BLOCK_H and num_warps=1 are pinned by launching the underlying JITFunction directly, so the autotuner
cannot vary the config, and no external reference implementation is involved.

1. The result depends on the block size — same kernel, same seeded input, two BLOCK_H values:

heads	BLOCK_H divides	BLOCK_H does not	before	after
12	4	8	differ	identical
24	8	16	differ	identical
96	32	64	differ	identical
16	8	4 — both divide	identical	identical

2. It writes past the end of the tensor — the tensor is a view of the head of a larger buffer whose tail is
poisoned; the tail is compared before and after the launch. The count matches the mechanism's prediction
exactly, (cdiv(H,BLOCK_H)*BLOCK_H - H) * emb_dim:

case	overflow	floats written past the end, before	after
H=12, BLOCK_H=8	4 heads	256	0
H=24, BLOCK_H=16	8 heads	512	0
H=48, BLOCK_H=32	16 heads	1024	0
H=96, BLOCK_H=64	32 heads	2048	0
H=12, BLOCK_H=4	none	0	0
H=16, BLOCK_H=8	none	0	0
H=12, BLOCK_H=16	grid collapses	0	0

3. The patch changes nothing where the old code was already correct — bitwise identical before and after
for H=16/BLOCK_H=8, H=12/BLOCK_H=4, H=32/BLOCK_H=8, H=128/BLOCK_H=64.

We could not run tests/unit_tests/transformer/test_mla_yarn_rope_apply.py as a whole here (no
transformer_engine on this box, and no wheel for it), so the kernels were exercised directly.

Not this bug: #4640 and #5963 report intermittent multi-warp mismatches in these same two in-place kernels
under remove_interleaving=False. Everything above pins num_warps=1 specifically so the two cannot be
confused, and the block-size comparison needs no oracle at all.

The strongest objection we know of, unanswered

An independent reviewer we ran this past put it this way, and we think it is fair:

The mask is provably correct for every configuration NVIDIA ships or tests, so this is a latent bug, not a
live one, and it can be legitimately closed as "add an assert". Because Megatron requires
num_attention_heads % TP == 0 and every published MLA model has a power-of-two head count (DeepSeek-V2/V3
128, V2-Lite 16, Kimi-K2 64), TP is forced to be a power of two and the per-rank count can never leave the
power-of-two family; every in-repo MLA config is 8/32/128 heads at TP 1-2 and the unit test uses 2 and 32.
Reaching the defect needs a hand-picked head count such as 96 on TP=8 or 12 on TP=1, and then only for the
single BLOCK_H that fails to divide it.

We have not found a shipped MLA recipe that triggers it either. What we can say is that there is no assert
on nheads anywhere in the file, that examples/gpt3/train_gpt3_175b_distributed.sh ships
--num-attention-heads 96 --tensor-model-parallel-size 8 (12 per partition) and
Nemotron-Mini-4B-Instruct.sh ships 24, so the shape exists in your own recipes even if not yet under MLA —
and that the correct form is already in this file from #6343.

If you would rather have an assert nheads % BLOCK_H == 0 or a constraint on nheads, say so and we will send
that instead.

Checklist

• -s sign-off, and the commit is SSH-signed
• black / isort / ruff per tools/autoformat.sh on the changed file; the two pre-existing F541
  hits elsewhere in the tree were left alone
• rebased on main; no out-of-scope changes
• no linked issue — this is a bug fix, which the template lists as recommended rather than required

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[huthvincent]

@NVIDIA/mcore-oncall — just marked this Ready. copy-pr-bot is waiting on validation, so NVIDIA's runners have not built it; could someone run /ok to test 844e8ca when convenient?

Five one-line mask fixes in megatron/core/fusions/fused_mla_yarn_rope_apply.py (+5 −5). The PR body has the mechanism, a standalone repro, and the check that matters most for merging it: the patch is bitwise identical to the current code for every head count and block size that already worked. Happy to rebase onto a newer main first if you would prefer that.

[wujingyue]

/claude fix

[svcnvidia-nemo-ci]

🛠️ Claude fix commit 115974900765 (attempt 1)

⚠️ This explanation is AI-generated and may be inaccurate; the exact commit is authoritative.

What changed
No changes made. The tree is exactly as I found it: no files edited, created, deleted, or renamed, and nothing committed.

Files changed by Claude

• No additional file edits; the pinned base was merged.

Why
Nothing in scope was repairable. The base merge is already fully auto-resolved: no unmerged paths, no conflict markers, and the base never touched the PR's only file, so a text conflict was impossible. No CI failure was actionable: the log dir is empty and the job name blank. The one runnable check passed (changed lines 75-76 chars vs the 100 limit; file byte-compiles; imports untouched). black, isort, ruff, pylint, torch and triton are all absent, so those checks are unsupported and left alone.

DCO
@huthvincent, please fix any DCO failures on your commits before merge. DCO does not block this workflow.

View exact commit

Sanitized and posted by svcnvidia-nemo-ci.

[svcnvidia-nemo-ci]

/ok to test 1159749

[wujingyue]

I'll defer approval to @FDecaYed

[svcnvidia-nemo-ci]

❌ Claude fix stopped because exact-SHA CI did not complete in time. View exact-SHA CI.

[wujingyue]

/claude fix

[svcnvidia-nemo-ci]

❌ Claude fix stopped because a workflow step failed. Inspect the run.

[wujingyue]

/ok to test

[copy-pr-bot]

/ok to test

@wujingyue, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/

[wujingyue]

/ok to test 2149f1e

[wujingyue]

The absolute-head masks look right, including using i in the reduction loop. Two points to address before merging; this is based on source inspection, not local GPU validation.

[wujingyue]

Please add other=0 to the four masked tl.load calls below this mask (both REMOVE_INTERLEAVING branches). With a partial final head block, masked-out values are undefined when other is omitted, but these values feed unconditional x_left_accum / x_right_accum additions and then tl.sum. The bounds fix therefore does not guarantee a correct dEMB reduction without zero-filling the invalid lanes. Unlike the other changed sites, a masked store cannot discard them after the reduction.

Triton documents this explicitly: https://triton-lang.org/main/python-api/generated/triton.language.load.html

[huthvincent]

Done -- other=0.0 on all four, with a comment on why the reduction needs it where the masked stores elsewhere in this file do not.

One thing to flag rather than let you infer it: on the hardware we have, the added tests pass with and without it. A direct probe (triton 3.7.0, one L4) has tl.load(ptr + off, mask=off < N) returning 0.0 in the masked lanes, so omitting other is indistinguishable from other=0.0 there. We are taking your citation -- the documented value is undefined -- as the reason for the change, not a measurement of our own. Restoring the old mask does fail all 16 of the new tests.

[wujingyue]

Could you commit regression coverage that forces a partial head block, e.g. head_num=12, BLOCK_H=8, rather than relying on the autotuner to select that configuration? The existing tests use power-of-two head counts, and the direct-kernel experiments described in the PR are not included as tests. Please cover forward output, backward gradients (especially the KV-split dEMB reduction), and guard-buffer integrity, exercising both interleaving branches. Adding a non-power-of-two head count alone can still miss the defect if autotuning picks a divisor.

[huthvincent]

Added TestFusedMLARopePartialHeadBlock. head_num=12 with BLOCK_H pinned to 8 by swapping each kernel's autotuner config list for a single entry, so the tuner cannot fall back on a divisor; BLOCK_H=4 is the control.

16 cases, over both REMOVE_INTERLEAVING branches:

• forward output and every gradient against the unfused reference, sbhd and thd, inverse both ways -- this is where the KV-split dEMB gradient is checked;
• BLOCK_H=8 vs BLOCK_H=4 bitwise on all four KV-split tensors, which needs no reference implementation and so catches a stray lane whatever value it held;
• guard-buffer integrity: the rotated tensor is a view of the head of a larger buffer with a poisoned tail, and the tail must come back untouched.

All 16 fail with this file restored to main's version and pass with the patch. Before the fix the guard buffer counts 255 elements written past the end of the tensor against the 256 the mechanism predicts -- one bf16 store landed on the poison value it replaced.

Unrelated, so a full-file run does not mislead you: four tests in this file (TestApplyRotaryPosEmbTHD, TestApplyRotaryPosEmbMlaFusionConflict) fail on our box because it has no transformer_engine. They fail identically on unmodified main.

[dimapihtar]

Hi @huthvincent
Could you resolve merge conflict, please?

[huthvincent]

Done — merged main; the conflict was with #5963's tests.

[Phlip79]

/ok to test f703e6f

[copy-pr-bot]

/ok to test f703e6f

@Phlip79, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[Phlip79]

/ok to test f703e6f

[wujingyue]

Hey @NVIDIA/ci, we may need an one-off to bump it up in the ci queue and/or the merge queue. It's intended for the 26.10 release but gets stuck in cicd-wait-in-queue.