[Heisenberg-Yin]

Thanks @TKONIY for the contribution to this PR!

• I, the PR author, have personally reviewed every line of this PR.

Summary

The MLA query RoPE and the DSA indexer RoPE (query and key) rotate their inputs with the in-place
Triton kernels of megatron/core/fusions/fused_mla_yarn_rope_apply.py, using
remove_interleaving=False:

• the MLA query RoPE is fused_apply_mla_rope_for_q in multi_latent_attention.py;
• the DSA indexer RoPE is DSAIndexer._apply_rope in dsa.py, which calls fused_mla_rope_inplace
  with rope_first=True for the indexer query [t, H_idx, 128] and key [t, 1, 128] when
  dsa_indexer_rope_interleaved=True. GLM-5 sets this option (its Hugging Face config has
  indexer_rope_interleave: true). With the default False, the indexer uses
  apply_rotary_pos_emb instead of these kernels.

Both callers use these kernels only with RoPE fusion on (apply_rope_fusion=True). In training,
their backward un-rotates the incoming gradient in place with the same kind of kernel.
With remove_interleaving=False, the kernels write each head's channels back to different positions
of the same buffer. Nothing makes a program finish its loads before it starts its stores. With
Triton 3.6, the backward kernel at BLOCK_H=2 splits every head between two warps. It can then
silently return wrong gradients that change from run to run. This happens when those warps drift
apart, which packed THD batches of many short sequences make likely.

The single-head DSA indexer key is the most exposed: its backward autotuner often chooses
BLOCK_H=2, and on packed THD batches the indexer key gradient then comes back wrong in nearly
every call.

This PR adds a tl.debug_barrier() before the permuting stores of both kernels, and a regression
test that fails without the barrier.

Related work

This race is the cause of the flaky test reported in #4640
(TestFusedApplyMLARope::test_forward_backward_for_q[thd], backward mismatch beyond bf16 tolerance),
and #5963 by @Sunt-ing fixes it on main with the same approach: a tl.debug_barrier() before the
permuting stores of the in-place forward and backward kernels. This PR is the dev counterpart; it
does not replace #5963:

• fused_mla_yarn_rope_apply.py has diverged substantially between main and dev, so Fix fused MLA Q-RoPE in-place load/store races #5963 does
  not apply to dev, and this patch does not apply to main (checked with git apply --check).
• On dev, the same in-place kernels also implement the DSA indexer RoPE (rope_first=True, used by
  GLM-5 with dsa_indexer_rope_interleaved=True). Its single-head key backward is the most exposed
  case, because the autotuner often settles on BLOCK_H=2 for it. The rope_first path does not
  exist on main.
• The regression test added here sweeps every BLOCK_H from 1 to 128 for an MLA query shape and a
  DSA indexer shape, checks the in-place forward and backward against an fp32 reference and bitwise
  across repeated runs, and includes an out= aliasing variant. It fails on unpatched dev and
  passes with this fix.

If #5963 lands on main first, the nightly main -> dev sync will touch the same lines. Both
changes add the barrier before the same stores, so the conflict can be resolved to either version.
I am happy to rebase this PR down to the dev-specific test coverage, or to close it, whichever the
maintainers prefer.

Root cause

Mechanism

With remove_interleaving=False, both kernels move the RoPE channels within each head of the same
buffer:

• The forward kernel _mla_rope_fwd_inplace_kernel loads the interleaved pairs (2i, 2i + 1) and
  stores the rotated values as two halves (i, d/2 + i).
• The backward kernel _mla_rope_bwd_kernel does the reverse: it loads the halves (i, d/2 + i) and
  stores interleaved pairs (2i, 2i + 1). _FusedMLARoPEInplace.backward reaches it through
  mla_rope_unapply_raw.

So element i of the [BLOCK_H, d/2] tile stores into channels that other elements load. Nothing
in the kernel orders one thread's loads before another thread's stores. Whether this matters depends
only on how Triton spreads the tile over the four warps.

Worked example: the GLM-5 DSA indexer key. It has one head of 128 channels, of which the first
64 are rotary, so d/2 = 32. At BLOCK_H=2, Triton 3.6 compiles the backward kernel with
threadsPerWarp = [2, 16], warpsPerCTA = [1, 4], order = [0, 1]:

