[Feature][OP] Add batch-invariant RMSNorm kernel and TP embedding Custom AR path

[gongweibao]

Motivation

In deterministic inference mode, batch composition changes (e.g., prefix cache hit vs miss) can cause different padding/concatenation patterns, leading to non-deterministic outputs from RMSNorm and TP all-reduce operations. This PR makes these operators M-invariant (batch-size independent) to ensure identical outputs regardless of batch composition.

Modifications

• Triton RMSNorm kernel (batch_invariant_ops.py): Add rms_norm_batch_invariant — a per-row Triton kernel that computes RMSNorm independently per row, eliminating cross-row reduction non-determinism.
• RMSNorm routing (normalization.py): When batch_invariant_mode is enabled, route through the Triton kernel instead of the default fused kernel.
• TP Embedding Custom AR (embeddings.py): In deterministic mode with TP>1, bypass Paddle's _mp_allreduce (NCCL) and use Custom All-Reduce for deterministic embedding all-reduce.
• Custom AR buffer size (envs.py): Increase FD_CUSTOM_AR_MAX_SIZE_MB default from 8 to 64 MB to accommodate larger tensors without NCCL fallback.
• CUDA kernel fix (multiquery_attention_c16_impl.cuh): Minor type adjustments in the append attention kernel.
• Unit tests: Add RMSNorm batch-invariance test and TP embedding deterministic test.

Usage or Command

Enable deterministic mode:

export FD_DETERMINISTIC_MODE=1

The batch-invariant RMSNorm is automatically activated when batch_invariant_mode is enabled during deterministic forward passes.

Accuracy Tests

• tests/batch_invariant/test_batch_invariance_op_rmsnorm.py: Verifies RMSNorm output is identical across different batch compositions (padding patterns).
• tests/e2e/4cards_cases/vocab_parallel_embedding_deterministic.py: Verifies TP embedding produces deterministic results with Custom AR path.

Checklist

• Add at least a tag in the PR title.
• Format your code, run pre-commit before commit.
• Add unit tests.
• Provide accuracy results.
• If the current PR is submitting to the release branch, make sure the PR has been submitted to the develop branch, then cherry-pick it to the release branch with the [Cherry-Pick] PR tag.

[paddle-bot]

Thanks for your contribution!

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

gongweibao seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[codecov-commenter]

Codecov Report

❌ Patch coverage is 36.58537% with 26 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (develop@28f7727). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
.../layers/batch_invariant_ops/batch_invariant_ops.py	39.28%	17 Missing ⚠️
fastdeploy/model_executor/layers/embeddings.py	28.57%	4 Missing and 1 partial ⚠️
fastdeploy/model_executor/layers/normalization.py	33.33%	3 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             develop    #6749   +/-   ##
==========================================
  Coverage           ?   72.03%           
==========================================
  Files              ?      392           
  Lines              ?    53919           
  Branches           ?     8475           
==========================================
  Hits               ?    38839           
  Misses             ?    12302           
  Partials           ?     2778           

Flag	Coverage Δ
GPU	72.03% <36.58%> (?)

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

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

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

[gongshaotian]

deterministic 的端到端单测尽量合并成一个？减少ci压力

[gongweibao]

这个比较小。CI 增加的有限。时间很短。
主要是跟其他的单测没有关系。