Faster Helion JSD

[PaulZhang12]

Stacked PRs:

• ->Faster Helion JSD #733

---

Faster Helion JSD

Benchmarking command: HELION_USE_DEFAULT_CONFIG=0 HELION_PRINT_OUTPUT_CODE=1 HELION_AUTOTUNE_ACCURACY_CHECK=1 HELION_AUTOTUNE_RANDOM_SEED=0 CUDA_VISIBLE_DEVICES=7 python benchmarks/run.py --op jsd --num-inputs 3 --metrics speedup 2>&1 | tee ~/test_helion.log

       (B, T, V)    liger_jsd-speedup    torch_compile_jsd-speedup    helion_jsd_tritonbench-speedup
----------------  -------------------  ---------------------------  --------------------------------
 (4, 2048, 4096)             3.07473                       8.75428                           8.20297
 (4, 2048, 8192)             3.33668                       9.41338                           9.06205
(4, 2048, 16384)             0.587061                      9.38754                           8.82071
         average             2.33282                       9.18507                           8.69524

[yf225]

The Triton misaligned address error / unspecified launch failure error can be repro'ed with this config: @helion.kernel(config=helion.Config(block_sizes=[4, 256], indexing='tensor_descriptor', num_stages=4, num_warps=4, pid_type='flat', range_flattens=[None, False], range_multi_buffers=[None, False], range_num_stages=[0, 4], range_unroll_factors=[0, 0], range_warp_specializes=[]), static_shapes=True)

[yf225]

@PaulZhang12 #792 is merged and should unblock this PR. Thanks!

[Sibylau]

LGTM!