warp	elements i	loads channels	stores channels
0	0-15	0-15 and 32-47	0-31
1	16-31	16-31 and 48-63	32-63
2, 3	copies of the above	(same)	none (stores are predicated off)

Warp 0's stores to channels 16-31 can land before warp 1 has loaded them. Warp 1's stores to
channels 32-47 can land before warp 0 has loaded those. In the compiled PTX nothing sits between the
global loads and the global stores (ld.global ... st.global, no bar.sync).

The other configurations are correct only because of the compiler's layout choices:

kernel	BLOCK_H	layout with Triton 3.6, 4 warps	result
backward	2	each head row is split between warps 0 and 1	races
backward	1	each head row stays within one warp, whose stores depend on its own loads	correct by layout
forward	1, 2	each head row stays within one warp (threadsPerWarp = [1, 32], warpsPerCTA = [4, 1])	correct by layout
both	4 to 128	the stores go through a shared-memory convert_layout; its bar.sync also orders the loads before the stores	correct by layout

How it shows up

Nothing fails. The affected channels come back as a mix of correct values and values computed from
inputs that were already overwritten. The errors are several units, up to about 8 in testing, while
bf16 rounding is about 0.03. Which elements are wrong changes from run to run.

In training, the result is a wrong gradient for the MLA query or, when the indexer loss is on, for
the DSA indexer query or key and, through them, for the indexer's projections. There is no error and
no NaN.

When it triggers

Configurations that run these kernels.

• RoPE fusion must be on: both callers check apply_rope_fusion=True through
  should_use_fused_mla_rope in rope_utils.py. The Megatron-LM training arguments turn it on by
  default; --no-rope-fusion turns it off.
• The DSA indexer also needs dsa_indexer_rope_interleaved=True. The default is False; GLM-5 sets
  it.
• The indexer's inputs are detached, and it records gradients only for its own KL loss: in training,
  with dsa_indexer_loss_coeff > 0. Only then does the indexer RoPE run backward, so only then can
  it race. The MLA query RoPE runs backward in every training step.

The race. It needs two things: the backward kernel runs at BLOCK_H=2, and the warps of a
program drift apart.

Inputs that make the warps drift. Packed THD batches of many short sequences make it likely,
because every program first walks cu_seqlens_q. The following was measured on the unfixed kernel
(upstream dev at ac100f773, B200, Triton 3.6.0, torch 2.11). Each cell is 20 calls on fresh
copies of one input. A call counts as wrong if it misses the fp32 formula by more than the bf16
tolerance of the test file, or differs bitwise from the first call.

shape (backward at BLOCK_H=2)	64 packed sequences	other inputs
GLM-5 DSA indexer query, 32 x 128, rope first	20/20 wrong at 512, 2048 and 8192 rows	one 4096-row sequence: 13/20; SBHD 2048 rows: 2/20
GLM-5 MLA query TP1, 64 x (192 + 64)	20/20 at 512, 2048 and 8192 rows	SBHD: 0/20
GLM-5 MLA query TP8, 8 x (192 + 64)	9/20 at 512 rows, 16/20 at 1024, 20/20 at 2048, 4096 and 8192	4 packed sequences, 4096 rows: 3/20 to 6/20 (three measurements); SBHD 512/1024 rows: 9/20
GLM-5 DSA indexer key, 1 x 128, rope first	8/20 at 2048 rows; 20/20 at 4096, 16384 and 65536; 19/20 or 20/20 at 8192 (four measurements)	one sequence of 512 to 262144 rows: 0/20
MLA query 128 x 192 and 16 x 192, DSA indexer 64 x 128 and 32 x 128	160/160 at 2048 and 8192 rows (all four shapes)

Cells with a range were measured several times. One of those measurements was an independent check
that compared every call bitwise with the same kernel's result into a separate out= buffer.

At 8192 rows, a wrong call had a median of about 100 wrong elements for the key and 7k to 337k for
the query shapes. In the scan of the last row, the other 47 configurations were wrong in 0 of 160
calls each. They cover the forward, the backward into a separate out=, both remove_interleaving
values and every other BLOCK_H.

