Fix memory requirements for benchmark_jsd.py and benchmark_distill_jsd_loss.py (#1050)

## Summary
In #1037 I mixed
`benchmark_jsd.py` and `benchmark_distill_jsd_loss.py`.

It is `benchmark_distill_jsd_loss.py` that requires 
39GBs of memory:
https://github.com/linkedin/Liger-Kernel/blob/0ea0b8ffcee27c5c94ffa87e480ea95036a0d2da/benchmark/data/all_benchmark_data.csv#L746

69GBs of memory:
https://github.com/linkedin/Liger-Kernel/blob/0ea0b8ffcee27c5c94ffa87e480ea95036a0d2da/benchmark/data/all_benchmark_data.csv#L747

`benchmark_jsd.py` requires just 53GBs:

https://github.com/linkedin/Liger-Kernel/blob/0ea0b8ffcee27c5c94ffa87e480ea95036a0d2da/benchmark/data/all_benchmark_data.csv#L459

## Testing Done
I run both benchmarks on XPU Intel GPU Max 1100 with 48 GBs of memory

- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

Author: Egor-Krivov

benchmark/scripts/benchmark_distill_jsd_loss.py

@@ -12,6 +12,7 @@
 from utils import run_benchmarks
 
 from liger_kernel.chunked_loss.jsd_loss import LigerFusedLinearJSDFunction
+from liger_kernel.utils import get_total_gpu_memory
 from liger_kernel.utils import infer_device
 
 device = infer_device()
@@ -224,12 +225,20 @@ def full():
 
 if __name__ == "__main__":
     args = parse_benchmark_script_args()
+    gpu_memory_gbs = get_total_gpu_memory()
+    # We know that the full test will require 69GBs for vocab size 2^13 and 39GBs for vocab size 2^12 on torch
+    if gpu_memory_gbs >= 69:
+        x_max = 13
+    elif gpu_memory_gbs >= 39:
+        x_max = 12
+    else:
+        x_max = 11
 
     common_configs = {
         "kernel_name": "distill_jsd_loss",
         "x_name": "BT",
         "x_label": "B x T",
-        "x_values": [2**i for i in range(10, 14)],
+        "x_values": [2**i for i in range(10, x_max + 1)],
         "kernel_providers": ["liger", "torch"],
         "extra_benchmark_configs": [
             {

benchmark/scripts/benchmark_jsd.py

@@ -125,13 +125,11 @@ def full():
 if __name__ == "__main__":
     args = parse_benchmark_script_args()
     gpu_memory_gbs = get_total_gpu_memory()
-    # We know that the full test will require 69GBs for vocab size 2^17 and 39GBs for vocab size 2^16 on torch
-    if gpu_memory_gbs >= 69:
+    # We know that the full test will require 54GBs for vocab size 2^17 on torch
+    if gpu_memory_gbs >= 54:
         x_max = 17
-    elif gpu_memory_gbs >= 39:
-        x_max = 16
     else:
-        x_max = 15
+        x_max = 16
     common_args = {
         "kernel_name": "jsd",
         "x_name": "V",