[chaofengw-nv]

Background

Fixes #1793. The pinned Moshi gating kernel (gating_forward_kernel) is torch.compiled. It takes the Half gate and up projections, computes
SiLU(gate) * up in FP32, and rounds the product to Half once. The family SwiGLU instead computed the sigmoid, the
SiLU product and the up product in the work dtype, so SiLU(gate) was rounded to Half before the multiply.

Exit Criteria

• In FP16 and BF16 networks, the SwiGLU pointwise math runs in FP32 and is rounded once to the projection dtype. This
  applies to both family implementations: graph_blocks.add_swiglu_mlp and the dual-profile _swiglu_mlp.
• FP32 networks are unchanged.
• A predeclared GPU counterexample fails on the base and passes on this head through the MLP's normal output, using
  exact comparison.
• Non-goal: this PR does not claim that the full-model audio-parity gate passes.

Implementation

• New family-local helper graph_blocks.add_swiglu_activation(network, gate, up). It casts both projections to FP32,
  applies sigmoid, multiplies by gate and then by up, and casts back once.

• add_swiglu_mlp and default_dual_profile_decoder._swiglu_mlp call the helper. The projection GEMMs and the down
  projection are unchanged.

• New test families/personaplex/tests/test_swiglu_precision.py (markers gpu and trt; it skips without CUDA). The
  setup is:

  • a one-hot Half input;
  • single-path unit weights, so every GEMM output is exact;
  • gate = 7/64 and up = 1.5, exact in Half and BF16.

  The expected value is the reference formula evaluated in CUDA FP32, Half(SiLU(gate) * up) = 0.08648681640625; the old
  rounding gives 0.0865478515625. The FP32 product sits about 600 FP32 ULPs from a Half rounding boundary. The test
  covers both implementations, with 3 fresh engine builds each, and compares with rtol=0, atol=0.

Change categories

• Model or runtime behavior
• Public API
• ABI
• Bundle or artifact format
• Dependencies
• Documentation only
• CI or developer tooling

Validation

Commands and Results

• Original/fixed control on GB300. The setup:

  • isolated trees of main 9a8b545 and of this change, each with the identical new test (sha256 be3c147d...);
  • separate processes, with a check that each tree's families.personaplex is the module imported;
  • command: python -m pytest -q families/personaplex/tests/test_swiglu_precision.py.

  Results:

  • Original: 6 failed, 0 errors, 0 skipped. Each failure is 1/128 elements, with a maximum absolute difference of
    6.1035e-05 (one Half ULP).
  • Fixed: 6 passed, 0 errors, 0 skipped.
  • Product file hashes equal this commit (graph_blocks.py bd4c21a1..., default_dual_profile_decoder.py
39f3b001...).

• Actual reference data (diagnostic, reference FP16 run of audio-parity clip 7, temporal layers 0/15/31, steps
  0/27/28):

  • Half projections plus the FP32 pointwise product reproduce the reference gating output bitwise in 7 of 9 cases.
    The other two (layer 0 at steps 0 and 27) differ in 53 and 59 elements, cause not verified.
  • The family graph before this change differs in 1,748-2,292 of 4,096 elements.
  • A TensorRT network with the FP32 pointwise math differs in 580-799 of 4,096. The remainder is consistent with GEMM
    accumulation-order differences in the projections, but that is unverified.

• ruff check and ruff format --check on the new test and graph_blocks.py: clean. git diff --check: clean.

Hardware, Environment, and Revisions

GB300 (SBSA), CUDA 13.0, TensorRT 11.2.0.113, PyTorch from the product container. Tested head
8a011cf9b2bf3d7f9d53f8d8c4b20754f107a330. Reference data: checkpoint
nvidia/personaplex-7b-v1 revision fdaf4090a61cb315c138a1faee287ffd6c716309, run with the pinned PersonaPlex
reference environment in FP16.

Not Run / Remaining Gaps

• BF16 networks, tensor-parallel builds and the full model are not covered by this control. The BF16 and TP builders use
  the same add_swiglu_mlp, but no separate BF16 or TP evidence was collected.
• Full model: a diagnostic true-FP16 whole-clip run still fails audio-parity clip 7 against the reference FP16 baseline
  (5.58 dB, RMS ratio 0.851), with the first token split unchanged (step 28, codebook 1). That run combined an inline
  diagnostic SwiGLU patch, which is implementation-equivalent but not this head's files, with the pending fix(personaplex): preserve reference RoPE precision #1791 (rotary),
  fix(personaplex): round the FP16 depth input like the reference #1789 (depth input) and FP32 RMSNorm epsilon changes. It is not exact-head model validation. No accuracy gate or baseline
  was changed.
• default_dual_profile_decoder.py already fails ruff format --check on main (blank lines elsewhere in the file).
  This PR leaves that unrelated formatting unchanged.
• No CI results yet for this head.

Contributor Self-Review

• I have completed a self-review of this change.

Notes For Future Readers

• The helper keeps the activation and the product together in FP32. Do not round SiLU(gate) to the work dtype
  separately. The reference is a fused kernel, not the eager composition.
• Related PersonaPlex precision fixes are fix(personaplex): round the FP16 depth input like the reference #1789 (depth input) and fix(personaplex): preserve reference RoPE precision #1791 (rotary).

Risk level

• Low
• Medium
• High

Medium: this changes the FP16 and BF16 SwiGLU arithmetic in the family toward the reference semantics. FP32 networks
are unchanged.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[chaofengw-nv]

Superseded by #1751. This change and its regression coverage are preserved in commit 7395a62 of the consolidated PR. Original issue links and acceptance criteria remain tracked there.