Whether BLOCK_H=2 runs. The autotuner decides. Its key is emb_dim, head_num and the
dtypes, not the row count, and its first tuning applies to every later input in the process.

• The DSA indexer key is the most exposed; it has one head at every tensor-parallel size. Its
  backward autotuner chose BLOCK_H=2:
  • in 18 of 40 fresh tunings on one 65536-row sequence, an input on which BLOCK_H=2 happens to
    stay correct;
  • in 14 of 40 tunings on 64 packed sequences of 8192 rows, an input on which the key gradient is
    then wrong in nearly every call.
• Tensor parallelism divides the MLA query heads (64 / TP for GLM-5) but not the 32 DSA indexer heads,
  whose projections are duplicated across TP ranks.
• For the multi-head shapes, the autotuner chose 4 or more on inputs of 16384 rows or more. On short
  inputs, its own (warm) timings of several configurations are within a few percent of each other.
  For the GLM-5 MLA query at TP 8 (8 heads) on a 4096-row sequence, the backward autotuner chose
  BLOCK_H=2 in one of the six performance processes below, both before and after this change.
• Context parallelism divides the rows per rank (sequence length / CP), which makes short local
  inputs more common.

How it was found

It was found in long-sequence tests, while validating the fix for inputs of 2**31 or more elements
(#7787).
We compared these kernels bitwise at every BLOCK_H on MLA query and DSA indexer shapes. The
in-place backward at BLOCK_H=2 gave different results from run to run, both before and after that
fix.

Fix

Add a tl.debug_barrier() (a bar.sync 0) right before the permuting stores of both kernels,
after the arithmetic that consumes the loads:

• forward: in the remove_interleaving=False branch;
• backward: under if not REMOVE_INTERLEAVING:. The backward also takes the barrier when it writes
  to out=, because out= may alias the input; the wrapper only checks shape, stride and dtype.

Why it is correct and complete.

• A CTA barrier orders every global load of the program before every global store, whatever layout
  Triton picks. In the compiled PTX of every variant we checked (both kernels, both head layouts,
  every BLOCK_H), the barrier now sits between the last global load and the first global store.
• It does not depend on BLOCK_H, on purpose. Limiting it to BLOCK_H <= 2 would rely on Triton
  keeping the shared-memory conversion for larger tiles.
• The forward kernel has the mirror-image hazard. It is safe today only because of its layout, so it
  gets the barrier as well.
• Only these two kernels write channels to different positions within the buffer they read. With
  remove_interleaving=True, every value is stored where it was loaded. The packing and key/value
  split kernels read and write different buffers.

What it does not change.

• Both barriers sit in constexpr branches. All 97 remove_interleaving=True variants of the two
  kernels compile to the same PTX as before. Each of the 48 remove_interleaving=False variants gains
  exactly one bar.sync 0 and keeps every other instruction, so results that were not racing do not
  change.
• There is no API, argument or launch change.

Tests and results

New GPU test: tests/unit_tests/fusions/test_mla_yarn_rope_apply.py::test_mla_rope_inplace_channel_permutation_is_race_free.

• 16 cases: the MLA query (16 x 192) and the DSA indexer (32 x 128, rope first), each at every
  BLOCK_H the autotuner offers. Each BLOCK_H is pinned through a single-config autotuner.
• Each case runs three variants, eight times each, on 64 packed THD sequences x 64 rows: the forward
  kernel in place, the backward kernel in place, and the backward kernel into out= aliasing its
  input.
• Every run is compared with an fp32 reference, using the bf16 tolerances of the file, and bitwise
  with the first run.
• The test takes about 2 s warm.

All numbers below: B200, Triton 3.6.0, torch 2.11.0+cu130. Each run is a fresh pytest process
(python -m torch.distributed.run --nproc-per-node 1 -m pytest ...).

kernel under test	runs	result
unfixed kernel (ac100f773) + new test	20	every run: exactly mla-query-2 and dsa-indexer-2 fail, at Mismatch in backward run 0 (6k-114k mismatched elements, greatest absolute difference 4.8-6.9); the other 14 cases pass
this PR	20	every run: 16 passed
this PR with the backward barrier limited to the in-place form (... and not COPY_NOPE)	10	every run: exactly mla-query-2 and dsa-indexer-2 fail, in the out= variant (Mismatch in backward into out=input run 0)

Without the forward barrier the test still passes. With Triton 3.6 the forward kernel never splits a
head row across warps, so no test can exercise the forward barrier today.

Existing tests, compared test by test before and after:
tests/unit_tests/fusions/test_mla_yarn_rope_apply.py, and
tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py (a
remove_interleaving=True caller of fused_mla_rope_inplace).

pass	before (ac100f773)	after (this PR)
prod (-m "not experimental and not flaky_in_dev")	17 passed, 1 skipped	17 passed, 1 skipped
--experimental -m "experimental and not flaky_in_dev"	24 passed	40 passed
--experimental -m flaky_in_dev (includes the main forward/backward test of fused_mla_rope_inplace)	4 passed	4 passed

No test changed outcome. The only new tests are the 16 above.

The fixed kernel was also run on every input of the table under "When it triggers" (backward at
BLOCK_H=2, 20 calls each), on the forward at BLOCK_H=2 (8-head MLA query, 64 packed sequences,
1024 and 4096 rows) and in the scan of the table's last row (48 configurations, 160 calls each):
0 wrong calls everywhere.

Performance (cold L2, kineto)

Method:

• B200 with SM clock 1927-1965 MHz. The old module (ac100f773) and the new module run in the same
  process, interleaved ABBA, in six processes that alternate which module goes first. Every process
  ran on the same GPU, alone.
• Before every timed call, a 1 GiB scratch buffer is filled, which evicts the 126 MB L2.
• Kernel durations come from kineto CUDA activity records: 20 calls per module per process, so 120
  samples per module per cell. The table shows medians.

Autotuned path, inputs of 16384 to 65536 rows. The autotuner made the same choice in all six
processes for both modules:

case	rows	kernel	BLOCK_H	old us	new us	new/old
GLM-5 DSA indexer query 32 x 128	65536	fwd / bwd	8 / 4	222.1 / 286.9	223.6 / 287.5	1.0063 / 1.0021
GLM-5 MLA query TP1 64 x (192 + 64)	16384	fwd / bwd	8 / 4	112.8 / 145.3	113.5 / 145.6	1.0065 / 1.0024
GLM-5 MLA query TP8 8 x (192 + 64)	65536	fwd / bwd	8 / 4	58.4 / 75.1	58.7 / 75.2	1.0055 / 1.0021
64-head DSA indexer query 64 x 128 (shape only)	65536	fwd / bwd	8 / 4	440.1 / 569.5	442.9 / 570.8	1.0065 / 1.0021
same, SBHD	65536	fwd / bwd	8 / 8	426.1 / 537.3	428.8 / 532.3	1.0063 / 0.9907
same, 64 packed sequences	65536	fwd / bwd	64 / 16	1200.5 / 1474.7	1202.8 / 1470.4	1.0019 / 0.9971
DeepSeek-V3 MLA query TP1 128 x (128 + 64)	16384	fwd / bwd	8 / 4	230.2 / 293.0	231.3 / 293.7	1.0049 / 1.0022
DeepSeek-V3 MLA query TP8 16 x (128 + 64)	65536	fwd / bwd	8 / 4	118.6 / 148.7	119.1 / 149.0	1.0038 / 1.0021
control: remove_interleaving=True call, 64 x 512	65536	fwd / bwd	4 / 4	542.6 / 548.8	542.6 / 548.8	1.0000 / 1.0000
control: MLA key/value split (unchanged kernels)	16384, 65536	fwd / bwd				0.9996-0.9997

Autotuned path where the choice varies: the key and 4096-row inputs. Each module was re-tuned
40 times from scratch on each input. The expected cold time weights each module's forced-BLOCK_H
medians by how often that module chose each BLOCK_H.

case	kernel	choices old	choices new	expected old us	expected new us	new/old
DSA indexer key 1 x 128, 65536 rows	fwd	1:22 2:17 4:1	1:18 2:20 4:2	36.28	36.38	1.0028
	bwd	1:18 2:18 4:4	1:20 2:18 4:2	36.28	36.22	0.9984
DSA indexer key 1 x 128, 8192 rows, 64 sequences	fwd	1:24 2:16	1:20 2:20	38.11	38.19	1.0021
	bwd	1:26 2:14	1:36 2:4	38.04	37.96	0.9979
GLM-5 DSA indexer query 32 x 128, 4096 rows	fwd	8:40	4:1 8:39	16.96	17.12	1.0096
	bwd	4:26 8:14	4:22 8:18	20.85	20.78	0.9965
GLM-5 MLA query TP8 8 x 256, 4096 rows	fwd	4:22 8:18	4:23 8:17	7.01	7.03	1.0029
	bwd	4:15 8:25	4:25 8:15	7.07	7.20	1.0190
64-head DSA indexer query 64 x 128 (shape only), 4096 rows	fwd	8:40	8:40	30.85	31.04	1.0062
	bwd	4:36 8:4	4:35 8:5	38.69	38.75	1.0016
DeepSeek-V3 MLA query TP8 16 x 192, 4096 rows	fwd	8:40	4:1 8:39	10.45	10.45	1.0004
	bwd	4:40	4:40	12.10	12.13	1.0026

The largest value, +1.9% for the GLM-5 TP8 MLA query backward, is +0.13 us on a 7 us kernel.
About half of it comes from the new module's autotuner choosing BLOCK_H=4 more often (25 vs 15 of
40 tunings). With the old module's choice frequencies applied to both modules, the change is +1.0%.

