[aw471]

Summary

Adds optional post-norm to TransformerBlock, through attn_residual_norm after the attention residual and mlp_residual_norm after the feedforward residual, and applies dropout to the attention output in training. Kumo-Anomaly's ResidualBlock now builds on a KumoAnomalyTransformerBlock instead of torch.nn.TransformerEncoderLayer.

Description of the change

Both norms apply on the eager and the in-place out= paths. mlp_residual_norm applies only when mlp is set, and passing it without mlp raises. dropout reuses the attention-weight dropout from #1084 and, as in torch.nn.TransformerEncoderLayer, also drops the attention output before the residual addition. The compiled, non-contiguous out= path rounds the feedforward residual to the dtype of out before mlp_residual_norm, so eager and compiled outputs agree when norms run in float32 under autocast.

KumoAnomalyTransformerBlock only configures TransformerBlock: a Linear/GELU/Dropout/Linear/Dropout feedforward, LayerNorm residual norms and torch.nn.MultiheadAttention-style initialization. The test fixture that disabled PyTorch's MHA fast path (#1105) is removed, since the block tests now compare against a sequence-first torch.nn.TransformerEncoderLayer with copied weights.

This PR is stacked on #1211 but does not depend on it.

Validation

With the released Kumo-Anomaly weights, KumoAnomalyTransformerBlock matches a torch.nn.TransformerEncoderLayer loaded with the same weights to 4.4e-16 in float64 and 2.4e-7 in float32. The compiled out= paths are covered by test/nn/test_compile.py, which runs only with FULL_TEST=1; those tests pass on CPU.

[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.

[aw471]

/ok to test aeb403f

[codecov]

Codecov Report

❌ Patch coverage is 83.33333% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.91%. Comparing base (a3ee9a3) to head (bbf07e1).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
sdm/nn/attention.py	72.22%	5 Missing ⚠️

Additional details and impacted files

@@                    Coverage Diff                    @@
##           aw471/kumo-anomaly-model    #1106   +/-   ##
=========================================================
  Coverage                     82.91%   82.91%           
=========================================================
  Files                           172      172           
  Lines                         11107    11133   +26     
=========================================================
+ Hits                           9209     9231   +22     
- Misses                         1898     1902    +4     

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