[Heisenberg-Yin]

Thanks @TKONIY for the contribution to this PR!

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

Summary

With apply_rope_fusion=True, MLASelfAttention builds its key and value with the fused MLA RoPE
key/value split, fused_apply_mla_rope_for_kv. This is an alias of fused_mla_rope_kv_split in
megatron/core/fusions/fused_mla_yarn_rope_apply.py. Its two Triton kernels compile only when
qk_head_dim, v_head_dim and half of qk_pos_emb_head_dim are powers of two.

GLM-5's MLA head dims are qk_head_dim 192, v_head_dim 256 and qk_pos_emb_head_dim 64. With
them, the first training forward pass of an MLA layer stops with:

triton.compiler.errors.CompilationError: at 62:24:
    k_in_off = kv_off + tl.arange(0, k_dim)[None, :]
                        ^
arange's range must be a power of 2

This PR pads the three widths to the next power of two inside the two kernels and masks the padding
lanes, as fused_mla_rope_concat in the same file already does. The masks are built only when a
width is actually padded. Power-of-two configurations, such as DeepSeek-V3's, therefore compile to
exactly the same kernels as before: identical PTX and SASS. A new GPU test covers GLM-5's dims and a
case in which all three widths are padded.

Problem

MLASelfAttention.get_query_key_value_tensors (multi_latent_attention.py) takes the fused path
when should_use_fused_mla_rope(config) is true:

key, value = fused_apply_mla_rope_for_kv(
    kv,                                # [tokens, heads, qk_head_dim + v_head_dim]
    k_pos_emb,                         # [tokens, 1, qk_pos_emb_head_dim]
    rotary_pos_cos,
    rotary_pos_sin,
    self.config.qk_pos_emb_head_dim,   # emb_dim
    self.config.qk_head_dim,           # k_dim
    self.config.v_head_dim,            # v_dim
    cu_seqlens_kv,
    cp_rank,
    cp_size,
)

With GLM-5's head dims, this call raises the CompilationError above instead of returning key and
value. The query RoPE just before it (fused_apply_mla_rope_for_q) works with the same dims.
Nothing checks the dims beforehand.

Reproduced on dev at ac100f773 (B200, Triton 3.6.0, torch 2.11):

• The public entry points fused_mla_rope_kv_split and fused_apply_mla_rope_for_kv with GLM-5's
  dims, at 64 heads (TP 1) and 8 heads (TP 8), in SBHD and in THD format: the forward kernel fails at
  line 62 of _mla_rope_fwd_kv_split_kernel. _FusedMLARoPEKVSplit.backward called on its own fails
  at line 54 of _mla_rope_bwd_kv_split_kernel, dk_out_off = dkv_off + tl.arange(0, k_dim)[None, :].
• An MLASelfAttention (Transformer Engine spec, bf16, 8 heads) with GLM-5's head dims and
  apply_rope_fusion=True, in SBHD and in packed THD format: the forward pass fails in
  qkv_up_proj_and_rope_apply, at the fused_apply_mla_rope_for_kv call.

Root cause

Mechanism

The key/value split kernels take the three widths as tl.constexpr and build their tiles with
tl.arange over them:

# _mla_rope_fwd_kv_split_kernel, before this PR
cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2))
...
k_in_off = kv_off + tl.arange(0, k_dim)[None, :]
v_in_off = kv_off + k_dim + tl.arange(0, v_dim)[None, :]