The six performance processes also timed the key and the 4096-row inputs with the autotuner left
on. Those cells compare whatever each module's autotuner picked in that process, so they measure
the picks, not the barrier. For the GLM-5 TP8 MLA query on 4096 rows, the picks differed from
process to process, and the new/old ratio of the medians of these cells is 1.17 (forward) and 1.22
(backward):

kernel	BLOCK_H picked by the old module, processes 1-6	BLOCK_H picked by the new module, processes 1-6
fwd	16, 32, 8, 16, 8, 32	16, 8, 16, 32, 2, 2
bwd	8, 4, 2, 8, 8, 4	4, 2, 4, 8, 16, 16

The 40 fresh tunings above chose only 4 or 8 for this input, and with BLOCK_H forced, this input
is within 0.9928-1.0272 new/old at every BLOCK_H. To re-measure the autotuned path on short
inputs, pin BLOCK_H or average over many tunings, as above.

Every configuration forced. There are 28 changed-kernel cells per BLOCK_H (14 cases x
forward/backward), and 6 control cells:

BLOCK_H	median new/old	min	max
1	1.0000	0.9995	1.0025
2	1.0000	0.9984	1.0046
4	1.0016	0.9958	1.0034
8	1.0036	0.9906	1.0162
16	1.0009	0.9972	1.0060
32	1.0034	0.9928	1.0118
64	1.0004	0.9984	1.0156
128	1.0003	0.9796	1.0272 (8-head MLA query fwd, 4096 rows; 3.4x slower than the best config, never chosen)
controls	0.9997-1.0001	0.9992	1.0005

