[PyTorch][CI] Fix test issues: connect test files that are not wired into qa

[pggPL]

Description

Part of the pytorch test suite is not connected to qa/ at all. Six test files are not referenced
by any script there, so they have never run in CI. This PR connects them and fixes the one that
turned out to be broken.

Confirmed with git log -S<name> --all -- qa/: none of these files was ever added and later
removed, and none of the PRs that introduced them touched qa/. They were simply never wired up.

file	status on current main
test_qk_norm.py	45 passed
test_float8_current_scaling_exact.py	5 passed
attention/test_cu_seqlens_cache.py	1 passed, 1 skipped (needs 2 GPUs)
test_nvfp4_fsdp2_hooks.py	16 skipped (needs sm_100+)
test_fused_router_perf.py	5 skipped (gated behind TE_RUN_PERF_TESTS)
layernorm_mlp/test_selective_activation_checkpoint.py	16 failed

All of them are single-GPU and self-skip on unsupported hardware, so they belong in L0.

test_fused_router_perf.py is deliberately left out of this PR — it is gated behind
TE_RUN_PERF_TESTS and, despite its name, contains no perf assertions at all (only
torch.testing.assert_close on correctness, with timings going to record_property). Whether to
ungate the correctness half deserves a separate decision.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

1. Connect four files to L0 (commit 2)

• test_qk_norm.py, test_float8_current_scaling_exact.py and attention/test_cu_seqlens_cache.py
  get an entry in qa/L0_pytorch_unittest/test.sh.
• test_nvfp4_fsdp2_hooks.py moves into tests/pytorch/nvfp4/, which L0 already runs as a whole
  directory. No qa/ change needed, and the file is now covered by the same rule as the other
  NVFP4 tests. It has no local imports, so the move is inert. The real multi-GPU FSDP2 path for
  NVFP4 is already covered separately by distributed/fsdp2_tests/run_fsdp2_model.py --recipe NVFP4BlockScaling in L1; this file is a single-GPU unit test of the shape and metadata
  arithmetic in fsdp_pre_all_gather / fsdp_post_all_gather.

Hardware guards are already in place in all four: test_float8_current_scaling_exact.py marks its
classes with skipif(not fp8_available), test_nvfp4_fsdp2_hooks.py requires sm_100+, and the one
multi-device case in test_cu_seqlens_cache.py checks device_count() first.

Added L0 cost is roughly 13 s.

2. Fix test_selective_activation_checkpoint.py (commit 1)

This is the file that never ran and does not pass. The failure is in the test, not in the feature.
Selective activation checkpointing is correct and does save memory:

• max|out_ckpt - out_no_ckpt| = 0.0 and max|grad_ckpt - grad_no_ckpt| = 0.0 for all six
  parameters, in every configuration that fits in memory.
• The forward memory ratio is a structural constant: 5.719 for small, 5.715 for medium /
  large / huge — independent of seq_len (128 → 65536) and hidden_size (128 → 2048), with no
  run-to-run variance.

The test asserted ln_fwd_mem > 6 * sln_fwd_mem. That threshold was never reachable, and it never
could have been, because the ratio is fixed by the test's own configuration rather than by anything
in TE. Both sides of the measurement are fully derivable, and the derivations reproduce the
measured peaks exactly:

ln_fwd_mem  = layers * (2*s*f + 2*s*h + 2*s) * itemsize
              fc1_out, act_out, ln_out, out, mu, rsigma per layer

sln_fwd_mem = ((layers+1)*s*h + 2*s*f + 2*s) * itemsize
              layer inputs shared along the Sequential chain, plus the transient of a
              single layer, since the checkpointed forward still computes fc1_out and
              act_out before freeing them

For small @ 128 this predicts 7876608 B and 1377280 B; measured 7876608 B and 1377280 B.
Substituting f = 4h and layers = 12 gives 120 / 21 = 5.714 in units of s*h, matching the
measurement. Reaching 6 would require a different model shape, e.g. layers = 16 gives
160 / 25 = 6.4.

I also checked that this is not a regression: the tensor lists in both branches of
_LayerNormMLP._forward are unchanged since the test was added (the only difference is the
fc1_wt_save / fc2_wt_save gating for FSDP2 from NVIDIA#2681, inactive for plain parameters), and the
test file itself has only been touched by the 2026 copyright update. Weights are allocated before
the start_mem snapshot, so they never entered this measurement in the first place.

Changes:

• Assert on the memory that recompute actually frees — fc1_out and act_out per layer, derived
  from the model config — instead of the ratio. Observed saving is 1.031-1.033x the derived value,
  so the 0.95 factor leaves ~8% headroom while still failing loudly if either recomputed tensor
  stops being freed (that would halve the saving). Unlike a whole-forward ratio, this does not
  depend on layers or on the ffn_hidden / hidden ratio.
• Move the output and gradient checks before the memory check. Previously a numerical
  regression would have been reported as a memory-ratio failure and the correctness comparison
  would never have run — which is exactly what happens on main today.
• Drop the bare assert ln_bwd_time < sln_bwd_time. The margin is as low as 13% on an idle GPU
  (huge @ 128: 11.2 ms vs 12.8 ms), which makes it a flake on a loaded CI runner. Forward and
  backward timings plus forward/backward memory ratios are now reported through record_property,
  so they land in the JUnit XML as trend data instead of a pass/fail gate.
• Skip parametrizations that do not fit in device memory. large @ 65536 and huge @ 65536 need
  more than 32 GiB for the non-checkpointed model alone and OOM on 48 GiB cards; the required
  budget is computed up front from the same analytic formula.

Test results

RTX 5880 Ada (48 GiB), current main + this PR:

test_selective_activation_checkpoint.py   12 passed, 4 skipped   (was 16 failed)
test_qk_norm.py                           45 passed
test_float8_current_scaling_exact.py       5 passed
attention/test_cu_seqlens_cache.py         1 passed, 1 skipped
nvfp4/test_nvfp4_fsdp2_hooks.py           16 skipped (sm_100+ required)

Skips in the SAC test are the memory guard doing its job:

SKIPPED [2] needs 30.0 GiB free device memory, only 29.5 GiB available
SKIPPED [1] needs 60.0 GiB free device memory, only 29.5 GiB available
SKIPPED [1] needs 120.0 GiB free device memory, only 29.5 GiB available

Sample of the recorded properties (JUnit XML):

[16384-large]  ln_fwd_ms=56.6  sln_fwd_ms=55.9  ln_bwd_ms=110.5  sln_bwd_ms=147.1  fwd_mem_ratio=5.715
[65536-small]  ln_fwd_ms=9.4   sln_fwd_ms=9.4   ln_bwd_ms=19.8   sln_bwd_ms=26.7   fwd_mem_ratio=5.719

Follow-up

Nothing in the repo prevents this from happening again — a new test file simply has to be
remembered in a shell script. A guard test asserting that every tests/**/test_*.py is covered by
at least one job in qa/, plus tests/ entries in CODEOWNERS (which currently only covers
transformer_engine/), would close the class of problem. Happy to send that separately if there is
appetite for it.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes