NVIDIA / NVIDIA/Megatron-LM

MLA pads V on the thd path unconditionally, which costs up to 87% of a step at long sequence

Open
#6,240 0 comments 0 reactions 0 assignees View on GitHub
community-request
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Is your feature request related to a problem? Please describe.**

On the packed (`thd`) path, `_prepare_mla_core_attention_value` pads V up to the QK width whenever the two differ, and `_trim_mla_core_attention_output` trims the result back. For MLA the widths always differ — the RoPE dims ride on q/k but not on v — so every MLA `thd` layer goes through the pad.

The pad and the trim cancel out. What they cost is a V that is 50% wider than it needs to be through the attention kernel, plus the pad and the trim, and on some stacks a fallback to a slower backend as well.

Measured on a small MLA model with DeepSeek-V2-Lite attention dims (`qk_head_dim` 128, `qk_pos_emb_head_dim` 64, `v_head_dim` 128, `kv_lora_rank` 512, `q_lora_rank` None), 8 layers, hidden 2048, `ffn_hidden_size` 10944, bf16, `attention_dropout` 0. Whole training step — forward, backward, Adam — median of 40 steps after 15 warmup, two rounds per cell with the arm order reversed between them:

| tokens | packs | tokens/seq | H200 native | H200 padded | H200 | B300 native | B300 padded | B300 |
|---|---|---|---|---|---|---|---|---|
| 4096 | 4 | 1024 | 57.7, 58.0 | 61.0, 61.3 | **5.7%** | 57.1, 58.0 | 57.8, 58.3 | — |
| 8192 | 4 | 2048 | 85.3, 85.8 | 91.4, 92.0 | **7.2%** | 62.5, 63.6 | 68.6, 69.9 | **10.0%** |
| 16384 | 4 | 4096 | 156.1, 157.3 | 168.9, 169.1 | **7.8%** | 97.2, 97.3 | 125.4, 125.5 | **29.0%** |
| 16384 | 1 | 16384 | 214.7, 215.0 | 239.9, 240.0 | **11.7%** | 116.4, 116.5 | 218.3, 218.3 | **87.4%** |

Both stacks improve monotonically with sequence length, which is what you would expect if the saving is in attention: attention is O(s²), so its share of the step grows and the wasted width costs more. On B300 the 4096-token cell is a wash — the two rounds disagree on which arm is faster, so it is noise rather than a small win.

The two shapes differ because the backends differ. On H200 both arms run cuDNN fused attention, so the gap is the wasted width alone. On B300 the padded shape gets no fused backend at all and falls back to FlashAttention:

```
# native V (192/128)
Disabling FlashAttention 2 as it does not support MLA.
Selected backend = FusedAttention (sub-backend 1)

# padded V (192/192)
Disabling FusedAttention as no backend supports the provided input
Selected backend = FlashAttention (2.7.4)
```

That message comes from `tex.get_fused_attn_backend` returning `No_Backend` for 192/192 while accepting 192/128 — so padding V to match QK does not widen backend support there, it narrows it. The two stacks differ in both compute capability and TE version, so this data does not isolate which accounts for that difference.

Tag @NVIDIA/mcore-oncall

**Describe the solution you'd like**

Let a model opt out of the pad. `need_v_pad` already has an escape hatch for `experimental_attention_variant`; this adds one more condition rather than a new mechanism.

A PR is attached that adds `MLATransformerConfig.mla_native_v_head_dim`, defaulting to `False` so nothing changes unless a model asks for it.

**Describe alternatives you've considered**

*Deciding automatically instead of by config.* Megatron-Core could skip the pad whenever the backend can take native V. That is the better end state, but it means querying TE's backend selection from Megatron-Core, and the answer varies with the stack — the two measured above already disagree about whether 192/192 is servable at all. Happy to pursue this instead if you prefer; the config flag is proposed because it is small and reviewable.

*Leaving it to FlashAttention 3.* FA3 does support mismatched head dims — TE's `_is_fa3_supported` admits `head_dim_qk` in `(128, 192]` with `head_dim_v` in `(96, 128]`, which covers this geometry. It does not address this. With FA3 3.0.0 installed and eligible on H200, TE still disables it in favour of fused attention ("Disabling FlashAttention to give FusedAttention preference on Hopper+ for performance reasons"), and TE gates FA3 to sm90 so it is unavailable on Blackwell.

*Keeping the pad as a fallback.* The pad is useful when fused attention cannot serve the layer, since it lets FlashAttention take over instead of falling to `UnfusedDotProductAttention`. That is why this is opt-in rather than a change of default.

**Additional context**

Pad-then-trim is an identity bit-exactly, not merely within tolerance: the padded columns of the attention output are sums of zeros, so the trim discards nothing, and every backward input matches. Verified in fp32 and bf16 across several shapes.

Where the two arms end up on different backends — the B300 stack — outputs and gradients agree to bf16 rounding rather than exactly: worst relative drift 5.7e-3 across forward, dq, dk and dv. On H200 both arms run the same kernel.

One measurement note, in case it saves someone else the detour: Megatron-Core defaults `attention_dropout` to 0.1, and on B300 a `thd` request with dropout is served by cuDNN's composite engine, where dropout mask generation costs several times the attention itself. Leaving that default in place reverses the comparison above. The numbers here use `attention_dropout=0`, which is what the models this geometry comes from actually set. Filed separately as NVIDIA/TransformerEngine#3312.

Contributor guide

Open the contributing guide

Research direction

Start at _prepare_mla_core_attention_value and _trim_mla_core_attention_output, then trace need_v_pad and MLATransformerConfig.mla_native_v_head_dim. Compare the native and padded thd paths and review the attached PR's coverage; done means the opt-in path preserves forward and backward results while avoiding unnecessary V padding.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, machine-learning, performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.