The barrier is free where it is the only barrier (BLOCK_H 1 and 2). Where Triton already converts
through shared memory, the median cost is 0.0-0.4%; the larger maxima are on small or unchosen
configurations.

Risk and compatibility

• No API change.
• remove_interleaving=True kernels compile to identical PTX.
• remove_interleaving=False kernels gain one bar.sync per program and keep every other
  instruction.
• Gradients of the MLA query RoPE and DSA indexer RoPE change only where the old kernel raced. There
  it returned wrong values that changed from run to run.
• tl.debug_barrier() is a plain CTA barrier. The kernels no longer depend on the compiler's layout
  choices, so a future Triton version that splits head rows differently is covered too.

Notes for reviewers

• The forward barrier is defensive (see above). It is free where it is the only barrier (BLOCK_H 1
  and 2).
• The test takes the kernel entry points (fused_mla_rope_inplace, mla_rope_unapply_raw) from the
  module whose autotuners it pins. So it does not touch the module-level import block. It sits next
  to the other in-place THD test.
• This keeps the PR independent of a companion fix to the same two files: 64-bit row offsets for
  inputs of 2**31 or more elements ([dev] Fix fused MLA RoPE for inputs of 2**31 or more elements #7787).
  • The two apply cleanly in either order. This was checked with git am and git cherry-pick;
    both orders give the same tree.
  • The tests pass on the combination.
• [dev] Fix fused MLA RoPE key/value split for non-power-of-two head dims #7803 (key/value split for non-power-of-two head dims) and [dev] Use binary search for the THD row lookup in fused MLA RoPE #7806 (binary search for the THD row
  lookup) also change these two files. All four PRs apply without conflicts, and the tests pass
  on the combined tree.
