[Performance] Balance the latent-MQA indexer's CP rows with a dual-chunk swap

[xxyux]

PR Category

Performance Optimization

PR Types

Improvements

Description

Under contiguous CP the latent-MQA indexer scores each rank's own row slice, so with a causal mask its cost grows with the rank index: measured at 256k/cp16, 2.2ms on cp0 against 66.8ms on cp15 per layer per pass (a 30x spread), which makes the slowest rank do 1.94x the average while every other rank waits for it at the next collective. Only the indexer has this shape -- attention already runs a fixed index_topk + window columns per row, and the indexer backward a fixed index_topk, both measured flat across ranks.

mqa_indexer_cp_mode="dualchunk_p2p" splits the global sequence into 2 * cp_size chunks and has rank r score chunks (2r, 2*cp_size-1-2r) instead of its own (2r, 2r+1). The ids sum to 2*cp_size-1 everywhere, and a causal row's candidate count grows linearly with its global position, so the work is equal on every rank.

Three properties keep the change contained:

• Only the indexer's per-row inputs move. query and the layer output stay put, which at s_local=16384/h=64 is 2.3GB of traffic not incurred.
• Rank r keeps the chunk contiguous CP already gave it and swaps the other with rank cp_size-1-r, so the exchange is one point-to-point sendrecv per tensor rather than an all-to-all, and it is an involution -- the same call undoes it.
• The permutation lives entirely inside paddle.no_grad() on detached inputs. The indexer gradient reaches the weights through the loss scaler applied to the unpermuted tensors, so nothing here has to be differentiable and the backward is byte-identical between the two modes.

The cuDNN indexer call is split in two because the kernel learns where its rows sit globally from q_causal_offsets, one scalar per batch, so a single affine map cannot describe two disjoint segments. Each chunk is internally contiguous and carries its own seq_offset -- the same shape as the query tiling the dense path already does. The results are concatenated back before the loss scaler sees them: applying it per chunk would halve target.shape[1], where its backward reads the row-count denominator, and double the gradient.

The layer output is bit-identical either way, which is what the new CP test asserts. That equality is only possible because the per-(row, column) score is itself bit-identical: a CTA covers q_stage * q_tokens_per_tile rows and both offsets are multiples of that, so a given global row lands in the same aligned tile and the kernel walks the same key blocks in the same order. Gradients are held to the rel ~2e-3 the shared cuDNN DSA backward already carries for dkv (atomics), with the off-vs-off pair checked against the same bound so the tolerance is visibly the kernel's and not this switch's.

是否引起精度变化

否

[risemeup1111]

序号	位置	优先级	状态
1	indexer loss 梯度回归		✅
2	CP 模式校验契约		✅

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

[xxyux]

/re-run all-failed

[xxyux]

/re-run all-failed

[codecov-commenter]

Codecov Report

❌ Patch coverage is 97.05882% with 2 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (develop@219e684). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
...rc/paddlefleet/transformer/mqa_latent_attention.py	93.10%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             develop    #1762   +/-   ##
==========================================
  Coverage           ?   97.05%           
==========================================
  Files              ?        3           
  Lines              ?       68           
  Branches           ?       13           
==========================================
  Hits               ?       66           
  Misses             ?        1           
  Partials           ?        1           

Flag	Coverage Δ
coverage_combine	97.05% <97.05%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/paddlefleet/transformer/cp_utils.py	100.00% <100.00%> (ø)
src/paddlefleet/transformer/transformer_config.py	100.00% <100.00%> (ø)
...rc/paddlefleet/transformer/mqa_latent_attention.py	93.10% <93.10%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ForFishes]

LGTM

[liuhao2638]

@risemeup1111 再次review

[sneaxiy]

LGTM for changes of transformer_config.py