The backward kernel has the same pattern: dk_out_off, dv_out_off, dk_in_off, dv_in_off, and
the rotary accumulators tl.zeros((BLOCK_H, emb_dim // 2)).

Triton tensor shapes must be powers of two. tl.arange(0, n) checks this when the kernel is
compiled, which happens at its first launch (here, the autotuner's first benchmark run). For any
other n it raises ValueError("arange's range must be a power of 2"), which Triton reports as a
CompilationError.

Worked example: GLM-5 at TP 8. The inputs are kv [tokens, 8, 192 + 256] and k_pos_emb
[tokens, 1, 64]; the outputs are key [tokens, 8, 192 + 64] and value [tokens, 8, 256].

tile	GLM-5	power of two?	DeepSeek-V3
key channels without RoPE, tl.arange(0, k_dim)	192	no (3 x 64)	128
value channels, tl.arange(0, v_dim)	256	yes	128
rotary pairs, tl.arange(0, emb_dim // 2)	32	yes	32

So the forward kernel fails at its key tile, and so does the backward kernel. DeepSeek-V3-style dims
are all powers of two, which is why this has not shown up before.

The other kernels of the file handle GLM-5's dims (checked on ac100f773, forward and backward):

• fused_mla_rope_concat (absorbed MLA) already pads its widths with NOPE_BLOCK and ROT_BLOCK
  and masks the rest.
• The in-place query kernels (fused_mla_rope_inplace: the MLA query, 64 x (192 + 64), and the DSA
  indexer query and key) skip the channels without RoPE with a runtime offset; when the backward
  copies them into out=, it already pads. Their only unpadded head-width extent is
  emb_dim // 2 = 32.

How it shows up

• It is a hard failure, not a silent one. The model builds; the first forward pass raises the
  CompilationError from inside the Triton compiler.
• No configuration check catches it earlier. should_use_fused_mla_rope looks only at
  apply_rope_fusion, rope_type and rotary_percent.
• The only workaround is to turn RoPE fusion off (--no-rope-fusion, or apply_rope_fusion=False),
  which runs the unfused PyTorch path instead. For GLM-5's shapes that path is several times slower;
  see Performance.

When it triggers

All of the following:

1. The layer is the standard MLA attention module, MLASelfAttention. This includes its
   mla_down_proj_fusion subclass FusedMLASelfAttention and the MLA layers of hybrid models, which
   share get_query_key_value_tensors.
2. Training. MLA RoPE fusion is not supported in inference; the code asserts
   inference_context is None.
3. should_use_fused_mla_rope(config) is true: apply_rope_fusion=True and either rope_type
"yarn" or "rope" with rotary_percent 1.0. In the Megatron-LM training scripts,
   apply_rope_fusion is on by default with --position-embedding-type rope (MLA then uses
   rope_type "yarn" unless --rope-type says otherwise); --no-rope-fusion turns it off, and
   any other --position-embedding-type except mrope turns it off too.
4. qk_head_dim or v_head_dim is not a power of two, as with GLM-5's qk_head_dim 192.

These are per-head widths, and parallelism never changes them:

• Tensor parallelism divides the heads: each GPU's kv is [tokens, 64 / TP, 192 + 256] for
  GLM-5, so [tokens, 64, 448] at TP 1 and [tokens, 8, 448] at TP 8.
• Context parallelism divides the tokens of each sequence. Sequence parallelism only splits the
  activations outside attention; inside it, kv covers all the tokens of the CP rank.
• Pipeline parallelism splits the layers.

The head width of every tile stays 192 (key), 256 (value) and 32 rotary pairs. GLM-5's dims
therefore fail at every parallel layout, and in both the SBHD and the packed THD format.

Not affected: DSA. With experimental_attention_variant="dsa" (GLM-5's own configuration), the
attention module is AbsorbedMLASelfAttention. It applies RoPE through fused_mla_rope_concat and
never calls the key/value split. So the failure hits GLM-5's head dims on the standard MLA module,
for example a GLM-5-shaped model trained with dense MLA attention.

How it was found

It was found while timing the fused MLA RoPE kernels on GLM-5 shapes: the key/value split case with
GLM-5's head dims did not compile.

Fix

• The two launches in _FusedMLARoPEKVSplit pass three new constexprs: K_BLOCK, V_BLOCK and
  ROT_BLOCK, the next powers of two of k_dim, v_dim and emb_dim // 2
  (triton.next_power_of_2).

• Both kernels range over these blocks instead of the widths and mask the lanes past each width. All
  offsets still use the real widths.

  rot_mask = None if ROT_BLOCK == emb_dim // 2 else tl.arange(0, ROT_BLOCK) < emb_dim // 2
  ...
  k_in_off = kv_off + tl.arange(0, K_BLOCK)[None, :]
  v_in_off = kv_off + k_dim + tl.arange(0, V_BLOCK)[None, :]
  k_mask = mask if K_BLOCK == k_dim else mask & (tl.arange(0, K_BLOCK)[None, :] < k_dim)
  v_mask = mask if V_BLOCK == v_dim else mask & (tl.arange(0, V_BLOCK)[None, :] < v_dim)
  k = tl.load(KV_ptr + k_in_off, mask=k_mask)

• The masks cover every access whose extent is a padded width:

  • forward: the cos/sin and k_pos_emb loads, the key and value loads and stores, and the
    rotary stores;
  • backward: the dK/dV loads, the dKV stores, the per-head rotary loads and the dK_POS_EMB
    stores. The rotary accumulators become [BLOCK_H, ROT_BLOCK].

Why it is correct and complete.

• Every tl.arange in the two kernels now has a power-of-two extent: K_BLOCK, V_BLOCK,
  ROT_BLOCK, or BLOCK_H from the autotuner configs. Any width compiles.
• Padded lanes are never loaded or stored, so no access goes past a width. compute-sanitizer
  memcheck, with the caching allocator off so that every tensor is its own allocation, reports no
  errors on padded widths at every BLOCK_H.
• Padded lanes never mix with valid lanes. The arithmetic is elementwise along the channel axis. The
  backward's only reduction, the rotary gradient summed over the heads, runs along the head axis, so
  a padded column never reaches a stored one.
• The valid lanes compute exactly what the old kernels compute for power-of-two widths. On GLM-5's
  widths (and others), the new kernels give bitwise the same key, value and gradients as the old
  kernels on the same data zero-padded to powers of two (192 + 256 run as 256 + 256), at every
  BLOCK_H.
• The rotary width is padded the same way. The key/value split therefore no longer needs any width
  to be a power of two.

Why the masks are built only for padded widths. An unconditional mask folds to "true" for
power-of-two widths, but it still turns the unmasked loads of cos, sin and k_pos_emb into
predicated loads, and the compiler then schedules the kernel differently. Compiled for B200, an
always-masked variant changed the SASS of 50 of the 256 power-of-two variants checked below. In 6 of
them (DeepSeek-V3 dims, forward kernel, BLOCK_H 64 and 128), it also changed the bf16 FMA
contraction of the rotary arithmetic (HFMA2 16 -> 15, HMUL2 4 -> 5), which can change results
bitwise. Building the masks only for padded widths keeps power-of-two kernels identical.

What it does not change.

• Power-of-two widths compile to identical PTX and SASS: 256 of 256 variants (DeepSeek-V3's 128 + 128
  with 64 rotary channels at 128 and 16 heads, the MLA unit tests' 16 + 16 with 8 at 8 heads, and
  64 + 512 with 64 at 1 head; forward and backward kernels; SBHD and THD; both remove_interleaving
  values; every BLOCK_H). Their outputs are bitwise the same as before (96 of 96 forced-BLOCK_H
  runs), and so is their speed.
• No API, argument or autotuner change.
• The in-place query kernels are unchanged. MLA calls them with the same rotary width, so a model
  whose qk_pos_emb_head_dim / 2 is not a power of two would still fail there. The 64 rotary
  channels of DeepSeek-V3 and GLM-5 give 32.

Tests and results

New GPU test
tests/unit_tests/fusions/test_mla_yarn_rope_apply.py::TestFusedApplyMLARope::test_kv_split_non_power_of_two_head_dims,
with the same markers as its neighbours (experimental, internal, skipped without CUDA or with
torch < 2.5):

• Two sets of dims:
  • glm5-tp8: 8 heads, qk_head_dim 192, v_head_dim 256, qk_pos_emb_head_dim 64;
  • k96-v160-rope48: 4 heads, with the key, value and rotary widths (96, 160 and 24 pairs) all
    padded.
• Each in SBHD and THD format, with remove_interleaving False and True: 8 cases.
• Each case runs the key/value split forward and backward and compares it with the unfused PyTorch
  path (apply_rotary_pos_emb plus split and concat, with autograd), using the bf16 tolerances of
  the file.
• To run with other dims, the existing helper _test_fused_mla_rope_kv_split now hands its body to
  _check_fused_mla_rope_kv_split(input_format, remove_interleaving, num_heads, k_dim, v_dim, emb_dim). The helper keeps its signature and its DeepSeek-V3 dims, so the existing test is
  unchanged.

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

tree	runs	result
ac100f773 + new test	3	every run: 8 failed, all with CompilationError: arange's range must be a power of 2 (glm5-tp8 at the key tile, k96-v160-rope48 at the rotary tile)
this PR	3	every run: 8 passed (about 10 s warm)

Existing tests, compared test by test before and after: test_mla_yarn_rope_apply.py and
test_csa_cp_utils.py (a remove_interleaving=True caller of the same file).

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	32 passed (the 8 new)
--experimental -m flaky_in_dev	4 passed	4 passed

No test changed outcome.

Additional checks, outside the unit tests:

• MLA module. An MLASelfAttention with GLM-5's head dims (8 and 64 heads, dropout off) now
  runs forward and backward with apply_rope_fusion=True, in SBHD and THD, with rope and yarn
  RoPE. The output and every gradient match the same module and inputs with
  apply_rope_fusion=False to within 0.72% of their largest value. The fused path keeps
  cos/sin in bf16, so the two paths differ by bf16 rounding; DeepSeek-V3-style dims on unmodified
  dev differ by the same amount (0.72%).
• Every BLOCK_H forced. 8 sets of dims × SBHD/THD × remove_interleaving × 8 BLOCK_H values
  = 256 runs, compared with an fp32 reference: all pass. The dims include GLM-5's at 8 and 64 heads,
  and sets in which the key, the value or the rotary width needs padding.
  • The copied channels are bitwise equal, and no output element is left unwritten.
  • The rotary outputs stay within a bf16 error bound that scales with the size of the products.
    They use at most 60% of it, the same as the unmodified kernels on power-of-two dims.
• Old kernels on zero-padded inputs. On non-power-of-two widths, the new kernels give bitwise the
  same key, value and gradients as the old kernels on the same data zero-padded to powers of two.
  This holds in 160 of 160 runs: 192 + 256 at 8 and 64 heads, 96 + 160, 72 + 48 and 192 + 128, each
  in both layouts, for both remove_interleaving values and every BLOCK_H.
• compute-sanitizer memcheck. Padded widths, every BLOCK_H, forward and backward, with the
  caching allocator off: 0 errors.

Performance

Cold L2 kernel times on B200:

• Method:
  • before every timed call, a 1 GiB buffer is written, which evicts the 126 MB L2;
  • kernel times come from kineto;
  • 6 processes, alternating which module runs first; 120 samples per module per cell; medians;
  • SM clock 1950-1965 MHz, and no other process on the GPU.
• Inputs: THD, bf16, one sequence. BLOCK_H is either forced or chosen by the autotuner ("auto").

Power-of-two dims: unchanged. The code is identical, so this only confirms that timing is too.

case	kernel	auto pick (old / new, all 6 runs)	before us	after us	after / before	8 forced BLOCK_H values
DeepSeek-V3 TP 1: 128 heads x (128 + 128), 16384 tokens	forward	128 / 128	361.3	361.2	0.9997	0.9998-1.0003
	backward	64 / 64	362.9	362.9	1.0000	0.9998-1.0002
DeepSeek-V3 TP 8: 16 heads x (128 + 128), 65536 tokens	forward	16 / 16	187.3	187.2	0.9996	0.9998-1.0007
	backward	16 / 16	184.4	184.4	1.0002	0.9996-1.0001

GLM-5's dims: now fused. The old kernels do not compile, so the comparison is with the path
that apply_rope_fusion=False runs: split, apply_rotary_pos_emb on the shared k_pos_emb,
expand, concatenate, contiguous.

case	fused forward us (auto pick)	fused backward us (auto pick)	effective bandwidth	unfused / fused, forward	unfused / fused, forward + backward
GLM-5 TP 1: 64 heads x (192 + 256), 16384 tokens	302.7 (64)	299.1 (64)	6.7-6.8 TB/s	6.3x (1 vs 31 kernels)	6.5x (2 vs 45 kernels)
GLM-5 TP 8: 8 heads x (192 + 256), 65536 tokens	152.4 (16)	153.2 (8)	6.7-6.8 TB/s	7.1x	7.1x

• The padding lanes are masked and generate no memory traffic. The padded kernels reach the same
  effective DRAM bandwidth as DeepSeek-V3's dims at the same TP: 6.6-6.7 TB/s.
• The largest configuration, BLOCK_H=128, spills registers with GLM-5's 256 + 256 lanes per head,
  as the old kernels do for 256 + 256 widths. It is far slower than the best configuration (for
  example 935 vs 152 us for the TP 8 forward), and the autotuner never picked it.

Risk and compatibility

• No API change. Power-of-two head dims compile to the same PTX and SASS as before, so their results
  and speed do not change.
• Head dims that are not powers of two used to fail at compile time. They now work.
• The padded tiles are larger than the real widths. For GLM-5's dims, the largest configuration
  (BLOCK_H=128, 256 + 256 lanes per head) spills registers; the autotuner does not pick it (see
  Performance).
• The kernels use only standard Triton features: masked loads and stores, and constexpr
  conditional expressions (x if CONSTEXPR_COND else y, rot_mask is None), as the file already
  does for position_ids is not None.

Notes for reviewers

• The constexpr choice between "no mask" and "mask" is deliberate. An always-on mask is simpler but
  changes the compiled power-of-two kernels, including the FMA contraction in some of them (see Fix).

• The test helper is extended below its first lines rather than given new parameters. [dev] Fix a race in in-place fused MLA RoPE with remove_interleaving=False #7789 (the
  in-place RoPE race fix) adds its test right above _test_fused_mla_rope_kv_split, and this way
  the two PRs do not touch neighbouring lines.

• Three other pending changes touch the same two files: [dev] Fix fused MLA RoPE for inputs of 2**31 or more elements #7787 (64-bit row offsets for inputs of
  2**31 or more elements), [dev] Fix a race in in-place fused MLA RoPE with remove_interleaving=False #7789, and [dev] Use binary search for the THD row lookup in fused MLA RoPE #7806 (a binary search for the THD row lookup). All four
  apply in any order without conflicts (git am without -3, git cherry-pick, git merge-tree), and the
  result does not depend on the order. On top of the 64-bit offset fix, this PR's kernels are again
  identical for power-of-two widths.

• The open head-mask fixes Offset the per-head bounds mask by the head block in fused MLA RoPE #6890, Fix out-of-bounds head masks in fused MLA YaRN RoPE kernels #7104 and test(mla): check every BLOCK_H candidate of the fused MLA RoPE kernels on a partial head block #7770 (against main; BLOCK_H not dividing the head
  count) rewrite the mask = ... line just above this PR's first change in both key/value split
  kernels and in the backward's per-head loop. Offset the per-head bounds mask by the head block in fused MLA RoPE #6890 and test(mla): check every BLOCK_H candidate of the fused MLA RoPE kernels on a partial head block #7770 also add other=0.0 to the four
  rotary-gradient loads that this PR re-masks. Whichever lands on dev second needs a mechanical
  rebase:

  • keep the new head mask and this PR's lines below it;
  • write those four loads as mask=x_mask, other=0.0.

  k_mask, v_mask and x_mask are built on mask, so they inherit the new head mask. This was
  checked on a port of test(mla): check every BLOCK_H candidate of the fused MLA RoPE kernels on a partial head block #7770 to dev with this PR on top: every BLOCK_H matches an fp64 reference
  with partial head blocks (12 and 127 heads) and padded widths, and the kernel tests of both PRs
  pass. This PR's tests use 8 and 4 heads, so they do not depend on the head-mask fix.

Pre-checks

• I have added relevant unit tests (test_kv_split_non_power_of_two_head_dims)
• 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: none (bug fix).

@hxbai @yuzhongw-nvidia @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 40b1cd3

[yaox12]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

Strict review passed: no significant issues found. LGTM.

Findings: 0 critical, 0 important, 0 suggestions.

Scope: This was a broad, strict review of both changed files. I read the two key/value split kernels and their launches in fused_mla_yarn_rope_apply.py and the test changes in test_mla_yarn_rope_apply.py. I also checked which files use these functions. I can't run tests or compile kernels here, so everything below comes from reading the source.

Checks on the kernel changes (reading only):

• Out-of-bounds accesses. Every padded lane that would read or write past its real width is masked:
  • k_mask covers the key lanes, which would otherwise spill into the value channels.
  • v_mask covers the value lanes, which would otherwise spill into the next head.
  • x_mask covers the rotary stores, where x_right_off = x_left_off + emb_dim // 2 would otherwise write past k_dim + emb_dim into the next head.
  • rot_mask covers the cos/sin/k_pos_emb loads and the dK_POS_EMB stores.
  • All offsets still use the real k_dim/v_dim/emb_dim.
• Backward reduction. The only reduction (L1269-L1270) sums over the head axis. Padded columns hold undefined values, but they never mix with valid columns and are never stored. Accumulation stays in fp32 and is cast to the dEMB dtype, as before.
• Power-of-two dims compile to the same kernel. ROT_BLOCK == emb_dim // 2, K_BLOCK == k_dim and V_BLOCK == v_dim are comparisons between constexprs, so the x if c else y and rot_mask is None choices are made at compile time. For power-of-two widths, mask is used unchanged and the cos/sin/k_pos_emb loads stay unmasked, so DeepSeek-V3-style dims build the same kernel as before. Inside the runtime if pid_head == 0: block, rot_mask and x_mask are only created and used locally. The variable that was already used before the block (mask) is reassigned with the same [BLOCK_H, 1] i1 type as before.
• Compatibility. The public API of fused_mla_rope_kv_split / fused_apply_mla_rope_for_kv is unchanged, and so are the autotuner configs and key. K_BLOCK, V_BLOCK and ROT_BLOCK are derived from values already in the key. The private kernels are launched only from _FusedMLARoPEKVSplit. Other users of this file (cp_utils.py, csa.py, deepseek_v4_hybrid_attention.py) import only the in-place and out-of-place helpers. multi_latent_attention.py goes through the alias, which is unchanged.
• New identifiers. K_BLOCK, V_BLOCK, ROT_BLOCK, k_mask, v_mask, rot_mask, x_mask and _check_fused_mla_rope_kv_split are all actually used; none is dead.

Tests: The existing _test_fused_mla_rope_kv_split keeps its 32 × (128 + 128) / 64 case. The new parametrized test uses the same class-level markers as its neighbours (experimental, internal, CUDA and torch-version skips). Between them, the two new parameter sets exercise padding of all three widths in SBHD and THD, with both remove_interleaving values.

Context, not counted as findings:

• Two head-mask issues in these kernels predate this PR. The block-relative head mask is wrong when BLOCK_H doesn't divide the head count, and the rotary-gradient loads have no other=0.0. The PR description acknowledges both and leaves them to the separate head-mask PRs; this PR doesn't make them worse.
• The MLA query in-place kernels still require emb_dim // 2 to be a power of two. The description already calls this out of scope.
• The rubric's bit-exact kernel replay test requirement doesn't apply to this branch. The base dev tree has neither tests/unit_tests/determinism/kernels/ nor tools/check_kernel_determinism_coverage.py.
• The rubric also points to a style-guide.md that doesn't exist in the base tree, so no style findings came from it.

Overall risk: low. The change is limited to two kernels and only affects head dims that are not powers of two, which previously failed to compile.

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

Review generated by claude