• Please merge this PR before [dev] Use binary search for the THD row lookup in fused MLA RoPE #7806. [dev] Use binary search for the THD row lookup in fused MLA RoPE #7806 does not touch the in-place kernels' stores, but it
  changes the timing of what runs before them: without this PR's barrier, the race then also
  shows up on the input of test_forward_backward_for_q[*-thd], the flaky test of Flaky: TestFusedApplyMLARope::test_forward_backward_for_q[thd] backward mismatch exceeds bf16 tolerances #4640: 18-74
  wrong backward calls of 200 at BLOCK_H=2, against 0 of 200 on dev. With this PR, 0 of 200
  with and without [dev] Use binary search for the THD row lookup in fused MLA RoPE #7806.

Pre-checks

• I have added relevant unit tests (test_mla_rope_inplace_channel_permutation_is_race_free, which fails without the barrier)
• I have added relevant functional tests (not applicable)
• I have added proper typing to my code (no new Python API)
• I have added relevant documentation (not applicable)
• I have run the autoformatter.sh on my PR (its checks in check mode on the two changed files,
  with the repository settings: black and isort leave them unchanged, ruff passes,
  pylint rates them 10.00/10, and git diff --check is clean)

Linked issue: related to #4640 (see Related work).

@hxbai @ghadiaravi13 @buptzyb, could you please take a look?

