Implement softcapping in fused jsd

[wheynelau]

trafficstars

Summary

Implements softcap in the fused linear jsd, so it can be used for gemma2 models

Details

Assumes same softcap for teacher and student model

Testing Done

• added tests for softcapping in test_fused_linear_jsd.py

• Hardware Type: L40S

• [x] run make test to ensure correctness

• [x] run make checkstyle to ensure code style

• [ ] run make test-convergence to ensure convergence

[yundai424]

@wheynelau i just fixed some conflict due to out of sync 😃 FYI test_correctness_functional is failing for me on A100 when softcap=50.0. Raw output here: https://gist.github.com/yundai424/7bbfa78f05667749ce189cd458cabf90

[wheynelau]

@yundai424 Okay thanks! Let me take a look at this

[wheynelau]

@yundai424 Have updated it!