[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.

[yaox12]

/ok to test f3b7f60

[yaox12]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

Strict review: approved

I found no problems in this PR's own changes. I'm approving it. The one item below was already in the code before this PR and is not part of the diff, so it does not block the merge.

This review is source reasoning only. I did not compile or run any kernels or tests.

What I checked in the diff

fused_mla_yarn_rope_apply.py

• Forward barrier (L195-L201). It sits only in the REMOVE_INTERLEAVING=False branch, which is the only path where a head's stores (halves i, d/2+i) land on channels that other lanes loaded (pairs 2i, 2i+1).
• Backward barrier (L319-L324). It is the mirror case and correctly does not depend on COPY_NOPE. mla_rope_unapply_raw only checks out's shape, stride and dtype (L467-L469), so out=t aliasing is reachable.
• The nope copy is safe. In both ROPE_FIRST layouts, the COPY_NOPE region (L326-L330) does not overlap the rope region. Each element is loaded and stored at the same place, so it needs no barrier.
• No deadlock risk. Both barriers are inside tl.constexpr branches, and nothing before them diverges: there is no early return, and the cu_seqlens walk is uniform scalar control flow.
• Ordering is sufficient. bar.sync makes each thread's earlier global loads complete for every participating thread. Different programs touch disjoint (row, head-block) tiles (but see the item below), so a CTA barrier is enough here.
• The remove_interleaving=True paths are unaffected. They store every value back where it was loaded.
• No API changes: no signature, argument or launch-grid change.

test_mla_yarn_rope_apply.py

• The fp32 reference matches the kernel math. It follows _mla_rope_fwd_inplace_kernel (x_left = x1·cos_l − x2·sin_l, x_right = x2·cos_r + x1·sin_r, stored as halves) and _mla_rope_bwd_kernel (stored as interleaved pairs) term by term.
• The THD positions are right. arange(rows) % seq_len matches _get_thd_token_idx for equal-length sequences with cp_size=1.
• Pinning BLOCK_H works. Replacing configs with a single entry bypasses Triton's benchmark and cache, monkeypatch restores the original configs, and the default num_warps=4 matches production.
• Markers follow the file's conventions (experimental and internal, the same as the neighbouring in-place THD test).
• No determinism-replay test is needed. The mcore-testing rule for a replay test under tests/unit_tests/determinism/kernels/ does not apply on dev: that directory and tools/check_kernel_determinism_coverage.py exist in neither tree.

Findings by severity

CRITICAL: 0 · IMPORTANT: 0 in the diff (1 that predates this PR, below) · SUGGESTION: 0

---

[IMPORTANT Correctness] Head-block mask ignores pid_head (predates this PR, not in the diff, non-blocking)

This is not introduced here: base has the same lines (fwd L178-L182, bwd L297-L300). I'm raising it because it limits the PR's claim that these kernels are race-free. It also hits the very tl.store(..., mask=mask) calls the new barriers protect.

The bug. Both in-place kernels add pid_head * BLOCK_H * stride_x_nheads to the base pointer, then build the mask from the local head index only:

• forward: L178-L182
• backward: L295-L302

mask = x_off < head_num * stride_x_nheads   # x_off = arange(BLOCK_H)*stride + rope_offset

The grid is cdiv(nheads, BLOCK_H) (L402, L472). So when head_num > BLOCK_H and head_num % BLOCK_H != 0, the last head block is never masked.

Smallest trigger. head_num=48, BLOCK_H=32 (the autotuner benchmarks every config, so it always runs this one):

• Program (m, 1) processes global heads 32..63.
• With a contiguous [rows, 48, D] layout, heads 48..63 are really row m+1's heads 0..15. Program (m, 1) rotates them with row m's cos/sin.
• Meanwhile program (m+1, 0) reads and writes those same channels. The result is wrong and changes from run to run.
• For the last row, the writes go past the end of the tensor.
• The same happens for 24 heads at BLOCK_H=16, 40 heads at 16/32, and similar cases.

Expected behaviour: heads >= head_num are fully masked.

The concat kernel already builds the mask correctly (L771-L772). The CTA barrier added here cannot fix overlap between different programs.

Who is affected. The common per-rank head counts (16/32/64/128, power-of-two TP) are always fully covered, so they are safe. Any non-power-of-two per-rank head count that is not a multiple of the chosen BLOCK_H is not.

Fix. Use the global head index in both in-place kernels. This is a two-line change per kernel, in this PR or a follow-up; it was not compiled or run here, so please verify it.

# _mla_rope_fwd_inplace_kernel, replace L182; _mla_rope_bwd_kernel, replace L302
head_idx = pid_head * BLOCK_H + tl.arange(0, BLOCK_H)
mask = (head_idx < head_num)[:, None]

The invariant to restore: a program touches only global heads < head_num of its own row.

What the fix keeps.

• The mask shape stays [BLOCK_H, 1], so the rope loads/stores and the COPY_NOPE mask & (nope_idx < nope_dim) broadcast exactly as before.
• For pid_head == 0 the new mask equals the old one, because 0 <= rope_offset < stride_x_nheads.
• So every currently-correct case (power-of-two heads at every BLOCK_H, _test_fused_mla_rope_inplace with 32 heads, the new race test) computes the same mask.
• Keep the barriers added by this PR.

Acceptance (proposed test).

• Add pytest.param(48, 192, False, id="mla-query-48-heads") and pytest.param(24, 128, True, id="dsa-indexer-24-heads") to test_mla_rope_inplace_channel_permutation_is_race_free.
• Allocate x as buf[:rows] from buf = torch.empty(rows + 1, H, D) and assert the pad row is unchanged. That way the last row's out-of-bounds writes show up as a mismatch instead of an illegal-address error.
• Expected today: failures at block_h=32 (48 heads) and block_h=16 (24 heads).
• Expected after the fix: every block_h passes, and the 16/32-head cases still pass.

The KV-split kernels use the same masking pattern (L1106-L1108, L1216-L1218, L1235-L1237). A follow-up should cover those too.

---

Overall risk: low

The change is minimal and confined to constexpr branches. It fixes silent, run-to-run-varying wrong gradients in the MLA query and DSA indexer RoPE backward, and the PR's measurements show negligible cost. The issue above predates this PR and does not get worse with it.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by claude