diff --git a/benchmark/data/all_benchmark_data.csv b/benchmark/data/all_benchmark_data.csv index 84e73583c..15ef7b8f9 100644 --- a/benchmark/data/all_benchmark_data.csv +++ b/benchmark/data/all_benchmark_data.csv @@ -2019,114 +2019,114 @@ fused_add_rms_norm,liger_rms_norm,full,memory,MB,H,hidden size,4096,416.11767578 fused_add_rms_norm,liger_rms_norm,full,memory,MB,H,hidden size,8192,832.22705078125,832.22705078125,832.22705078125,"{""M"": 2048, ""dtype"": ""torch.float32"", ""eps"": 1e-06}",NVIDIA H100 80GB HBM3,2026-04-07 21:27:45,0.7.0 fused_add_rms_norm,liger_rms_norm,full,memory,MB,H,hidden size,16384,1544.44580078125,1544.44580078125,1544.44580078125,"{""M"": 2048, ""dtype"": ""torch.float32"", ""eps"": 1e-06}",NVIDIA H100 80GB HBM3,2026-04-07 21:27:45,0.7.0 fused_add_rms_norm,liger_rms_norm,full,memory,MB,H,hidden size,32768,2960.8837890625,2960.8837890625,2960.8837890625,"{""M"": 2048, ""dtype"": ""torch.float32"", ""eps"": 1e-06}",NVIDIA H100 80GB HBM3,2026-04-07 21:27:45,0.7.0 -fused_moe,liger,full,speed,ms,T,num_tokens,128,3.86625599861145,3.8625791549682615,3.869932842254639,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:32,0.7.0 -fused_moe,liger,full,speed,ms,T,num_tokens,256,3.926751971244812,3.925888013839722,3.927615928649902,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:32,0.7.0 -fused_moe,liger,full,speed,ms,T,num_tokens,512,4.070320129394531,4.0700225830078125,4.07061767578125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:32,0.7.0 -fused_moe,liger,full,speed,ms,T,num_tokens,1024,4.403903961181641,4.402329540252686,4.405478382110596,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:32,0.7.0 -fused_moe,liger,full,speed,ms,T,num_tokens,2048,5.0946879386901855,5.0946879386901855,5.0946879386901855,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:32,0.7.0 -fused_moe,liger,full,speed,ms,T,num_tokens,4096,6.476992130279541,6.476992130279541,6.476992130279541,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:32,0.7.0 -fused_moe,liger,full,speed,ms,T,num_tokens,8192,9.547840118408203,9.547840118408203,9.547840118408203,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:32,0.7.0 -fused_moe,liger,full,speed,ms,T,num_tokens,16384,15.691679954528809,15.691679954528809,15.691679954528809,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:32,0.7.0 -fused_moe,liger,full,speed,ms,T,num_tokens,32768,28.224639892578125,28.224639892578125,28.224639892578125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:32,0.7.0 -fused_moe,huggingface,full,speed,ms,T,num_tokens,128,244.2610626220703,244.2610626220703,244.2610626220703,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:50,0.7.0 -fused_moe,huggingface,full,speed,ms,T,num_tokens,256,244.83847045898438,244.83847045898438,244.83847045898438,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:50,0.7.0 -fused_moe,huggingface,full,speed,ms,T,num_tokens,512,244.74520874023438,244.74520874023438,244.74520874023438,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:50,0.7.0 -fused_moe,huggingface,full,speed,ms,T,num_tokens,1024,246.53123474121094,246.53123474121094,246.53123474121094,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:50,0.7.0 -fused_moe,huggingface,full,speed,ms,T,num_tokens,2048,248.0272674560547,248.0272674560547,248.0272674560547,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:50,0.7.0 -fused_moe,huggingface,full,speed,ms,T,num_tokens,4096,250.1736297607422,250.1736297607422,250.1736297607422,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:50,0.7.0 -fused_moe,huggingface,full,speed,ms,T,num_tokens,8192,255.07052612304688,255.07052612304688,255.07052612304688,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:50,0.7.0 -fused_moe,huggingface,full,speed,ms,T,num_tokens,16384,266.1387634277344,266.1387634277344,266.1387634277344,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:50,0.7.0 -fused_moe,huggingface,full,speed,ms,T,num_tokens,32768,287.9275207519531,287.9275207519531,287.9275207519531,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:50,0.7.0 -fused_moe,liger,forward,speed,ms,T,num_tokens,128,0.5888479948043823,0.5869632005691529,0.590726375579834,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:51,0.7.0 -fused_moe,liger,forward,speed,ms,T,num_tokens,256,0.5958719849586487,0.5934015989303589,0.5983423948287964,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:51,0.7.0 -fused_moe,liger,forward,speed,ms,T,num_tokens,512,0.6133120059967041,0.6098751902580262,0.6163007855415344,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:51,0.7.0 -fused_moe,liger,forward,speed,ms,T,num_tokens,1024,0.6712960004806519,0.6674431920051576,0.6750207901000976,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:51,0.7.0 -fused_moe,liger,forward,speed,ms,T,num_tokens,2048,0.8744959831237793,0.8713215827941895,0.8767807960510254,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:51,0.7.0 -fused_moe,liger,forward,speed,ms,T,num_tokens,4096,1.2509759664535522,1.249452829360962,1.2559295892715454,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:51,0.7.0 -fused_moe,liger,forward,speed,ms,T,num_tokens,8192,2.0957120656967163,2.0833152770996093,2.0995328426361084,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:51,0.7.0 -fused_moe,liger,forward,speed,ms,T,num_tokens,16384,3.8910878896713257,3.8891006946563724,3.8930750846862794,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:51,0.7.0 -fused_moe,liger,forward,speed,ms,T,num_tokens,32768,7.342400074005127,7.342400074005127,7.342400074005127,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:51,0.7.0 -fused_moe,huggingface,forward,speed,ms,T,num_tokens,128,30.31283187866211,30.31283187866211,30.31283187866211,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:53,0.7.0 -fused_moe,huggingface,forward,speed,ms,T,num_tokens,256,30.287456512451172,30.287456512451172,30.287456512451172,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:53,0.7.0 -fused_moe,huggingface,forward,speed,ms,T,num_tokens,512,30.208576202392578,30.208576202392578,30.208576202392578,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:53,0.7.0 -fused_moe,huggingface,forward,speed,ms,T,num_tokens,1024,31.234560012817383,31.234560012817383,31.234560012817383,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:53,0.7.0 -fused_moe,huggingface,forward,speed,ms,T,num_tokens,2048,31.455904006958008,31.455904006958008,31.455904006958008,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:53,0.7.0 -fused_moe,huggingface,forward,speed,ms,T,num_tokens,4096,31.446624755859375,31.446624755859375,31.446624755859375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:53,0.7.0 -fused_moe,huggingface,forward,speed,ms,T,num_tokens,8192,31.467967987060547,31.467967987060547,31.467967987060547,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:53,0.7.0 -fused_moe,huggingface,forward,speed,ms,T,num_tokens,16384,33.299617767333984,33.299617767333984,33.299617767333984,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:53,0.7.0 -fused_moe,huggingface,forward,speed,ms,T,num_tokens,32768,37.644065856933594,37.644065856933594,37.644065856933594,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:53,0.7.0 -fused_moe,liger,backward,speed,ms,T,num_tokens,128,3.256111979484558,3.2556800365448,3.2565439224243162,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:54,0.7.0 -fused_moe,liger,backward,speed,ms,T,num_tokens,256,3.3059359788894653,3.3057343482971193,3.3061376094818113,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:54,0.7.0 -fused_moe,liger,backward,speed,ms,T,num_tokens,512,3.431328058242798,3.4307904720306395,3.4318656444549562,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:54,0.7.0 -fused_moe,liger,backward,speed,ms,T,num_tokens,1024,3.710319995880127,3.7087359905242923,3.711904001235962,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:54,0.7.0 -fused_moe,liger,backward,speed,ms,T,num_tokens,2048,4.176991939544678,4.175801467895508,4.1781824111938475,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:54,0.7.0 -fused_moe,liger,backward,speed,ms,T,num_tokens,4096,5.179903984069824,5.179903984069824,5.179903984069824,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:54,0.7.0 -fused_moe,liger,backward,speed,ms,T,num_tokens,8192,7.382239818572998,7.382239818572998,7.382239818572998,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:54,0.7.0 -fused_moe,liger,backward,speed,ms,T,num_tokens,16384,11.77836799621582,11.77836799621582,11.77836799621582,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:54,0.7.0 -fused_moe,liger,backward,speed,ms,T,num_tokens,32768,20.486303329467773,20.486303329467773,20.486303329467773,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:57:54,0.7.0 -fused_moe,huggingface,backward,speed,ms,T,num_tokens,128,214.22210693359375,214.22210693359375,214.22210693359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:10,0.7.0 -fused_moe,huggingface,backward,speed,ms,T,num_tokens,256,214.73184204101562,214.73184204101562,214.73184204101562,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:10,0.7.0 -fused_moe,huggingface,backward,speed,ms,T,num_tokens,512,214.63980102539062,214.63980102539062,214.63980102539062,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:10,0.7.0 -fused_moe,huggingface,backward,speed,ms,T,num_tokens,1024,215.37744140625,215.37744140625,215.37744140625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:10,0.7.0 -fused_moe,huggingface,backward,speed,ms,T,num_tokens,2048,216.5009002685547,216.5009002685547,216.5009002685547,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:10,0.7.0 -fused_moe,huggingface,backward,speed,ms,T,num_tokens,4096,218.47984313964844,218.47984313964844,218.47984313964844,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:10,0.7.0 -fused_moe,huggingface,backward,speed,ms,T,num_tokens,8192,223.47523498535156,223.47523498535156,223.47523498535156,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:10,0.7.0 -fused_moe,huggingface,backward,speed,ms,T,num_tokens,16384,232.13894653320312,232.13894653320312,232.13894653320312,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:10,0.7.0 -fused_moe,huggingface,backward,speed,ms,T,num_tokens,32768,249.76806640625,249.76806640625,249.76806640625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:10,0.7.0 -fused_moe,liger,full,memory,MB,T,num_tokens,128,5061.021484375,5061.021484375,5061.021484375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:11,0.7.0 -fused_moe,liger,full,memory,MB,T,num_tokens,256,5066.041015625,5066.041015625,5066.041015625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:11,0.7.0 -fused_moe,liger,full,memory,MB,T,num_tokens,512,5076.080078125,5076.080078125,5076.080078125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:11,0.7.0 -fused_moe,liger,full,memory,MB,T,num_tokens,1024,5096.1591796875,5096.1591796875,5096.1591796875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:11,0.7.0 -fused_moe,liger,full,memory,MB,T,num_tokens,2048,5136.31640625,5136.31640625,5136.31640625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:11,0.7.0 -fused_moe,liger,full,memory,MB,T,num_tokens,4096,5216.630859375,5216.630859375,5216.630859375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:11,0.7.0 -fused_moe,liger,full,memory,MB,T,num_tokens,8192,5569.384765625,5569.384765625,5569.384765625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:11,0.7.0 -fused_moe,liger,full,memory,MB,T,num_tokens,16384,6466.767578125,6466.767578125,6466.767578125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:11,0.7.0 -fused_moe,liger,full,memory,MB,T,num_tokens,32768,8261.533203125,8261.533203125,8261.533203125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:11,0.7.0 -fused_moe,huggingface,full,memory,MB,T,num_tokens,128,5466.693359375,5466.693359375,5466.693359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,huggingface,full,memory,MB,T,num_tokens,256,5487.25390625,5487.25390625,5487.25390625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,huggingface,full,memory,MB,T,num_tokens,512,5528.3955078125,5528.3955078125,5528.3955078125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,huggingface,full,memory,MB,T,num_tokens,1024,5610.61962890625,5610.61962890625,5610.61962890625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,huggingface,full,memory,MB,T,num_tokens,2048,5775.12841796875,5775.12841796875,5775.12841796875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,huggingface,full,memory,MB,T,num_tokens,4096,6105.1943359375,6105.1943359375,6105.1943359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,huggingface,full,memory,MB,T,num_tokens,8192,6763.20654296875,6763.20654296875,6763.20654296875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,huggingface,full,memory,MB,T,num_tokens,16384,8082.79736328125,8082.79736328125,8082.79736328125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,huggingface,full,memory,MB,T,num_tokens,32768,10711.859375,10711.859375,10711.859375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,liger,forward,memory,MB,T,num_tokens,128,2377.52001953125,2377.52001953125,2377.52001953125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,liger,forward,memory,MB,T,num_tokens,256,2387.03759765625,2387.03759765625,2387.03759765625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,liger,forward,memory,MB,T,num_tokens,512,2406.07275390625,2406.07275390625,2406.07275390625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,liger,forward,memory,MB,T,num_tokens,1024,2444.14404296875,2444.14404296875,2444.14404296875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,liger,forward,memory,MB,T,num_tokens,2048,2520.28564453125,2520.28564453125,2520.28564453125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,liger,forward,memory,MB,T,num_tokens,4096,2672.56884765625,2672.56884765625,2672.56884765625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,liger,forward,memory,MB,T,num_tokens,8192,2977.13525390625,2977.13525390625,2977.13525390625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,liger,forward,memory,MB,T,num_tokens,16384,3586.26806640625,3586.26806640625,3586.26806640625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,liger,forward,memory,MB,T,num_tokens,32768,4804.53369140625,4804.53369140625,4804.53369140625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:34,0.7.0 -fused_moe,huggingface,forward,memory,MB,T,num_tokens,128,2388.1318359375,2388.1318359375,2388.1318359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:37,0.7.0 -fused_moe,huggingface,forward,memory,MB,T,num_tokens,256,2408.1376953125,2408.1376953125,2408.1376953125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:37,0.7.0 -fused_moe,huggingface,forward,memory,MB,T,num_tokens,512,2448.177734375,2448.177734375,2448.177734375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:37,0.7.0 -fused_moe,huggingface,forward,memory,MB,T,num_tokens,1024,2528.26513671875,2528.26513671875,2528.26513671875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:37,0.7.0 -fused_moe,huggingface,forward,memory,MB,T,num_tokens,2048,2688.43798828125,2688.43798828125,2688.43798828125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:37,0.7.0 -fused_moe,huggingface,forward,memory,MB,T,num_tokens,4096,3009.0703125,3009.0703125,3009.0703125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:37,0.7.0 -fused_moe,huggingface,forward,memory,MB,T,num_tokens,8192,3650.37158203125,3650.37158203125,3650.37158203125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:37,0.7.0 -fused_moe,huggingface,forward,memory,MB,T,num_tokens,16384,4934.99169921875,4934.99169921875,4934.99169921875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:37,0.7.0 -fused_moe,huggingface,forward,memory,MB,T,num_tokens,32768,7497.27587890625,7497.27587890625,7497.27587890625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:37,0.7.0 -fused_moe,liger,backward,memory,MB,T,num_tokens,128,5061.021484375,5061.021484375,5061.021484375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:38,0.7.0 -fused_moe,liger,backward,memory,MB,T,num_tokens,256,5066.041015625,5066.041015625,5066.041015625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:38,0.7.0 -fused_moe,liger,backward,memory,MB,T,num_tokens,512,5076.080078125,5076.080078125,5076.080078125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:38,0.7.0 -fused_moe,liger,backward,memory,MB,T,num_tokens,1024,5096.1591796875,5096.1591796875,5096.1591796875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:38,0.7.0 -fused_moe,liger,backward,memory,MB,T,num_tokens,2048,5136.31640625,5136.31640625,5136.31640625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:38,0.7.0 -fused_moe,liger,backward,memory,MB,T,num_tokens,4096,5216.630859375,5216.630859375,5216.630859375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:38,0.7.0 -fused_moe,liger,backward,memory,MB,T,num_tokens,8192,5569.384765625,5569.384765625,5569.384765625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:38,0.7.0 -fused_moe,liger,backward,memory,MB,T,num_tokens,16384,6466.767578125,6466.767578125,6466.767578125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:38,0.7.0 -fused_moe,liger,backward,memory,MB,T,num_tokens,32768,8261.533203125,8261.533203125,8261.533203125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:38,0.7.0 -fused_moe,huggingface,backward,memory,MB,T,num_tokens,128,5466.693359375,5466.693359375,5466.693359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:58,0.7.0 -fused_moe,huggingface,backward,memory,MB,T,num_tokens,256,5487.25390625,5487.25390625,5487.25390625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:58,0.7.0 -fused_moe,huggingface,backward,memory,MB,T,num_tokens,512,5528.3955078125,5528.3955078125,5528.3955078125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:58,0.7.0 -fused_moe,huggingface,backward,memory,MB,T,num_tokens,1024,5610.61962890625,5610.61962890625,5610.61962890625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:58,0.7.0 -fused_moe,huggingface,backward,memory,MB,T,num_tokens,2048,5775.12841796875,5775.12841796875,5775.12841796875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:58,0.7.0 -fused_moe,huggingface,backward,memory,MB,T,num_tokens,4096,6104.3583984375,6104.3583984375,6104.3583984375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:58,0.7.0 -fused_moe,huggingface,backward,memory,MB,T,num_tokens,8192,6762.88623046875,6762.88623046875,6762.88623046875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:58,0.7.0 -fused_moe,huggingface,backward,memory,MB,T,num_tokens,16384,8081.99072265625,8081.99072265625,8081.99072265625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:58,0.7.0 -fused_moe,huggingface,backward,memory,MB,T,num_tokens,32768,10712.716796875,10712.716796875,10712.716796875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:58:58,0.7.0 +fused_moe,liger,full,speed,ms,T,num_tokens,128,3.252880096435547,3.252524900436401,3.2532352924346926,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:20,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,256,3.322848081588745,3.320467281341553,3.3252288818359377,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:20,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,512,3.4244319200515747,3.4236927032470703,3.425171136856079,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:20,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,1024,3.63755202293396,3.637292814254761,3.637811231613159,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:20,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,2048,4.166431903839111,4.164281558990479,4.168582248687744,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:20,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,4096,5.229824066162109,5.229824066162109,5.229824066162109,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:20,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,8192,7.380095958709717,7.380095958709717,7.380095958709717,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:20,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,16384,11.771743774414062,11.771743774414062,11.771743774414062,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:20,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,32768,21.382400512695312,21.382400512695312,21.382400512695312,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:20,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,128,243.1138916015625,243.1138916015625,243.1138916015625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,256,245.09808349609375,245.09808349609375,245.09808349609375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,512,245.65708923339844,245.65708923339844,245.65708923339844,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,1024,245.98092651367188,245.98092651367188,245.98092651367188,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,2048,246.71241760253906,246.71241760253906,246.71241760253906,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,4096,248.5768585205078,248.5768585205078,248.5768585205078,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,8192,255.05027770996094,255.05027770996094,255.05027770996094,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,16384,264.7030944824219,264.7030944824219,264.7030944824219,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,32768,286.83026123046875,286.83026123046875,286.83026123046875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,128,0.4447999894618988,0.44380800724029545,0.4462591886520386,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,256,0.46931201219558716,0.4688639879226685,0.4703359961509705,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,512,0.47775998711586,0.47673600912094116,0.47891199588775635,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,1024,0.5133920013904572,0.5127040147781372,0.5141119956970215,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,2048,0.6037120223045349,0.6027136087417603,0.6050432085990906,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,4096,0.8442560136318207,0.8422399878501893,0.847705602645874,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,8192,1.5113440155982971,1.5077439546585083,1.512992024421692,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,16384,2.6998400688171387,2.698150491714477,2.700147199630737,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,32768,5.086463928222656,5.086463928222656,5.086463928222656,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:39,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,128,28.65065574645996,28.65065574645996,28.65065574645996,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:41,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,256,28.873247146606445,28.873247146606445,28.873247146606445,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:41,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,512,28.469600677490234,28.469600677490234,28.469600677490234,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:41,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,1024,29.30975914001465,29.30975914001465,29.30975914001465,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:41,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,2048,29.72710418701172,29.72710418701172,29.72710418701172,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:41,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,4096,29.906591415405273,29.906591415405273,29.906591415405273,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:41,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,8192,30.028287887573242,30.028287887573242,30.028287887573242,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:41,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,16384,31.180736541748047,31.180736541748047,31.180736541748047,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:41,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,32768,35.575199127197266,35.575199127197266,35.575199127197266,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:41,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,128,2.812191963195801,2.8102912425994875,2.81267204284668,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:42,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,256,2.8557119369506836,2.8534271240234377,2.855942392349243,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:42,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,512,2.956160068511963,2.9550848960876466,2.9563135623931887,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:42,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,1024,3.1283841133117676,3.125542402267456,3.1284992694854736,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:42,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,2048,3.538655996322632,3.5377536296844485,3.5395583629608156,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:42,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,4096,4.334159851074219,4.333516693115234,4.334803009033203,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:42,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,8192,5.943359851837158,5.943359851837158,5.943359851837158,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:42,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,16384,9.129568099975586,9.129568099975586,9.129568099975586,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:42,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,32768,16.25609588623047,16.25609588623047,16.25609588623047,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:42,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,128,214.5557098388672,214.5557098388672,214.5557098388672,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:58,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,256,214.8544921875,214.8544921875,214.8544921875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:58,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,512,214.719482421875,214.719482421875,214.719482421875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:58,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,1024,215.3389129638672,215.3389129638672,215.3389129638672,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:58,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,2048,216.3684539794922,216.3684539794922,216.3684539794922,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:58,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,4096,218.29324340820312,218.29324340820312,218.29324340820312,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:58,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,8192,223.3937530517578,223.3937530517578,223.3937530517578,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:58,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,16384,231.8494110107422,231.8494110107422,231.8494110107422,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:58,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,32768,249.27987670898438,249.27987670898438,249.27987670898438,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:58,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,128,5061.0234375,5061.0234375,5061.0234375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:59,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,256,5066.04296875,5066.04296875,5066.04296875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:59,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,512,5076.08203125,5076.08203125,5076.08203125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:59,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,1024,5096.16015625,5096.16015625,5096.16015625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:59,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,2048,5136.31640625,5136.31640625,5136.31640625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:59,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,4096,5216.6298828125,5216.6298828125,5216.6298828125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:59,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,8192,5569.6318359375,5569.6318359375,5569.6318359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:59,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,16384,6467.2607421875,6467.2607421875,6467.2607421875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:59,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,32768,8262.5185546875,8262.5185546875,8262.5185546875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:15:59,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,128,5466.693359375,5466.693359375,5466.693359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,256,5487.25390625,5487.25390625,5487.25390625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,512,5528.3955078125,5528.3955078125,5528.3955078125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,1024,5610.61962890625,5610.61962890625,5610.61962890625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,2048,5775.12841796875,5775.12841796875,5775.12841796875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,4096,6105.1943359375,6105.1943359375,6105.1943359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,8192,6762.67529296875,6762.67529296875,6762.67529296875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,16384,8081.0048828125,8081.0048828125,8081.0048828125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,32768,10709.560546875,10709.560546875,10709.560546875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,128,2377.52197265625,2377.52197265625,2377.52197265625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,256,2387.03955078125,2387.03955078125,2387.03955078125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,512,2406.07470703125,2406.07470703125,2406.07470703125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,1024,2444.14501953125,2444.14501953125,2444.14501953125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,2048,2520.28564453125,2520.28564453125,2520.28564453125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,4096,2672.56787109375,2672.56787109375,2672.56787109375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,8192,2977.13232421875,2977.13232421875,2977.13232421875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,16384,3586.26123046875,3586.26123046875,3586.26123046875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,32768,4804.51904296875,4804.51904296875,4804.51904296875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:22,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,128,2388.1318359375,2388.1318359375,2388.1318359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,256,2408.1376953125,2408.1376953125,2408.1376953125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,512,2448.177734375,2448.177734375,2448.177734375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,1024,2528.26513671875,2528.26513671875,2528.26513671875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,2048,2688.43798828125,2688.43798828125,2688.43798828125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,4096,3009.0703125,3009.0703125,3009.0703125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,8192,3649.57080078125,3649.57080078125,3649.57080078125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,16384,4934.091796875,4934.091796875,4934.091796875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,32768,7496.35107421875,7496.35107421875,7496.35107421875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,128,5061.0234375,5061.0234375,5061.0234375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,256,5066.04296875,5066.04296875,5066.04296875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,512,5076.08203125,5076.08203125,5076.08203125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,1024,5096.16015625,5096.16015625,5096.16015625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,2048,5136.31640625,5136.31640625,5136.31640625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,4096,5216.6298828125,5216.6298828125,5216.6298828125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,8192,5569.6318359375,5569.6318359375,5569.6318359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,16384,6467.2607421875,6467.2607421875,6467.2607421875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,32768,8262.5185546875,8262.5185546875,8262.5185546875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:25,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,128,5466.693359375,5466.693359375,5466.693359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:46,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,256,5487.25390625,5487.25390625,5487.25390625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:46,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,512,5528.3955078125,5528.3955078125,5528.3955078125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:46,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,1024,5610.61962890625,5610.61962890625,5610.61962890625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:46,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,2048,5775.12841796875,5775.12841796875,5775.12841796875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:46,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,4096,6104.3583984375,6104.3583984375,6104.3583984375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:46,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,8192,6762.08544921875,6762.08544921875,6762.08544921875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:46,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,16384,8081.0908203125,8081.0908203125,8081.0908203125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:46,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,32768,10711.791015625,10711.791015625,10711.791015625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-07-08 08:16:46,0.8.0 fused_moe,liger,full,speed,ms,E,num_experts,8,6.273632049560547,6.273632049560547,6.273632049560547,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:59:40,0.7.0 fused_moe,liger,full,speed,ms,E,num_experts,16,6.440032005310059,6.440032005310059,6.440032005310059,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:59:40,0.7.0 fused_moe,liger,full,speed,ms,E,num_experts,32,6.884575843811035,6.884575843811035,6.884575843811035,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA H100 80GB HBM3,2026-04-02 23:59:40,0.7.0 @@ -2351,3 +2351,171 @@ megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,16384,1024.1 megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,32768,2048.16552734375,2048.16552734375,2048.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0 megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,65536,4096.16552734375,4096.16552734375,4096.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0 megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,131072,8192.166015625,8192.166015625,8192.166015625,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,128,1.990608036518097,1.988102412223816,1.999232006072998,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:28,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,256,2.100543975830078,1.8793535709381102,2.200608015060425,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:28,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,512,2.348847985267639,2.3457280158996583,2.3566271305084228,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:28,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,1024,2.3433600664138794,2.3175424575805663,2.370086431503296,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:28,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,2048,2.5968639850616455,2.5476735591888424,2.826630449295044,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:28,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,4096,3.416048049926758,3.347974443435669,3.4841216564178468,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:28,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,8192,4.754303932189941,4.754303932189941,4.754303932189941,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:28,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,16384,7.7433600425720215,7.7433600425720215,7.7433600425720215,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:28,0.8.0 +fused_moe,liger,full,speed,ms,T,num_tokens,32768,13.487071990966797,13.487071990966797,13.487071990966797,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:28,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,128,811.599853515625,811.599853515625,811.599853515625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,256,494.4394226074219,494.4394226074219,494.4394226074219,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,512,145.98655700683594,145.98655700683594,145.98655700683594,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,1024,147.2841033935547,147.2841033935547,147.2841033935547,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,2048,148.05398559570312,148.05398559570312,148.05398559570312,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,4096,149.03208923339844,149.03208923339844,149.03208923339844,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,8192,151.2623291015625,151.2623291015625,151.2623291015625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,16384,157.14816284179688,157.14816284179688,157.14816284179688,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,huggingface,full,speed,ms,T,num_tokens,32768,169.1504669189453,169.1504669189453,169.1504669189453,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,128,0.3019520044326782,0.2850815951824188,0.32361599802970886,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,256,0.3097760081291199,0.2948479950428009,0.32047998905181885,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,512,0.30004799365997314,0.2826560139656067,0.3112959861755371,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,1024,0.2969599962234497,0.2959359884262085,0.3035840094089508,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,2048,0.3676159977912903,0.36748799681663513,0.3687936127185822,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,4096,0.5484640002250671,0.5477120280265808,0.5499200224876404,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,8192,0.9251199960708618,0.9245695829391479,0.9266944169998169,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,16384,1.669119954109192,1.6566207885742188,1.671430444717407,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,liger,forward,speed,ms,T,num_tokens,32768,3.205024003982544,3.2013760566711422,3.209382343292236,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:54,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,128,27.944063186645508,27.944063186645508,27.944063186645508,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:56,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,256,28.172576904296875,28.172576904296875,28.172576904296875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:56,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,512,27.962175369262695,27.962175369262695,27.962175369262695,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:56,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,1024,29.042879104614258,29.042879104614258,29.042879104614258,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:56,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,2048,29.31395149230957,29.31395149230957,29.31395149230957,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:56,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,4096,29.45020866394043,29.45020866394043,29.45020866394043,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:56,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,8192,30.506528854370117,30.506528854370117,30.506528854370117,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:56,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,16384,29.378463745117188,29.378463745117188,29.378463745117188,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:56,0.8.0 +fused_moe,huggingface,forward,speed,ms,T,num_tokens,32768,31.45155143737793,31.45155143737793,31.45155143737793,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:56,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,128,1.4304799437522888,1.4286079406738281,1.4307199716567993,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:57,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,256,1.4474560022354126,1.4467840194702148,1.4487359523773193,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:57,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,512,1.5062080025672913,1.506111979484558,1.5085439682006836,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:57,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,1024,1.6168320178985596,1.616377592086792,1.6171711921691894,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:57,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,2048,1.8176000118255615,1.8173951625823974,1.8177983999252318,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:57,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,4096,2.1983999013900757,2.197593545913696,2.1984895706176757,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:57,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,8192,2.960383892059326,2.9602879047393795,2.963475227355957,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:57,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,16384,4.612560033798218,4.60177927017212,4.623340797424317,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:57,0.8.0 +fused_moe,liger,backward,speed,ms,T,num_tokens,32768,7.944032192230225,7.944032192230225,7.944032192230225,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:05:57,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,128,118.43968200683594,118.43968200683594,118.43968200683594,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,256,118.71110534667969,118.71110534667969,118.71110534667969,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,512,118.2043228149414,118.2043228149414,118.2043228149414,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,1024,118.9240951538086,118.9240951538086,118.9240951538086,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,2048,119.47103881835938,119.47103881835938,119.47103881835938,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,4096,120.70182037353516,120.70182037353516,120.70182037353516,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,8192,122.90444946289062,122.90444946289062,122.90444946289062,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,16384,128.1175994873047,128.1175994873047,128.1175994873047,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,huggingface,backward,speed,ms,T,num_tokens,32768,137.93463134765625,137.93463134765625,137.93463134765625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,128,5061.0234375,5061.0234375,5061.0234375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,256,5066.04296875,5066.04296875,5066.04296875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,512,5076.08203125,5076.08203125,5076.08203125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,1024,5096.16015625,5096.16015625,5096.16015625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,2048,5136.31640625,5136.31640625,5136.31640625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,4096,5216.6298828125,5216.6298828125,5216.6298828125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,8192,5569.6318359375,5569.6318359375,5569.6318359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,16384,6467.2607421875,6467.2607421875,6467.2607421875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,liger,full,memory,MB,T,num_tokens,32768,8262.5185546875,8262.5185546875,8262.5185546875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:06,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,128,5466.697265625,5466.697265625,5466.697265625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,256,5487.24609375,5487.24609375,5487.24609375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,512,5528.40087890625,5528.40087890625,5528.40087890625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,1024,5610.642578125,5610.642578125,5610.642578125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,2048,5775.08203125,5775.08203125,5775.08203125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,4096,6105.361328125,6105.361328125,6105.361328125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,8192,6762.3828125,6762.3828125,6762.3828125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,16384,8081.48876953125,8081.48876953125,8081.48876953125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,huggingface,full,memory,MB,T,num_tokens,32768,10710.45703125,10710.45703125,10710.45703125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,128,2377.52197265625,2377.52197265625,2377.52197265625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,256,2387.03955078125,2387.03955078125,2387.03955078125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,512,2406.07470703125,2406.07470703125,2406.07470703125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,1024,2444.14501953125,2444.14501953125,2444.14501953125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,2048,2520.28564453125,2520.28564453125,2520.28564453125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,4096,2672.56787109375,2672.56787109375,2672.56787109375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,8192,2977.13232421875,2977.13232421875,2977.13232421875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,16384,3586.26123046875,3586.26123046875,3586.26123046875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,liger,forward,memory,MB,T,num_tokens,32768,4804.51904296875,4804.51904296875,4804.51904296875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:20,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,128,2388.1318359375,2388.1318359375,2388.1318359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,256,2408.1376953125,2408.1376953125,2408.1376953125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,512,2448.17529296875,2448.17529296875,2448.17529296875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,1024,2528.2646484375,2528.2646484375,2528.2646484375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,2048,2688.4384765625,2688.4384765625,2688.4384765625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,4096,3009.1396484375,3009.1396484375,3009.1396484375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,8192,3650.0859375,3650.0859375,3650.0859375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,16384,4934.6767578125,4934.6767578125,4934.6767578125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,huggingface,forward,memory,MB,T,num_tokens,32768,7494.8193359375,7494.8193359375,7494.8193359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,128,5061.0234375,5061.0234375,5061.0234375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,256,5066.04296875,5066.04296875,5066.04296875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,512,5076.08203125,5076.08203125,5076.08203125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,1024,5096.16015625,5096.16015625,5096.16015625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,2048,5136.31640625,5136.31640625,5136.31640625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,4096,5216.6298828125,5216.6298828125,5216.6298828125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,8192,5569.6318359375,5569.6318359375,5569.6318359375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,16384,6467.2607421875,6467.2607421875,6467.2607421875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,liger,backward,memory,MB,T,num_tokens,32768,8262.5185546875,8262.5185546875,8262.5185546875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:23,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,128,5466.697265625,5466.697265625,5466.697265625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:34,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,256,5487.24609375,5487.24609375,5487.24609375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:34,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,512,5528.40087890625,5528.40087890625,5528.40087890625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:34,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,1024,5610.642578125,5610.642578125,5610.642578125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:34,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,2048,5775.08203125,5775.08203125,5775.08203125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:34,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,4096,6104.427734375,6104.427734375,6104.427734375,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:34,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,8192,6762.5419921875,6762.5419921875,6762.5419921875,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:34,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,16384,8081.67578125,8081.67578125,8081.67578125,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:34,0.8.0 +fused_moe,huggingface,backward,memory,MB,T,num_tokens,32768,10710.166015625,10710.166015625,10710.166015625,"{""sweep_dim"": ""T"", ""T"": null, ""E"": 128, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:06:34,0.8.0 +fused_moe,liger,full,speed,ms,E,num_experts,8,2.539520025253296,2.535219192504883,2.551078414916992,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:35,0.8.0 +fused_moe,liger,full,speed,ms,E,num_experts,16,2.714400053024292,2.7120384216308597,2.716281604766846,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:35,0.8.0 +fused_moe,liger,full,speed,ms,E,num_experts,32,2.8434560298919678,2.838598442077637,2.8453184127807614,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:35,0.8.0 +fused_moe,liger,full,speed,ms,E,num_experts,64,3.180351972579956,3.179776048660278,3.181542444229126,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:35,0.8.0 +fused_moe,liger,full,speed,ms,E,num_experts,128,3.942336082458496,3.9417216777801514,3.942950487136841,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:35,0.8.0 +fused_moe,huggingface,full,speed,ms,E,num_experts,8,6.617983818054199,6.617983818054199,6.617983818054199,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,huggingface,full,speed,ms,E,num_experts,16,11.182016372680664,11.182016372680664,11.182016372680664,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,huggingface,full,speed,ms,E,num_experts,32,20.890623092651367,20.890623092651367,20.890623092651367,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,huggingface,full,speed,ms,E,num_experts,64,51.23072052001953,51.23072052001953,51.23072052001953,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,huggingface,full,speed,ms,E,num_experts,128,151.31954956054688,151.31954956054688,151.31954956054688,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,liger,forward,speed,ms,E,num_experts,8,0.7546720206737518,0.7545920014381409,0.7561535835266113,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,liger,forward,speed,ms,E,num_experts,16,0.795520007610321,0.7917439937591553,0.7976319789886475,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,liger,forward,speed,ms,E,num_experts,32,0.822272002696991,0.8213120102882385,0.824288010597229,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,liger,forward,speed,ms,E,num_experts,64,0.853408008813858,0.8361535906791687,0.8572352170944214,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,liger,forward,speed,ms,E,num_experts,128,0.9215520024299622,0.8926911950111389,0.9491583824157714,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,huggingface,forward,speed,ms,E,num_experts,8,2.994784116744995,2.9763328075408935,3.017324829101563,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,huggingface,forward,speed,ms,E,num_experts,16,4.736224174499512,4.728985786437988,4.743462562561035,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,huggingface,forward,speed,ms,E,num_experts,32,8.254560470581055,8.254560470581055,8.254560470581055,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,huggingface,forward,speed,ms,E,num_experts,64,15.176287651062012,15.176287651062012,15.176287651062012,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,huggingface,forward,speed,ms,E,num_experts,128,29.438528060913086,29.438528060913086,29.438528060913086,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:38,0.8.0 +fused_moe,liger,backward,speed,ms,E,num_experts,8,1.7520639896392822,1.7515583992004395,1.7525760173797607,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:39,0.8.0 +fused_moe,liger,backward,speed,ms,E,num_experts,16,1.8400640487670898,1.837875247001648,1.842630434036255,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:39,0.8.0 +fused_moe,liger,backward,speed,ms,E,num_experts,32,1.9733279943466187,1.9725119590759277,1.974988842010498,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:39,0.8.0 +fused_moe,liger,backward,speed,ms,E,num_experts,64,2.312160015106201,2.3107264041900635,2.3132352352142336,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:39,0.8.0 +fused_moe,liger,backward,speed,ms,E,num_experts,128,2.978816032409668,2.977510404586792,2.9795264244079593,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:39,0.8.0 +fused_moe,huggingface,backward,speed,ms,E,num_experts,8,3.568511962890625,3.5671871662139893,3.5698367595672607,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:40,0.8.0 +fused_moe,huggingface,backward,speed,ms,E,num_experts,16,6.3691840171813965,6.3691840171813965,6.3691840171813965,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:40,0.8.0 +fused_moe,huggingface,backward,speed,ms,E,num_experts,32,12.523391723632812,12.523391723632812,12.523391723632812,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:40,0.8.0 +fused_moe,huggingface,backward,speed,ms,E,num_experts,64,36.244384765625,36.244384765625,36.244384765625,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:40,0.8.0 +fused_moe,huggingface,backward,speed,ms,E,num_experts,128,123.16547393798828,123.16547393798828,123.16547393798828,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:40,0.8.0 +fused_moe,liger,full,memory,MB,E,num_experts,8,1249.630859375,1249.630859375,1249.630859375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:40,0.8.0 +fused_moe,liger,full,memory,MB,E,num_experts,16,1537.630859375,1537.630859375,1537.630859375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:40,0.8.0 +fused_moe,liger,full,memory,MB,E,num_experts,32,2113.630859375,2113.630859375,2113.630859375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:40,0.8.0 +fused_moe,liger,full,memory,MB,E,num_experts,64,3265.630859375,3265.630859375,3265.630859375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:40,0.8.0 +fused_moe,liger,full,memory,MB,E,num_experts,128,5569.6318359375,5569.6318359375,5569.6318359375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:40,0.8.0 +fused_moe,huggingface,full,memory,MB,E,num_experts,8,1761.765625,1761.765625,1761.765625,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,huggingface,full,memory,MB,E,num_experts,16,2074.2548828125,2074.2548828125,2074.2548828125,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,huggingface,full,memory,MB,E,num_experts,32,2739.43310546875,2739.43310546875,2739.43310546875,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,huggingface,full,memory,MB,E,num_experts,64,4080.96044921875,4080.96044921875,4080.96044921875,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,huggingface,full,memory,MB,E,num_experts,128,6763.693359375,6763.693359375,6763.693359375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,liger,forward,memory,MB,E,num_experts,8,817.13134765625,817.13134765625,817.13134765625,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,liger,forward,memory,MB,E,num_experts,16,961.13134765625,961.13134765625,961.13134765625,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,liger,forward,memory,MB,E,num_experts,32,1249.13134765625,1249.13134765625,1249.13134765625,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,liger,forward,memory,MB,E,num_experts,64,1825.13134765625,1825.13134765625,1825.13134765625,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,liger,forward,memory,MB,E,num_experts,128,2977.13232421875,2977.13232421875,2977.13232421875,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,huggingface,forward,memory,MB,E,num_experts,8,1429.50048828125,1429.50048828125,1429.50048828125,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,huggingface,forward,memory,MB,E,num_experts,16,1577.9365234375,1577.9365234375,1577.9365234375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,huggingface,forward,memory,MB,E,num_experts,32,1875.61279296875,1875.61279296875,1875.61279296875,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,huggingface,forward,memory,MB,E,num_experts,64,2466.8408203125,2466.8408203125,2466.8408203125,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,huggingface,forward,memory,MB,E,num_experts,128,3651.3759765625,3651.3759765625,3651.3759765625,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:43,0.8.0 +fused_moe,liger,backward,memory,MB,E,num_experts,8,1249.630859375,1249.630859375,1249.630859375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:44,0.8.0 +fused_moe,liger,backward,memory,MB,E,num_experts,16,1537.630859375,1537.630859375,1537.630859375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:44,0.8.0 +fused_moe,liger,backward,memory,MB,E,num_experts,32,2113.630859375,2113.630859375,2113.630859375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:44,0.8.0 +fused_moe,liger,backward,memory,MB,E,num_experts,64,3265.630859375,3265.630859375,3265.630859375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:44,0.8.0 +fused_moe,liger,backward,memory,MB,E,num_experts,128,5569.6318359375,5569.6318359375,5569.6318359375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:44,0.8.0 +fused_moe,huggingface,backward,memory,MB,E,num_experts,8,1761.765625,1761.765625,1761.765625,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:45,0.8.0 +fused_moe,huggingface,backward,memory,MB,E,num_experts,16,2072.60009765625,2072.60009765625,2072.60009765625,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:45,0.8.0 +fused_moe,huggingface,backward,memory,MB,E,num_experts,32,2738.3193359375,2738.3193359375,2738.3193359375,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:45,0.8.0 +fused_moe,huggingface,backward,memory,MB,E,num_experts,64,4078.7666015625,4078.7666015625,4078.7666015625,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:45,0.8.0 +fused_moe,huggingface,backward,memory,MB,E,num_experts,128,6763.83203125,6763.83203125,6763.83203125,"{""sweep_dim"": ""E"", ""T"": 8192, ""E"": null, ""H"": 2048, ""intermediate_dim"": 768, ""K"": 8, ""dtype"": ""torch.bfloat16""}",NVIDIA B300 SXM6 AC,2026-07-08 15:09:45,0.8.0 diff --git a/benchmark/scripts/benchmark_chunked_grpo_loss_head_to_head.py b/benchmark/scripts/benchmark_chunked_grpo_loss_head_to_head.py new file mode 100644 index 000000000..b723e3840 --- /dev/null +++ b/benchmark/scripts/benchmark_chunked_grpo_loss_head_to_head.py @@ -0,0 +1,191 @@ +"""Head-to-head GRPO loss benchmark: torch chunked vs triton unchunked vs triton chunked. + +Compares the three Liger GRPO loss implementations from the hidden-state +boundary (lm_head projection included, since the chunked variants fuse it): + + chunked_torch: LigerFusedLinearGRPOLoss (fused linear, torch/cuBLAS chunking) + triton: triton_grpo_loss on materialized (B, L+1, V) logits + chunked_triton: chunked_triton_grpo_loss (fused linear, Triton kernels) + +Measures forward+backward wall time (CUDA events) and peak memory above the +resident inputs, per micro-batch. Config mirrors GRPO training on Qwen3.5-MoE: +dapo loss, sequence-level importance sampling, beta=0, temperature 1.0, +eps 0.2/0.2, hidden 2048, vocab 248320, bf16. + +Run from the repo root: + PYTHONPATH=src python benchmark/scripts/benchmark_chunked_grpo_loss_head_to_head.py +""" + +import argparse + +import torch + +from liger_kernel.chunked_loss import LigerFusedLinearGRPOLoss +from liger_kernel.transformers.chunked_grpo_loss import chunked_triton_grpo_loss +from liger_kernel.transformers.grpo_loss import triton_grpo_loss + +HIDDEN_SIZE = 2048 +VOCAB_SIZE = 248320 +LOSS_KWARGS = dict( + temperature=1.0, + beta=0.0, + eps_low=0.2, + eps_high=0.2, + loss_type="dapo", + importance_sampling_level="sequence", +) + + +def make_inputs(batch, seq_len, device, seed=0): + gen = torch.Generator(device=device).manual_seed(seed) + hidden = torch.randn(batch, seq_len + 1, HIDDEN_SIZE, device=device, generator=gen).to(torch.bfloat16).mul_(0.02) + weight = torch.randn(VOCAB_SIZE, HIDDEN_SIZE, device=device, generator=gen).to(torch.bfloat16).mul_(0.02) + completion_ids = torch.randint(0, VOCAB_SIZE, (batch, seq_len), device=device, generator=gen) + lengths = torch.randint(seq_len // 2, seq_len + 1, (batch,), device=device, generator=gen) + mask = (torch.arange(seq_len, device=device).unsqueeze(0) < lengths.unsqueeze(1)).float() + advantages = torch.randn(batch, device=device, dtype=torch.float32, generator=gen) + return { + "hidden": hidden, + "weight": weight, + "completion_ids": completion_ids, + "mask": mask, + "advantages": advantages, + "num_items_in_batch": mask.sum(), + } + + +def run_variant(variant, inputs, chunked_torch_module): + hidden, weight = inputs["hidden"], inputs["weight"] + common = dict(num_items_in_batch=inputs["num_items_in_batch"]) + if variant == "chunked_torch": + loss, _ = chunked_torch_module( + hidden[:, :-1, :], + weight, + inputs["completion_ids"], + inputs["mask"], + inputs["advantages"], + **common, + ) + elif variant == "triton": + logits = hidden @ weight.t() + loss, _ = triton_grpo_loss( + logits, + None, + None, + inputs["completion_ids"], + inputs["advantages"], + inputs["mask"], + inplace=True, + reduce=True, + **LOSS_KWARGS, + **common, + ) + elif variant == "chunked_triton": + loss, _ = chunked_triton_grpo_loss( + hidden[:, :-1, :].contiguous(), + weight, + None, + None, + inputs["completion_ids"], + inputs["advantages"], + inputs["mask"], + reduce=True, + **LOSS_KWARGS, + **common, + ) + else: + raise ValueError(variant) + return loss + + +def bench(variant, inputs, chunked_torch_module, warmup, iters): + hidden, weight = inputs["hidden"], inputs["weight"] + + def once(): + hidden.grad = None + weight.grad = None + hidden.requires_grad_(True) + weight.requires_grad_(True) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + baseline = torch.cuda.memory_allocated() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + loss = run_variant(variant, inputs, chunked_torch_module) + loss.backward() + end.record() + torch.cuda.synchronize() + peak = (torch.cuda.max_memory_allocated() - baseline) / 1024**3 + hidden.requires_grad_(False) + weight.requires_grad_(False) + return start.elapsed_time(end), peak + + try: + for _ in range(warmup): + once() + times, peaks = zip(*[once() for _ in range(iters)]) + except torch.cuda.OutOfMemoryError: + torch.cuda.empty_cache() + return "OOM" + except RuntimeError as err: + return f"FAIL: {err}" + t = torch.tensor(times) + return t.mean().item(), t.std().item(), max(peaks) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch-size", type=int, default=4) + parser.add_argument("--seq-lens", type=int, nargs="+", default=[1024, 4096, 16384, 32768, 65535, 65536]) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iters", type=int, default=5) + args = parser.parse_args() + + device = torch.device("cuda:0") + torch.cuda.set_device(device) + chunked_torch_module = LigerFusedLinearGRPOLoss( + beta=0.0, + compiled=False, + use_ref_model=False, + epsilon_low=LOSS_KWARGS["eps_low"], + epsilon_high=LOSS_KWARGS["eps_high"], + loss_type=LOSS_KWARGS["loss_type"], + importance_sampling_level=LOSS_KWARGS["importance_sampling_level"], + temperature=LOSS_KWARGS["temperature"], + ) + print(f"Device: {torch.cuda.get_device_name(device)}") + print(f"Hidden {HIDDEN_SIZE}, vocab {VOCAB_SIZE}, batch {args.batch_size}, bf16") + print(f"Config: {LOSS_KWARGS}\n") + + # quick loss parity sanity check + inputs = make_inputs(args.batch_size, 1024, device) + losses = { + v: run_variant(v, inputs, chunked_torch_module).item() for v in ["chunked_torch", "triton", "chunked_triton"] + } + print(f"Loss parity @1024: {losses}\n") + del inputs + + variants = ["chunked_torch", "triton", "chunked_triton"] + header = f"{'seq_len':>8} {'logits_GiB':>11}" + "".join(f" {v + '_ms':>22} {v + '_peak_GiB':>18}" for v in variants) + print(header) + print("-" * len(header)) + for seq_len in args.seq_lens: + inputs = make_inputs(args.batch_size, seq_len, device) + logits_gib = args.batch_size * (seq_len + 1) * VOCAB_SIZE * 2 / 1024**3 + row = f"{seq_len:>8} {logits_gib:>11.1f}" + for variant in variants: + result = bench(variant, inputs, chunked_torch_module, args.warmup, args.iters) + if isinstance(result, str): + label = result if len(result) < 20 else "LAUNCH FAIL" + row += f" {label:>22} {'-':>18}" + else: + mean_ms, std_ms, peak = result + row += f" {mean_ms:>14.1f} ±{std_ms:>5.1f} {peak:>18.2f}" + print(row) + del inputs + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/benchmark/scripts/benchmark_fused_moe.py b/benchmark/scripts/benchmark_fused_moe.py index c4e83e28e..d81f90d30 100644 --- a/benchmark/scripts/benchmark_fused_moe.py +++ b/benchmark/scripts/benchmark_fused_moe.py @@ -25,6 +25,7 @@ from utils import run_speed_benchmark from liger_kernel.ops import LigerFusedMoEFunction +from liger_kernel.ops.fused_moe import _pick_block_m_token from liger_kernel.utils import get_total_gpu_memory from liger_kernel.utils import infer_device @@ -133,12 +134,11 @@ def bench_memory_fused_moe(input: SingleBenchmarkRunInput) -> SingleBenchmarkRun def _warmup_liger(T, E, H, intermediate_dim, K, dtype, sweep_dim): - """Run one full fwd+bwd to exhaust Triton autotune for (H, intermediate_dim). + """Run one full fwd+bwd to exhaust Triton autotune for one autotune key. - Triton autotune key is (H_dim, I_dim), so a single call is sufficient to - cache the best config for all subsequent calls with the same H and intermediate_dim. - For the num_experts sweep we also call this once per E value to warm up - CUDA caches for each expert count before do_bench starts timing. + The GEMM autotune key is (H_dim, I_dim, BLOCK_M[, USE_TMA]) where BLOCK_M is + picked adaptively from tokens-per-expert, so the caller warms one + representative point per distinct BLOCK_M bucket of the sweep. """ warmup_input = SingleBenchmarkRunInput( x=T if sweep_dim == "T" else E, @@ -221,34 +221,6 @@ def _probe(): peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) kernel_bpt = peak_bytes // probe_T - # Pre-warm Liger's Triton autotune before benchmarks start. - # - # Autotune key is (H_dim, I_dim) — one warmup per (H, intermediate_dim) pair is sufficient - # to cache the best config for the entire sweep. - # - # For num_tokens sweep: one pass with the model's base T is enough. - # For num_experts sweep: one pass per E value in EXPERT_SWEEP_VALUES to also - # warm up CUDA caches for each expert count, since weight tensor sizes differ. - print(f"Pre-warming Liger autotune (H={H}, intermediate_dim={intermediate_dim})...") - - if args.sweep_dim == "num_tokens": - _warmup_liger(probe_T, E, H, intermediate_dim, K, dtype, sweep_dim="T") - else: # num_experts - for e_val in EXPERT_SWEEP_VALUES: - print(f" warmup E={e_val}...") - _warmup_liger(probe_T, e_val, H, intermediate_dim, K, dtype, sweep_dim="E") - - if device == "cuda": - torch.cuda.synchronize() - elif device == "npu": - torch.npu.synchronize() - elif device == "xpu": - torch.xpu.synchronize() - else: - torch.cpu.synchronize() - - print("Autotune warmup complete.\n") - if args.sweep_dim == "num_tokens": # Derive a memory-safe upper bound for T from the probe measurement. # Target 40% GPU memory utilisation to leave headroom for framework overhead. @@ -284,6 +256,37 @@ def _probe(): ] x_name, x_label = "E", "num_experts" + # Pre-warm Liger's Triton autotune before benchmarks start. + # + # The GEMM autotune key includes the adaptive BLOCK_M (a function of tokens per + # expert), so warm one representative x-value per distinct BLOCK_M bucket. + # For the num_experts sweep this also warms CUDA caches per expert count. + print(f"Pre-warming Liger autotune (H={H}, intermediate_dim={intermediate_dim})...") + + if args.sweep_dim == "num_tokens": + warmed = set() + for t_val in x_values: + bucket = _pick_block_m_token(t_val * K, E) + if bucket not in warmed: + print(f" warmup T={t_val} (BLOCK_M={bucket})...") + _warmup_liger(t_val, E, H, intermediate_dim, K, dtype, sweep_dim="T") + warmed.add(bucket) + else: # num_experts + for e_val in EXPERT_SWEEP_VALUES: + print(f" warmup E={e_val}...") + _warmup_liger(probe_T, e_val, H, intermediate_dim, K, dtype, sweep_dim="E") + + if device == "cuda": + torch.cuda.synchronize() + elif device == "npu": + torch.npu.synchronize() + elif device == "xpu": + torch.xpu.synchronize() + else: + torch.cpu.synchronize() + + print("Autotune warmup complete.\n") + common_configs = { "kernel_name": "fused_moe", "x_name": x_name, diff --git a/src/liger_kernel/chunked_loss/fused_linear_ppo.py b/src/liger_kernel/chunked_loss/fused_linear_ppo.py index 28bd5f11c..aa4654488 100644 --- a/src/liger_kernel/chunked_loss/fused_linear_ppo.py +++ b/src/liger_kernel/chunked_loss/fused_linear_ppo.py @@ -4,8 +4,12 @@ import torch import torch._dynamo.config -_SELECTIVE_LOGPROB_VOCAB_CHUNK_SIZE = 4096 -_SELECTIVE_LOGPROB_SEQ_CHUNK_SIZE = 2048 +# Chunk temporaries are seq_chunk x vocab_chunk fp32 (~128 MB at these sizes) — +# negligible next to the fp32 grad_weight buffer, and larger chunks amortize the +# per-chunk elementwise/launch overhead: 4096x8192 measured 21% faster than the +# previous 2048x4096 at identical peak memory (B300, V=248320, 65K tokens). +_SELECTIVE_LOGPROB_VOCAB_CHUNK_SIZE = 8192 +_SELECTIVE_LOGPROB_SEQ_CHUNK_SIZE = 4096 def _maybe_mark_dynamic_dim1(tensor): @@ -72,6 +76,16 @@ def _selective_logprob_backward(hidden, weight, targets, bias, log_z, grad_logpr """Dual-chunked (sequence × vocab) backward for selective logprob. Recomputes logits per chunk for memory efficiency. + + The two grad GEMMs run with operands in ``hidden.dtype`` (bf16/fp16 in + practice) rather than fp32: fp32×fp32 matmuls dispatch to SIMT CUDA-core + kernels (~57 TFLOPS on B300, ~20x below the bf16 tensor-core rate) and were + ~78% of this backward's runtime. Precision is preserved where it matters: + cuBLAS accumulates each chunk GEMM in fp32 internally, and the cross-chunk + accumulation buffers (grad_hidden/grad_weight) stay fp32. Only the per-chunk + GEMM inputs/outputs round to the compute dtype — the same rounding a + non-chunked autograd backward through a bf16 lm_head applies everywhere. + For fp32 inputs the casts are no-ops and behavior is unchanged. """ inv_t = 1.0 / temperature n_rows, _ = hidden.shape @@ -109,8 +123,11 @@ def _selective_logprob_backward(hidden, weight, targets, bias, log_z, grad_logpr grad_logits[row_idx, local_idx] += grad_chunk * in_chunk grad_logits.mul_(inv_t) - grad_hidden[seq_start:seq_end].add_(grad_logits @ weight_chunk.float()) - grad_weight[vocab_start:vocab_end].add_(grad_logits.t() @ hidden_chunk.float()) + # Tensor-core GEMMs in the input dtype; fp32 accumulation across + # chunks via the fp32 grad buffers (see docstring). + grad_logits_lp = grad_logits.to(hidden.dtype) + grad_hidden[seq_start:seq_end].add_(grad_logits_lp @ weight_chunk.to(hidden.dtype)) + grad_weight[vocab_start:vocab_end].add_(grad_logits_lp.t() @ hidden_chunk) if has_bias: grad_bias[vocab_start:vocab_end].add_(grad_logits.sum(dim=0)) diff --git a/src/liger_kernel/ops/chunked_grpo_loss.py b/src/liger_kernel/ops/chunked_grpo_loss.py new file mode 100644 index 000000000..6a42ef47d --- /dev/null +++ b/src/liger_kernel/ops/chunked_grpo_loss.py @@ -0,0 +1,330 @@ +"""Chunked selective log-softmax through the lm_head for GRPO-style losses. + +Computes per-token log-probabilities (and logsumexp) directly from hidden +states and the lm_head weight without ever materializing the (N, V) logits +tensor, giving the memory profile of the chunked fused-linear GRPO loss with +Triton-kernel speed: + + forward: one fused kernel streams vocab tiles per 128-token program, + computing each logits tile with tl.dot GEMMs into a fresh + accumulator and maintaining an online logsumexp (flash-attention + style). Only per-token logp / lse (fp32) are written to HBM. + backward: a fused kernel recomputes logits tiles and emits grad_logits for + one sequence chunk into a reusable buffer in the input dtype; the + two large grad GEMMs (grad_hidden, grad_weight) run in cuBLAS with + an fp32 cross-chunk grad_weight accumulator. No atomics. + +sm_103 (B300) notes: + - triton-lang/triton#10821: two tl.dot calls chained through one accumulator + in a K-loop miscompile on tcgen05 tiles (BLOCK_M >= 64). Every K-loop here + issues exactly ONE tl.dot per accumulator per iteration and each vocab tile + gets a fresh accumulator, so the buggy pattern never occurs. + - Launch configs are fixed (no autotune) and there are no atomics, so results + are bitwise reproducible run-to-run and rank-to-rank. +""" + +import torch +import triton +import triton.language as tl + +# Tile sizes. BLOCK_M >= 64 selects the tcgen05 MMA path on sm_103, which is +# safe given the single-dot-per-accumulator structure (see module docstring). +# Chosen by a manual sweep on B300 at V=248320, H=2048, 65k tokens: 249 ms +# fwd+bwd vs 377 ms for (64, 128, 128); larger BN/BK exhaust shared memory at +# num_stages=3, and num_warps=16 regresses ~2x. +_BM = 128 +_BN = 256 +_BK = 64 +_NUM_WARPS = 8 +_NUM_STAGES = 3 +# fp32 inputs double the tile footprint; 3 stages exceeds sm_103's 228 KB SMEM +# (needs ~288 KB), so drop to 2 pipeline stages for fp32. +_NUM_STAGES_FP32 = 2 + + +def _num_stages(dtype: torch.dtype) -> int: + return _NUM_STAGES_FP32 if dtype == torch.float32 else _NUM_STAGES + + +# Sequence-chunk size for the backward grad_logits buffer +# (4096 x 248320 bf16 ~= 1.9 GiB). +_BWD_SEQ_CHUNK = 4096 + + +@triton.jit +def _chunked_selective_logp_fwd_kernel( + HIDDEN, + W, + TARGETS, + LOGP, + LSE, + N, + V, + stride_hn, + stride_wv, + inv_temp, + H: tl.constexpr, + V_CEIL: tl.constexpr, + EVEN_V: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, +): + pid = tl.program_id(0) + rows = pid * BM + tl.arange(0, BM) + row_mask = rows < N + tgt = tl.load(TARGETS + rows, mask=row_mask, other=-1) + + m_i = tl.full((BM,), float("-inf"), tl.float32) + l_i = tl.zeros((BM,), tl.float32) + t_logit = tl.zeros((BM,), tl.float32) + + for v0 in tl.range(0, V_CEIL, BN): + cols = v0 + tl.arange(0, BN) + col_mask = cols < V + acc = tl.zeros((BM, BN), tl.float32) + for k0 in tl.range(0, H, BK): + k = k0 + tl.arange(0, BK) + a = tl.load( + HIDDEN + rows[:, None] * stride_hn + k[None, :], + mask=row_mask[:, None], + other=0.0, + ) + if EVEN_V: + b = tl.load(W + cols[:, None] * stride_wv + k[None, :]) + else: + b = tl.load( + W + cols[:, None] * stride_wv + k[None, :], + mask=col_mask[:, None], + other=0.0, + ) + acc = tl.dot(a, tl.trans(b), acc) # single dot per acc per iteration + logits = acc * inv_temp + if not EVEN_V: + logits = tl.where(col_mask[None, :], logits, float("-inf")) + tile_max = tl.max(logits, axis=1) + m_new = tl.maximum(m_i, tile_max) + l_i = l_i * tl.exp(m_i - m_new) + tl.sum(tl.exp(logits - m_new[:, None]), axis=1) + m_i = m_new + is_tgt = tgt[:, None] == cols[None, :] + t_logit += tl.sum(tl.where(is_tgt, logits, 0.0), axis=1) + + lse = m_i + tl.log(l_i) + tl.store(LOGP + rows, t_logit - lse, mask=row_mask) + tl.store(LSE + rows, lse, mask=row_mask) + + +@triton.jit +def _chunked_grad_logits_kernel( + HIDDEN, + W, + TARGETS, + LSE, + GRAD_LOGP, + GRAD_LOGITS, + N, + V, + row0, + stride_hn, + stride_wv, + stride_gl, + inv_temp, + H: tl.constexpr, + EVEN_V: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + rows = row0 + pid_m * BM + tl.arange(0, BM) + row_mask = rows < N + cols = pid_n * BN + tl.arange(0, BN) + col_mask = cols < V + + acc = tl.zeros((BM, BN), tl.float32) + for k0 in tl.range(0, H, BK): + k = k0 + tl.arange(0, BK) + a = tl.load( + HIDDEN + rows[:, None] * stride_hn + k[None, :], + mask=row_mask[:, None], + other=0.0, + ) + if EVEN_V: + b = tl.load(W + cols[:, None] * stride_wv + k[None, :]) + else: + b = tl.load( + W + cols[:, None] * stride_wv + k[None, :], + mask=col_mask[:, None], + other=0.0, + ) + acc = tl.dot(a, tl.trans(b), acc) # single dot per acc per iteration + logits = acc * inv_temp + + lse = tl.load(LSE + rows, mask=row_mask, other=0.0) + p = tl.exp(logits - lse[:, None]) + if not EVEN_V: + p = tl.where(col_mask[None, :], p, 0.0) + tgt = tl.load(TARGETS + rows, mask=row_mask, other=-1) + g = tl.load(GRAD_LOGP + rows, mask=row_mask, other=0.0) + is_tgt = tgt[:, None] == cols[None, :] + # d logp_target / d raw_logit_v = (delta_{v==target} - softmax_v) / T + gl = (tl.where(is_tgt, 1.0, 0.0) - p) * (g * inv_temp)[:, None] + + lrow = pid_m * BM + tl.arange(0, BM) + store_mask = row_mask[:, None] if EVEN_V else (row_mask[:, None] & col_mask[None, :]) + tl.store( + GRAD_LOGITS + lrow[:, None] * stride_gl + cols[None, :], + gl.to(GRAD_LOGITS.dtype.element_ty), + mask=store_mask, + ) + + +class ChunkedSelectiveLogPFunction(torch.autograd.Function): + """Per-token selective log-softmax through the lm_head, chunked over vocab. + + forward(hidden (N, H), weight (V, H), targets (N,)) -> logp (N,) fp32 + """ + + @staticmethod + def forward(ctx, hidden, weight, targets, temperature): + assert hidden.is_contiguous(), "hidden must be contiguous" + assert weight.is_contiguous(), "weight must be contiguous" + n_tokens, h = hidden.shape + v = weight.shape[0] + assert h % _BK == 0, f"hidden size {h} must be divisible by {_BK}" + even_v = v % _BN == 0 + v_ceil = triton.cdiv(v, _BN) * _BN + + logp = torch.empty(n_tokens, dtype=torch.float32, device=hidden.device) + lse = torch.empty(n_tokens, dtype=torch.float32, device=hidden.device) + grid = (triton.cdiv(n_tokens, _BM),) + _chunked_selective_logp_fwd_kernel[grid]( + hidden, + weight, + targets, + logp, + lse, + n_tokens, + v, + hidden.stride(0), + weight.stride(0), + 1.0 / temperature, + H=h, + V_CEIL=v_ceil, + EVEN_V=even_v, + BM=_BM, + BN=_BN, + BK=_BK, + num_warps=_NUM_WARPS, + num_stages=_num_stages(hidden.dtype), + ) + ctx.save_for_backward(hidden, weight, targets, lse) + ctx.temperature = temperature + # Record dims at forward time: under sharded training (ZeRO-3 / FSDP) the + # weight tensor's storage may be repartitioned between forward and backward, + # so its shape cannot be re-read in backward. + ctx.vocab_size = v + return logp + + @staticmethod + def backward(ctx, grad_logp): + hidden, weight, targets, lse = ctx.saved_tensors + n_tokens, h = hidden.shape + v = ctx.vocab_size + if weight.numel() != v * h: + raise RuntimeError( + f"lm_head weight is not materialized in backward (numel={weight.numel()}, expected {v * h}). " + "Under ZeRO-3/FSDP the fused loss must be invoked inside the weight-owning module's forward " + "(e.g. via forward redirection through lm_head) so the sharding wrapper re-gathers it for the " + "backward window." + ) + even_v = v % _BN == 0 + grad_logp = grad_logp.contiguous() + + grad_hidden = torch.empty_like(hidden) + # fp32 cross-chunk accumulator; the chunk GEMMs run in the input dtype. + grad_weight = torch.zeros(v, h, dtype=torch.float32, device=weight.device) + chunk = min(_BWD_SEQ_CHUNK, n_tokens) + buf = torch.empty(chunk, v, dtype=hidden.dtype, device=hidden.device) + + for row0 in range(0, n_tokens, chunk): + n = min(chunk, n_tokens - row0) + grid = (triton.cdiv(n, _BM), triton.cdiv(v, _BN)) + _chunked_grad_logits_kernel[grid]( + hidden, + weight, + targets, + lse, + grad_logp, + buf, + n_tokens, + v, + row0, + hidden.stride(0), + weight.stride(0), + buf.stride(0), + 1.0 / ctx.temperature, + H=h, + EVEN_V=even_v, + BM=_BM, + BN=_BN, + BK=_BK, + num_warps=_NUM_WARPS, + num_stages=_num_stages(hidden.dtype), + ) + gl = buf[:n] + torch.matmul(gl, weight, out=grad_hidden[row0 : row0 + n]) + grad_weight.add_(gl.t() @ hidden[row0 : row0 + n]) + + return grad_hidden, grad_weight.to(weight.dtype), None, None + + +@torch.no_grad() +def chunked_selective_log_softmax_with_lse(hidden, weight, targets, temperature=1.0): + """No-grad variant returning (logp, lse), both (N,) fp32. For testing/inspection.""" + assert hidden.is_contiguous() and weight.is_contiguous() + n_tokens, h = hidden.shape + v = weight.shape[0] + assert h % _BK == 0, f"hidden size {h} must be divisible by {_BK}" + even_v = v % _BN == 0 + v_ceil = triton.cdiv(v, _BN) * _BN + logp = torch.empty(n_tokens, dtype=torch.float32, device=hidden.device) + lse = torch.empty(n_tokens, dtype=torch.float32, device=hidden.device) + grid = (triton.cdiv(n_tokens, _BM),) + _chunked_selective_logp_fwd_kernel[grid]( + hidden, + weight, + targets, + logp, + lse, + n_tokens, + v, + hidden.stride(0), + weight.stride(0), + 1.0 / temperature, + H=h, + V_CEIL=v_ceil, + EVEN_V=even_v, + BM=_BM, + BN=_BN, + BK=_BK, + num_warps=_NUM_WARPS, + num_stages=_num_stages(hidden.dtype), + ) + return logp, lse + + +def chunked_selective_log_softmax(hidden, weight, targets, temperature=1.0): + """Differentiable selective log-softmax through the lm_head without logits. + + Args: + hidden: (N, H) hidden states (bf16/fp16/fp32, contiguous). + weight: (V, H) lm_head weight. + targets: (N,) token ids to select. + temperature: softmax temperature applied to logits. + + Returns: + (N,) fp32 per-token log-probabilities of the targets. + """ + return ChunkedSelectiveLogPFunction.apply(hidden, weight, targets, temperature) diff --git a/src/liger_kernel/ops/fused_moe.py b/src/liger_kernel/ops/fused_moe.py index 7496b7c00..9c61904e2 100644 --- a/src/liger_kernel/ops/fused_moe.py +++ b/src/liger_kernel/ops/fused_moe.py @@ -1,13 +1,37 @@ """ Fused MoE expert computation via Triton grouped GEMM. -Forward: routing metadata (3 kernels) → fused gather+GEMM+SwiGLU → down-proj → token aggregation -Backward: memory-efficient — recomputes dA' = dO@W2^T to avoid caching Y (TK×H bytes) +Forward: routing metadata (3 kernels, sync-free) → fused gather+GEMM+SwiGLU → +down-proj → token aggregation. +Backward: memory-efficient — recomputes dA' = dO@W2^T instead of caching Y (TK×H) +and accumulates router-score gradients in fp32. + +Runtime properties: +- No host↔device sync anywhere on the hot path (the m-tile count is upper-bounded + host-side and GEMM CTAs early-exit past the device-side actual count), so the op + is CUDA-graph-capturable. +- Tile sizes adapt to tokens-per-expert; grouped GEMMs autotune per + (H, I, BLOCK_M, TMA) and the dW kernels per (H, I, tokens-per-expert bucket). +- Expert-weight loads use TMA descriptors on Hopper+ when shapes are 16B-aligned. +- Blackwell datacenter parts (sm100/sm103) tune over an extended config space + (wide-N/deep-stage tiles enabled by TMEM accumulators); other archs are unchanged. +- Inference (no input requires grad) skips saving/storing pre-activations. + +Env flags: +- LIGER_FUSED_MOE_AUTOTUNE=0: pin one config per kernel (skip tuning; see #1246). +- LIGER_FUSED_MOE_MEMORY_EFFICIENT=1: backward writes SwiGLU gradients in place + over the saved pre-activations (saves TK*2I*itemsize bytes) and drops the + (TK, I) weighted_act buffer (dW2 recomputes s_k*silu(gate)*up on the fly; + saves TK*I*itemsize bytes). Combined ≈ −1.2 GB peak at T=32768/I=768/bf16 for a + ~10-15% slower backward. In this mode a second backward over the same graph + (retain_graph) raises a version-counter error — by design. """ import torch import triton +import liger_kernel.ops.fused_moe_kernels as _kernels_mod + from liger_kernel.ops.fused_moe_kernels import _fused_down_proj_kernel from liger_kernel.ops.fused_moe_kernels import _fused_up_proj_swiglu_kernel from liger_kernel.ops.fused_moe_kernels import _moe_bwd_down_proj_kernel @@ -20,13 +44,50 @@ from liger_kernel.ops.fused_moe_kernels import _token_gather_weighted_sum_kernel from liger_kernel.ops.utils import ensure_contiguous -# Token-dimension tile size for M. -# Not in the inner-loop autotune because tile_row_start/tile_expert and the -# grid dim-0 (num_m_tiles) must be recomputed for every candidate value. -# To tune: change this constant and re-run benchmarks. +# LIGER_FUSED_MOE_MEMORY_EFFICIENT=1 → backward trades a little speed and +# retain_graph support for TK*(2I + I) bytes of peak memory (see module docstring). +# Must be set before importing liger_kernel (the kernels module reads it at import +# to configure the autotuner's restore_value for the in-place alias). +_MEMORY_EFFICIENT = _kernels_mod._MEMORY_EFFICIENT + + +# Device-side TMA descriptors (tl.make_tensor_descriptor) need a global-memory +# scratch allocator. Triton stores it in a ContextVar, which does NOT propagate to +# the autograd engine's backward thread — so it is (re)registered at the top of +# both forward() and backward() instead of only at import time. +def _tma_alloc_fn(size: int, alignment: int, stream): + return torch.empty(size, device="cuda", dtype=torch.int8) + + +def _ensure_tma_allocator(): + triton.set_allocator(_tma_alloc_fn) + + +# Fallback M-tile size (used when callers pass no explicit block_m_token). +# The autograd function picks it adaptively per call via _pick_block_m_token. BLOCK_M_TOKEN = 64 +def _tma_eligibility(t, H: int, intermediate_dim: int, E: int): + """TMA needs sm90+, 16-byte-aligned rows, and int32-addressable row counts. + W1 view is (E*2I, H) (row stride H), W2 view is (E*H, I) (row stride I).""" + if t.device.type != "cuda" or torch.cuda.get_device_capability(t.device)[0] < 9: + return False, False + itemsize = t.element_size() + w1_ok = (H * itemsize) % 16 == 0 and E * 2 * intermediate_dim < 2**31 + w2_ok = (intermediate_dim * itemsize) % 16 == 0 and E * H < 2**31 + return w1_ok, w2_ok + + +def _pick_block_m_token(TK: int, E: int) -> int: + """Match the M-tile to the expected expert segment length: large tiles at high + occupancy amortize weight re-reads; small tiles at low occupancy avoid running + mostly-padded MMAs (e.g. T=128,K=8,E=128 → 8 tokens/expert → 87% padding at 64).""" + avg = max(1, TK // max(1, E)) + b = triton.next_power_of_2(avg) + return max(16, min(128, b)) + + # --------------------------------------------------------------------------- # Routing metadata # --------------------------------------------------------------------------- @@ -35,8 +96,10 @@ def compute_routing_metadata(topk_indices: torch.Tensor, E: int, block_m_token: int = BLOCK_M_TOKEN): """Compute token→expert routing permutation metadata via 3 Triton kernels. - Also computes GPU tile metadata (tile_row_start, tile_expert) inside - Kernel 3 — no CPU loop, one .item() sync for num_m_tiles allocation. + Fully sync-free: tile metadata is allocated at a host-computable upper bound + (TK//block_m_token + min(E, TK)); the actual m-tile count stays on device in + expert_tile_offset[E] and the GEMM kernels early-exit CTAs past it. No CPU + loop, no .item(), so the whole path is CUDA-graph-capturable. Args: topk_indices: (T, K) int32 — pre-computed top-k expert indices per token @@ -44,13 +107,15 @@ def compute_routing_metadata(topk_indices: torch.Tensor, E: int, block_m_token: block_m_token: BLOCK_M for token-dimension tiling (default BLOCK_M_TOKEN) Returns: - expert_token_count: (E,) int32 - expert_start_idx: (E+1,) int32 - x_gather_idx: (TK,) int32 - s_scatter_idx: (TK,) int32 - s_reverse_scatter_idx: (TK,) int32 - tile_row_start: (num_m_tiles,) int32 — absolute row_start per M-tile - tile_expert: (num_m_tiles,) int32 — expert index per M-tile + expert_token_count: (E,) int32 + expert_start_idx: (E+1,) int32 + x_gather_idx: (TK,) int32 + s_scatter_idx: (TK,) int32 + s_reverse_scatter_idx: (TK,) int32 + tile_row_start: (num_m_tiles_max,) int32 — absolute row_start per M-tile + tile_expert: (num_m_tiles_max,) int32 — expert index per M-tile + expert_tile_offset: (E+1,) int32 — cumsum of per-expert tile + counts; [E] holds the actual total m-tile count """ T, K = topk_indices.shape TK = T * K @@ -92,11 +157,13 @@ def compute_routing_metadata(topk_indices: torch.Tensor, E: int, block_m_token: BLOCK_M_TOKEN=block_m_token, ) - # One sync to get num_m_tiles for buffer allocation and GEMM grid. - num_m_tiles = int(expert_tile_offset[-1].item()) + # No host sync: allocate tile metadata at the worst-case bound and let GEMM + # CTAs past the actual count (expert_tile_offset[E], read on device) exit early. + # Bound: sum_e ceil(f_e / B) <= floor(TK / B) + #nonempty_experts <= TK//B + min(E, TK). + num_m_tiles_max = TK // block_m_token + min(E, TK) - tile_row_start = torch.empty(num_m_tiles, dtype=torch.int32, device=device) - tile_expert = torch.empty(num_m_tiles, dtype=torch.int32, device=device) + tile_row_start = torch.empty(num_m_tiles_max, dtype=torch.int32, device=device) + tile_expert = torch.empty(num_m_tiles_max, dtype=torch.int32, device=device) # Kernel 3: sort by expert + scatter permutation arrays + tile metadata s_scatter_idx = torch.empty(TK, dtype=torch.int32, device=device) @@ -130,13 +197,14 @@ def compute_routing_metadata(topk_indices: torch.Tensor, E: int, block_m_token: s_reverse_scatter_idx, tile_row_start, tile_expert, + expert_tile_offset, ) def _token_aggregation(Y, topk_weights_flat, s_reverse_scatter_idx, T, K, H): """Weighted gather-sum: out[t] = sum_k w[t,k] * Y[s_rev[t*K+k]].""" out = torch.empty(T, H, dtype=Y.dtype, device=Y.device) - _token_gather_weighted_sum_kernel[(T,)]( + _token_gather_weighted_sum_kernel[lambda meta: (T, triton.cdiv(H, meta["BLOCK_H"]))]( Y, topk_weights_flat, s_reverse_scatter_idx, @@ -160,8 +228,8 @@ def _token_aggregation(Y, topk_weights_flat, s_reverse_scatter_idx, T, K, H): class LigerFusedMoEFunction(torch.autograd.Function): """Fused grouped GEMM MoE forward + memory-efficient backward. - Forward: routing metadata → fused gather+GEMM+SwiGLU → down-proj → token aggregation - Backward: avoids caching Y (TK×H) by recomputing dA' = dO@W2^T in backward + Forward: routing metadata → fused gather+GEMM+SwiGLU → down-proj → token aggregation. + Backward: avoids caching Y (TK×H) by recomputing dA' = dO@W2^T. Troubleshooting: If Triton's autotune ``do_bench`` loop OOMs (each config holds its own @@ -169,6 +237,10 @@ class LigerFusedMoEFunction(torch.autograd.Function): importing liger_kernel to pin each kernel to a single config and skip the benchmark loop. Temporary escape hatch until triton's autotuner handles such errors itself. + + Set ``LIGER_FUSED_MOE_MEMORY_EFFICIENT=1`` to shave another ~TK*3I*itemsize + bytes off backward peak memory (in-place SwiGLU backward + weighted_act + recompute) at ~10-15% slower backward; retain_graph re-backward then raises. """ @staticmethod @@ -190,6 +262,15 @@ def forward(ctx, x, gate_up_proj, down_proj, top_k_index, top_k_weights): intermediate_dim = gate_up_proj.shape[1] // 2 TK = T * K + block_m_token = _pick_block_m_token(TK, E) + use_tma_w1, use_tma_w2 = _tma_eligibility(x, H, intermediate_dim, E) + if use_tma_w1 or use_tma_w2: + _ensure_tma_allocator() + # Inference: no input needs grad → skip saving/storing pre_act (TK×2I bytes). + needs_grad = ( + x.requires_grad or gate_up_proj.requires_grad or down_proj.requires_grad or (top_k_weights.requires_grad) + ) + with torch.no_grad(): ( _, @@ -199,23 +280,29 @@ def forward(ctx, x, gate_up_proj, down_proj, top_k_index, top_k_weights): s_reverse_scatter_idx, tile_row_start, tile_expert, - ) = compute_routing_metadata(top_k_index, E) + expert_tile_offset, + ) = compute_routing_metadata(top_k_index, E, block_m_token) - num_m_tiles = tile_row_start.shape[0] + num_m_tiles = tile_row_start.shape[0] # upper bound; actual count lives on device + total_tiles_dev = expert_tile_offset[E:] - pre_act = torch.empty(TK, 2 * intermediate_dim, dtype=x.dtype, device=x.device) post_act = torch.empty(TK, intermediate_dim, dtype=x.dtype, device=x.device) + # pre_act only exists in training; in inference the kernel skips the store + # entirely (post_act doubles as a dummy pointer that is never written). + pre_act = torch.empty(TK, 2 * intermediate_dim, dtype=x.dtype, device=x.device) if needs_grad else post_act if num_m_tiles > 0: - _fused_up_proj_swiglu_kernel[lambda meta: (num_m_tiles, triton.cdiv(intermediate_dim, meta["BLOCK_N"]))]( + _fused_up_proj_swiglu_kernel[lambda meta: (num_m_tiles * triton.cdiv(intermediate_dim, meta["BLOCK_N"]),)]( x, gate_up_proj, x_gather_idx, expert_start_idx, tile_row_start, tile_expert, + total_tiles_dev, pre_act, post_act, + w_rows=E * 2 * intermediate_dim, H_dim=H, I_dim=intermediate_dim, stride_x_T=x.stride(0), @@ -227,19 +314,23 @@ def forward(ctx, x, gate_up_proj, down_proj, top_k_index, top_k_weights): stride_pre_N=pre_act.stride(1), stride_post_TK=post_act.stride(0), stride_post_N=post_act.stride(1), - BLOCK_M=BLOCK_M_TOKEN, + BLOCK_M=block_m_token, + USE_TMA=use_tma_w1, + STORE_PREACT=needs_grad, ) Y = torch.empty(TK, H, dtype=x.dtype, device=x.device) if num_m_tiles > 0: - _fused_down_proj_kernel[lambda meta: (num_m_tiles, triton.cdiv(H, meta["BLOCK_N"]))]( + _fused_down_proj_kernel[lambda meta: (num_m_tiles * triton.cdiv(H, meta["BLOCK_N"]),)]( post_act, down_proj, expert_start_idx, tile_row_start, tile_expert, + total_tiles_dev, Y, + w_rows=E * H, H_dim=H, I_dim=intermediate_dim, stride_post_TK=post_act.stride(0), @@ -249,25 +340,28 @@ def forward(ctx, x, gate_up_proj, down_proj, top_k_index, top_k_weights): stride_w_I=down_proj.stride(2), stride_Y_TK=Y.stride(0), stride_Y_H=Y.stride(1), - BLOCK_M=BLOCK_M_TOKEN, + BLOCK_M=block_m_token, + USE_TMA=use_tma_w2, ) topk_weights_flat = top_k_weights.flatten().contiguous() out = _token_aggregation(Y, topk_weights_flat, s_reverse_scatter_idx, T, K, H) - ctx.save_for_backward( - x, - gate_up_proj, - down_proj, - pre_act, - topk_weights_flat, - expert_start_idx, - x_gather_idx, - s_scatter_idx, - s_reverse_scatter_idx, - tile_row_start, - tile_expert, - ) + if needs_grad: + ctx.save_for_backward( + x, + gate_up_proj, + down_proj, + pre_act, + topk_weights_flat, + expert_start_idx, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + tile_row_start, + tile_expert, + total_tiles_dev, + ) ctx.T = T ctx.K = K ctx.E = E @@ -275,6 +369,7 @@ def forward(ctx, x, gate_up_proj, down_proj, top_k_index, top_k_weights): ctx.intermediate_dim = intermediate_dim ctx.TK = TK ctx.num_m_tiles = num_m_tiles + ctx.block_m_token = block_m_token ctx.mark_non_differentiable(top_k_index) ctx.set_materialize_grads(False) @@ -298,6 +393,7 @@ def backward(ctx, dO): s_reverse_scatter_idx, tile_row_start, tile_expert, + total_tiles_dev, ) = ctx.saved_tensors T = ctx.T @@ -307,14 +403,69 @@ def backward(ctx, dO): intermediate_dim = ctx.intermediate_dim TK = ctx.TK num_m_tiles = ctx.num_m_tiles + block_m_token = ctx.block_m_token + use_tma_w1, use_tma_w2 = _tma_eligibility(dO, H, intermediate_dim, E) + if use_tma_w1 or use_tma_w2: + _ensure_tma_allocator() + + mem_eff = _MEMORY_EFFICIENT + + # Tokens-per-expert bucket for the dW autotune keys: the best dW tile is + # regime-dependent (output-write-bound at small TK/E, K-loop-bound at + # large TK/E), but H_dim/I_dim alone can't see the difference — without + # this key the config tuned at the first-seen T is reused for every T. + # Clamped: below 16 / beyond 4096 tokens-per-expert the optimum stops + # moving, so extreme sizes share the edge buckets instead of retuning. + tpe_bucket = max(16, min(4096, triton.next_power_of_2(max(1, TK // max(1, E))))) + + # ---- dW2 = (s_k * y1)^T @ dO_gathered (memory-efficient order) ------ + # In memory-efficient mode s_k*y1 is recomputed from pre_act inside dW2, so + # dW2 MUST run before the bwd-down-proj kernel overwrites pre_act in place. + # empty (not zeros): the kernel writes every element, storing 0 for empty experts. + ddown_proj = torch.empty_like(down_proj) + if mem_eff: + # No weighted_act buffer, and d_pre_act aliases pre_act (in-place): + # each (row, n) element is consumed and produced by the same CTA, so + # the alias is race-free and saves a (TK, 2I) buffer. Costs support + # for a second backward over the same graph (version check raises). + weighted_act = pre_act[:0] # dummy, never read + d_pre_act = pre_act + _moe_bwd_dW2_kernel[ + lambda meta: (E * triton.cdiv(intermediate_dim, meta["BLOCK_M"]) * triton.cdiv(H, meta["BLOCK_N"]),) + ]( + weighted_act, + pre_act, + s_scatter_idx, + topk_weights_flat, + dO, + x_gather_idx, + expert_start_idx, + ddown_proj, + H_dim=H, + I_dim=intermediate_dim, + stride_wact_TK=0, + stride_wact_I=1, + stride_pre_TK=pre_act.stride(0), + stride_pre_N=pre_act.stride(1), + stride_dout_T=dO.stride(0), + stride_dout_H=dO.stride(1), + stride_dW2_E=ddown_proj.stride(0), + stride_dW2_H=ddown_proj.stride(1), + stride_dW2_I=ddown_proj.stride(2), + RECOMPUTE_WACT=True, + TPE_BUCKET=tpe_bucket, + ) + else: + weighted_act = torch.empty(TK, intermediate_dim, dtype=dO.dtype, device=dO.device) + d_pre_act = torch.empty(TK, 2 * intermediate_dim, dtype=dO.dtype, device=dO.device) - # dA' = dO @ W2^T, SwiGLU backward, write d_pre_act and dS - d_pre_act = torch.empty(TK, 2 * intermediate_dim, dtype=dO.dtype, device=dO.device) - weighted_act = torch.empty(TK, intermediate_dim, dtype=dO.dtype, device=dO.device) - dS = torch.zeros(TK, dtype=dO.dtype, device=dO.device) # zeros: atomic_add in kernel accumulates across N-tiles + # ---- dA' = dO @ W2^T, SwiGLU backward → d_pre_act, dS --------------- + # fp32 dS: atomic_add accumulates ceil(I/BLOCK_N) partials per element; + # low-precision atomics would round every partial. + dS = torch.zeros(TK, dtype=torch.float32, device=dO.device) if num_m_tiles > 0: - _moe_bwd_down_proj_kernel[lambda meta: (num_m_tiles, triton.cdiv(intermediate_dim, meta["BLOCK_N"]))]( + _moe_bwd_down_proj_kernel[lambda meta: (num_m_tiles * triton.cdiv(intermediate_dim, meta["BLOCK_N"]),)]( dO, x_gather_idx, s_scatter_idx, @@ -324,9 +475,11 @@ def backward(ctx, dO): expert_start_idx, tile_row_start, tile_expert, + total_tiles_dev, d_pre_act, weighted_act, dS, + w_rows=E * H, H_dim=H, I_dim=intermediate_dim, stride_dO_T=dO.stride(0), @@ -338,46 +491,58 @@ def backward(ctx, dO): stride_pre_N=pre_act.stride(1), stride_d_pre_TK=d_pre_act.stride(0), stride_d_pre_N=d_pre_act.stride(1), - stride_wact_TK=weighted_act.stride(0), - stride_wact_I=weighted_act.stride(1), - BLOCK_M=BLOCK_M_TOKEN, + stride_wact_TK=weighted_act.stride(0) if not mem_eff else 0, + stride_wact_I=weighted_act.stride(1) if not mem_eff else 1, + BLOCK_M=block_m_token, + USE_TMA=use_tma_w2, + WRITE_WACT=not mem_eff, ) - # dW2 = (s_k * y1)^T @ dO_gathered - ddown_proj = torch.zeros_like(down_proj) - _moe_bwd_dW2_kernel[ - lambda meta: ( - E * triton.cdiv(intermediate_dim, meta["BLOCK_M"]), - triton.cdiv(H, meta["BLOCK_N"]), + if mem_eff: + # pre_act now holds d_pre_act; bump its autograd version so a second + # backward through the same graph errors out instead of silently + # producing garbage. + torch.autograd.graph.increment_version(pre_act) + else: + _moe_bwd_dW2_kernel[ + lambda meta: (E * triton.cdiv(intermediate_dim, meta["BLOCK_M"]) * triton.cdiv(H, meta["BLOCK_N"]),) + ]( + weighted_act, + pre_act, + s_scatter_idx, + topk_weights_flat, + dO, + x_gather_idx, + expert_start_idx, + ddown_proj, + H_dim=H, + I_dim=intermediate_dim, + stride_wact_TK=weighted_act.stride(0), + stride_wact_I=weighted_act.stride(1), + stride_pre_TK=pre_act.stride(0), + stride_pre_N=pre_act.stride(1), + stride_dout_T=dO.stride(0), + stride_dout_H=dO.stride(1), + stride_dW2_E=ddown_proj.stride(0), + stride_dW2_H=ddown_proj.stride(1), + stride_dW2_I=ddown_proj.stride(2), + RECOMPUTE_WACT=False, + TPE_BUCKET=tpe_bucket, ) - ]( - weighted_act, - dO, - x_gather_idx, - expert_start_idx, - ddown_proj, - H_dim=H, - I_dim=intermediate_dim, - stride_wact_TK=weighted_act.stride(0), - stride_wact_I=weighted_act.stride(1), - stride_dout_T=dO.stride(0), - stride_dout_H=dO.stride(1), - stride_dW2_E=ddown_proj.stride(0), - stride_dW2_H=ddown_proj.stride(1), - stride_dW2_I=ddown_proj.stride(2), - ) # dx_expanded = d_pre_act @ W1^T dx_expanded = torch.empty(TK, H, dtype=dO.dtype, device=dO.device) if num_m_tiles > 0: - _moe_bwd_dX_expanded_kernel[lambda meta: (num_m_tiles, triton.cdiv(H, meta["BLOCK_N"]))]( + _moe_bwd_dX_expanded_kernel[lambda meta: (num_m_tiles * triton.cdiv(H, meta["BLOCK_N"]),)]( d_pre_act, gate_up_proj, expert_start_idx, tile_row_start, tile_expert, + total_tiles_dev, dx_expanded, + w_rows=E * 2 * intermediate_dim, H_dim=H, I_dim=intermediate_dim, stride_d_pre_TK=d_pre_act.stride(0), @@ -387,13 +552,15 @@ def backward(ctx, dO): stride_w_K=gate_up_proj.stride(2), stride_dxe_TK=dx_expanded.stride(0), stride_dxe_H=dx_expanded.stride(1), - BLOCK_M=BLOCK_M_TOKEN, + BLOCK_M=block_m_token, + USE_TMA=use_tma_w1, ) # dx = unweighted gather-sum of dx_expanded - dx = torch.zeros(T, H, dtype=dO.dtype, device=dO.device) + # empty (not zeros): the gather-sum kernel stores every (t, h) element. + dx = torch.empty(T, H, dtype=dO.dtype, device=dO.device) if TK > 0: - _token_gather_weighted_sum_kernel[(T,)]( + _token_gather_weighted_sum_kernel[lambda meta: (T, triton.cdiv(H, meta["BLOCK_H"]))]( dx_expanded, dS, # dummy w_ptr — never loaded when w_is_None=True s_reverse_scatter_idx, @@ -408,12 +575,10 @@ def backward(ctx, dO): ) # dW1 = X_gathered^T @ d_pre_act - dgate_up_proj = torch.zeros_like(gate_up_proj) + # empty (not zeros): the kernel writes every element, storing 0 for empty experts. + dgate_up_proj = torch.empty_like(gate_up_proj) _moe_bwd_dW1_kernel[ - lambda meta: ( - E * triton.cdiv(H, meta["BLOCK_M"]), - triton.cdiv(2 * intermediate_dim, meta["BLOCK_N"]), - ) + lambda meta: (E * triton.cdiv(H, meta["BLOCK_M"]) * triton.cdiv(2 * intermediate_dim, meta["BLOCK_N"]),) ]( x, d_pre_act, @@ -429,6 +594,7 @@ def backward(ctx, dO): stride_dW1_E=dgate_up_proj.stride(0), stride_dW1_N=dgate_up_proj.stride(1), stride_dW1_H=dgate_up_proj.stride(2), + TPE_BUCKET=tpe_bucket, ) - return dx, dgate_up_proj, ddown_proj, None, dS.view(T, K) + return dx, dgate_up_proj, ddown_proj, None, dS.to(topk_weights_flat.dtype).view(T, K) diff --git a/src/liger_kernel/ops/fused_moe_kernels.py b/src/liger_kernel/ops/fused_moe_kernels.py index 11f62d412..27f2fa028 100644 --- a/src/liger_kernel/ops/fused_moe_kernels.py +++ b/src/liger_kernel/ops/fused_moe_kernels.py @@ -5,20 +5,56 @@ # Copyright 2025 Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao # # Grouped GEMM kernels and backward kernels are new Triton implementations -# inspired by the SonicMoE paper (arXiv:2512.14080), ported to portable Triton -# (no Hopper-specific WGMMA/TMA) for general GPU support. +# inspired by the SonicMoE paper (arXiv:2512.14080). Expert-weight loads use TMA +# tensor descriptors on Hopper+ when eligible (USE_TMA), with a portable +# pointer-load path for other GPUs and non-16B-aligned shapes. import os import triton import triton.language as tl +from liger_kernel.utils import infer_device_arch + # LIGER_FUSED_MOE_AUTOTUNE=0 pins each kernel to one config, skipping Triton's # `do_bench` loop whose per-config working sets can OOM (see issue #1246). Must # be set before importing liger_kernel. Temporary escape hatch until triton's # autotuner handles such errors itself. _AUTOTUNE_DISABLED = os.environ.get("LIGER_FUSED_MOE_AUTOTUNE", "1").lower() in ("0", "false", "no") +# LIGER_FUSED_MOE_MEMORY_EFFICIENT=1 (must be set before import): backward writes +# SwiGLU gradients in place over the saved pre-activations and skips the (TK, I) +# weighted_act buffer. See liger_kernel.ops.fused_moe for details. Read at import +# because it decides the autotuner's restore_value list: when the in-place alias is +# active, tuning runs must restore pre_act between configs; when it is not, the +# restore's copy_() would needlessly version-bump a saved tensor and break +# retain_graph double-backward for everyone. +_MEMORY_EFFICIENT = os.environ.get("LIGER_FUSED_MOE_MEMORY_EFFICIENT", "0").lower() in ("1", "true", "yes") + + +def _is_blackwell_datacenter() -> bool: + """True on sm100/sm103 (B200/B300 class) — the parts with tcgen05 MMA + TMEM. + + Used to extend (never replace) the autotune config spaces: TMEM accumulators + make wide-N tiles cheap there, and tcgen05 MMAs are issued by a single warp + so num_warps=4 competes with 8. Deliberately excludes consumer Blackwell + (sm120, "blackwell_consumer"): it has neither TMEM nor 228 KB smem, so the + wide-tile configs would only waste tuning time as OutOfResources skips. + + Only called from the autotuners' early_config_prune hooks (i.e. at first + kernel launch), NOT at import: infer_device_arch() initializes CUDA, and + importing liger_kernel must stay side-effect-free for fork-based workers. + """ + return infer_device_arch() in ("blackwell", "blackwell_ultra") + + +def _blackwell_config(kwargs, num_warps, num_stages): + """A triton.Config that the prune hooks drop on non-sm100/sm103 devices.""" + cfg = triton.Config(kwargs, num_warps=num_warps, num_stages=num_stages) + cfg.liger_blackwell_only = True + return cfg + + # --------------------------------------------------------------------------- # Routing metadata overview # @@ -314,7 +350,7 @@ def _moe_router_scatter_kernel( def _get_gemm_autotune_configs(): if _AUTOTUNE_DISABLED: - return [triton.Config({"BLOCK_N": 128, "BLOCK_K": 64}, num_warps=8, num_stages=2)] + return [triton.Config({"BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=8, num_stages=2)] configs = [] for bn in [64, 128]: for bk in [32, 64]: @@ -322,25 +358,130 @@ def _get_gemm_autotune_configs(): for ns in [2, 3, 4, 5]: configs.append( triton.Config( - {"BLOCK_N": bn, "BLOCK_K": bk}, + {"BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": 8}, num_warps=nw, num_stages=ns, ) ) + # Wider tiles for compute-bound large-T shapes (H100 has 228 KB smem). + for bn, bk, nw, ns in [ + (256, 32, 8, 3), + (256, 64, 8, 2), + (256, 64, 8, 3), + (256, 64, 8, 4), + (128, 128, 8, 2), + (128, 128, 8, 3), + (128, 128, 4, 3), + (256, 128, 8, 2), + ]: + configs.append(triton.Config({"BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": 8}, num_warps=nw, num_stages=ns)) + # sm100/sm103 extras (B300-measured), pruned away at launch on other archs: + # TMEM holds the fp32 accumulator, so BN=256 with deep pipelines wins where + # Hopper would spill registers — a gathered-GEMM probe hit 1103 TFLOPS at + # BN=256/BK=64/ns=3 vs 808 at the Hopper-favored BN=128/BK=64/ns=4. + # GROUP_M=16 probes exploit the 132 MB L2 (H100: 50 MB). Configs that exceed + # smem fail compile and are skipped by the autotuner (OutOfResources → inf). + for bn, bk, nw, ns in [ + (256, 64, 4, 3), + (256, 64, 4, 4), + (256, 64, 8, 5), + (256, 128, 8, 3), + (256, 128, 4, 2), + (128, 128, 8, 4), + (128, 128, 4, 4), + (256, 32, 4, 4), + ]: + configs.append(_blackwell_config({"BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": 8}, num_warps=nw, num_stages=ns)) + for bn, bk, nw, ns in [ + (256, 64, 8, 3), + (128, 64, 8, 4), + ]: + configs.append(_blackwell_config({"BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": 16}, num_warps=nw, num_stages=ns)) return configs def _get_dW_autotune_configs(): """Configs for backward weight-grad kernels (dW1, dW2): include BLOCK_M sweep.""" if _AUTOTUNE_DISABLED: - return [triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=8, num_stages=2)] - return [ - triton.Config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk}, num_warps=nw, num_stages=2) + return [triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_M": 4}, num_warps=8, num_stages=2)] + configs = [ + triton.Config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": 4}, num_warps=nw, num_stages=ns) for bm in [64, 128] - for bn in [64, 128] - for bk in [16, 32] + for bn in [64, 128, 256] + for bk in [32, 64] for nw in [4, 8] + for ns in [2, 3] ] + # sm100/sm103 extras (B300-measured, +22-26% backward at TK/E >= 512 together + # with the TPE_BUCKET key), pruned away at launch on other archs: deep-K + # (BK=128) reads token rows in fewer longer bursts (neutral on H100 per E13), + # ns=4 pipelines are free with TMEM accumulators, GROUP_M=8 halves re-reads + # of the per-expert row window in the larger L2. + for bm, bn, bk, gm, nw, ns in [ + (128, 256, 128, 4, 8, 2), + (128, 256, 128, 4, 8, 3), + (128, 256, 64, 4, 8, 4), + (128, 256, 64, 4, 4, 3), + (128, 128, 128, 4, 8, 3), + (128, 256, 64, 8, 8, 3), + (128, 128, 64, 8, 8, 3), + (64, 256, 128, 4, 8, 3), + ]: + configs.append( + _blackwell_config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": gm}, num_warps=nw, num_stages=ns) + ) + return configs + + +def _make_tile_prune(n_extent_of, m_extent_of=None): + """Shape/arch-aware early config pruning. + + Drops (a) tiles wider than the (padded) problem — cuts first-run + compile+tune cost dramatically for small shapes (e.g. unit tests) — and + (b) the Blackwell-datacenter-only configs when not on sm100/sm103, keeping + other architectures' search spaces exactly as tuned on H100. Falls back to + the full list if pruning would empty it. + """ + + def prune(configs, nargs, **kwargs): + args = {**nargs, **kwargs} + if not _is_blackwell_datacenter(): + configs = [c for c in configs if not getattr(c, "liger_blackwell_only", False)] + n_extent = triton.next_power_of_2(max(64, n_extent_of(args))) + m_extent = triton.next_power_of_2(max(64, m_extent_of(args))) if m_extent_of is not None else None + pruned = [ + c + for c in configs + if c.kwargs["BLOCK_N"] <= n_extent and (m_extent is None or c.kwargs.get("BLOCK_M", 0) <= m_extent) + ] + return pruned or configs + + return prune + + +# N-extent per kernel family (the dimension BLOCK_N tiles over). +_prune_gemm_n_is_I = _make_tile_prune(lambda a: a["I_dim"]) +_prune_gemm_n_is_H = _make_tile_prune(lambda a: a["H_dim"]) +_prune_dW2 = _make_tile_prune(lambda a: a["H_dim"], lambda a: a["I_dim"]) +_prune_dW1 = _make_tile_prune(lambda a: 2 * a["I_dim"], lambda a: a["H_dim"]) + + +@triton.jit +def _grouped_pid_swizzle(pid, num_pid_m, NUM_PID_N: tl.constexpr, GROUP_M: tl.constexpr): + """L2-friendly CTA remap (triton matmul-tutorial style). + + Flat 1D pid → (pid_m, pid_n) such that GROUP_M consecutive m-tiles are + visited for every n before advancing: CTAs resident together then share + both the x rows of those m-tiles and the (expert) weight n-tiles in L2. + m-tiles are expert-major (sorted), so a group usually stays inside one expert. + """ + num_pid_in_group = GROUP_M * NUM_PID_N + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_M + group_size_m = tl.minimum(num_pid_m - first_pid_m, GROUP_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + return pid_m, pid_n # --------------------------------------------------------------------------- @@ -351,7 +492,8 @@ def _get_dW_autotune_configs(): @triton.autotune( configs=_get_gemm_autotune_configs(), - key=["H_dim", "I_dim"], + key=["H_dim", "I_dim", "BLOCK_M", "USE_TMA", "STORE_PREACT"], + prune_configs_by={"early_config_prune": _prune_gemm_n_is_I}, ) @triton.jit def _fused_up_proj_swiglu_kernel( @@ -361,8 +503,10 @@ def _fused_up_proj_swiglu_kernel( expert_start_ptr, # (E+1,) int32 tile_row_start_ptr, # (num_m_tiles,) int32 — row_start per M-tile tile_expert_ptr, # (num_m_tiles,) int32 — expert index per M-tile + total_tiles_ptr, # (1,) int32 — actual number of m-tiles (device scalar) pre_act_ptr, # (TK, 2*I) pre-SwiGLU activations [saved for backward] post_act_ptr, # (TK, I) post-SwiGLU activations + w_rows, # E * 2*I — rows of the flattened (E*2I, H) weight view (TMA) H_dim: tl.constexpr, I_dim: tl.constexpr, stride_x_T, @@ -377,11 +521,22 @@ def _fused_up_proj_swiglu_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, + USE_TMA: tl.constexpr, + STORE_PREACT: tl.constexpr, ): - """Grid: (num_m_tiles, ceil(I/BLOCK_N)). - pid_m selects M-tile via tile_row_start/tile_expert; pid_n selects N-tile.""" - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) + """Grid: 1D (num_m_tiles_max * ceil(I/BLOCK_N),), L2-swizzled into (pid_m, pid_n). + pid_m selects M-tile via tile_row_start/tile_expert; pid_n selects N-tile. + Grid is an upper bound; CTAs past the actual tile count exit early. + USE_TMA: weight tiles load through a TMA descriptor over the flattened + (E*2I, H) view — async bulk copies that bypass L1 and free LSU issue slots. + STORE_PREACT=False (inference): skip writing the (TK, 2I) pre-activation.""" + pid = tl.program_id(0) + NUM_PID_N: tl.constexpr = (I_dim + BLOCK_N - 1) // BLOCK_N + num_pid_m = tl.load(total_tiles_ptr) + if pid >= num_pid_m * NUM_PID_N: + return + pid_m, pid_n = _grouped_pid_swizzle(pid, num_pid_m, NUM_PID_N, GROUP_M) row_start = tl.load(tile_row_start_ptr + pid_m) # int64 prevents expert_idx * stride_w_E overflow at large E*I*H (see #1246). @@ -402,6 +557,16 @@ def _fused_up_proj_swiglu_kernel( n_idx = n_start + n_offs n_mask = n_idx < I_dim + if USE_TMA: + w_desc = tl.make_tensor_descriptor( + gate_up_proj_ptr, + shape=[w_rows, H_dim], + strides=[stride_w_N, 1], + block_shape=[BLOCK_N, BLOCK_K], + ) + # Row of this expert's gate block in the flattened view. OOB rows at the + # tail expert zero-fill; garbage lanes are masked at the store. + w_row0 = (expert_idx * (2 * I_dim)).to(tl.int32) + n_start # int64 prevents token_idx * stride_T overflow at large T*H (see #1246). token_idx = tl.load(x_gather_idx_ptr + row_offs, mask=row_mask, other=0).to(tl.int64) for k in tl.range(0, H_dim, BLOCK_K): @@ -417,32 +582,39 @@ def _fused_up_proj_swiglu_kernel( eviction_policy="evict_first", # token rows not reused; free L2 for weights ) - w_mask = n_mask[:, None] & k_mask[None, :] - w_gate_ptrs = ( - gate_up_proj_ptr + expert_idx * stride_w_E + n_idx[:, None] * stride_w_N + k_idx[None, :] * stride_w_K - ) - w_gate = tl.load( - w_gate_ptrs, - mask=w_mask, - other=0.0, - ) - acc_gate = tl.dot(x_tile, tl.trans(w_gate), acc=acc_gate) - - w_up_ptrs = w_gate_ptrs + I_dim * stride_w_N - w_up = tl.load( - w_up_ptrs, - mask=w_mask, - other=0.0, - ) - - acc_up = tl.dot(x_tile, tl.trans(w_up), acc=acc_up) + if USE_TMA: + w_gate = w_desc.load([w_row0, k]) + acc_gate = tl.dot(x_tile, tl.trans(w_gate), acc=acc_gate) + w_up = w_desc.load([w_row0 + I_dim, k]) + acc_up = tl.dot(x_tile, tl.trans(w_up), acc=acc_up) + else: + w_mask = n_mask[:, None] & k_mask[None, :] + w_gate_ptrs = ( + gate_up_proj_ptr + expert_idx * stride_w_E + n_idx[:, None] * stride_w_N + k_idx[None, :] * stride_w_K + ) + w_gate = tl.load( + w_gate_ptrs, + mask=w_mask, + other=0.0, + ) + acc_gate = tl.dot(x_tile, tl.trans(w_gate), acc=acc_gate) + + w_up_ptrs = w_gate_ptrs + I_dim * stride_w_N + w_up = tl.load( + w_up_ptrs, + mask=w_mask, + other=0.0, + ) + + acc_up = tl.dot(x_tile, tl.trans(w_up), acc=acc_up) out_mask = row_mask[:, None] & n_mask[None, :] - pre_gate_ptrs = pre_act_ptr + row_offs[:, None] * stride_pre_TK + n_idx[None, :] * stride_pre_N - pre_up_ptrs = pre_gate_ptrs + I_dim * stride_pre_N - tl.store(pre_gate_ptrs, acc_gate.to(pre_act_ptr.dtype.element_ty), mask=out_mask) - tl.store(pre_up_ptrs, acc_up.to(pre_act_ptr.dtype.element_ty), mask=out_mask) + if STORE_PREACT: + pre_gate_ptrs = pre_act_ptr + row_offs[:, None] * stride_pre_TK + n_idx[None, :] * stride_pre_N + pre_up_ptrs = pre_gate_ptrs + I_dim * stride_pre_N + tl.store(pre_gate_ptrs, acc_gate.to(pre_act_ptr.dtype.element_ty), mask=out_mask) + tl.store(pre_up_ptrs, acc_up.to(pre_act_ptr.dtype.element_ty), mask=out_mask) sig_gate = tl.sigmoid(acc_gate) silu_gate = acc_gate * sig_gate @@ -460,7 +632,8 @@ def _fused_up_proj_swiglu_kernel( @triton.autotune( configs=_get_gemm_autotune_configs(), - key=["H_dim", "I_dim"], + key=["H_dim", "I_dim", "BLOCK_M", "USE_TMA"], + prune_configs_by={"early_config_prune": _prune_gemm_n_is_H}, ) @triton.jit def _fused_down_proj_kernel( @@ -469,7 +642,9 @@ def _fused_down_proj_kernel( expert_start_ptr, # (E+1,) int32 tile_row_start_ptr, # (num_m_tiles,) int32 tile_expert_ptr, # (num_m_tiles,) int32 + total_tiles_ptr, # (1,) int32 — actual number of m-tiles (device scalar) Y_ptr, # (TK, H) + w_rows, # E * H — rows of the flattened (E*H, I) weight view (TMA) H_dim: tl.constexpr, I_dim: tl.constexpr, stride_post_TK, @@ -482,11 +657,17 @@ def _fused_down_proj_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, + USE_TMA: tl.constexpr, ): - """Grid: (num_m_tiles, ceil(H/BLOCK_N)). + """Grid: 1D (num_m_tiles_max * ceil(H/BLOCK_N),), L2-swizzled into (pid_m, pid_n). Each CTA: one (BLOCK_M, BLOCK_N) tile of Y = post_act @ down_proj[e]^T.""" - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) + pid = tl.program_id(0) + NUM_PID_N: tl.constexpr = (H_dim + BLOCK_N - 1) // BLOCK_N + num_pid_m = tl.load(total_tiles_ptr) + if pid >= num_pid_m * NUM_PID_N: + return + pid_m, pid_n = _grouped_pid_swizzle(pid, num_pid_m, NUM_PID_N, GROUP_M) row_start = tl.load(tile_row_start_ptr + pid_m) # int64 prevents expert_idx * stride_w_E overflow at large E*I*H (see #1246). @@ -506,6 +687,15 @@ def _fused_down_proj_kernel( acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + if USE_TMA: + w_desc = tl.make_tensor_descriptor( + down_proj_ptr, + shape=[w_rows, I_dim], + strides=[stride_w_H, 1], + block_shape=[BLOCK_N, BLOCK_K], + ) + w_row0 = (expert_idx * H_dim).to(tl.int32) + n_start + for k in tl.range(0, I_dim, BLOCK_K): k_idx = k + k_offs k_mask = k_idx < I_dim @@ -514,12 +704,15 @@ def _fused_down_proj_kernel( # Keep bf16 for dot operands → tensor cores. acc stays fp32. a_tile = tl.load(a_ptrs, mask=row_mask[:, None] & k_mask[None, :], other=0.0) - w_ptrs = down_proj_ptr + expert_idx * stride_w_E + n_idx[:, None] * stride_w_H + k_idx[None, :] * stride_w_I - w_tile = tl.load( - w_ptrs, - mask=n_mask[:, None] & k_mask[None, :], - other=0.0, - ) + if USE_TMA: + w_tile = w_desc.load([w_row0, k]) + else: + w_ptrs = down_proj_ptr + expert_idx * stride_w_E + n_idx[:, None] * stride_w_H + k_idx[None, :] * stride_w_I + w_tile = tl.load( + w_ptrs, + mask=n_mask[:, None] & k_mask[None, :], + other=0.0, + ) acc = tl.dot(a_tile, tl.trans(w_tile), acc=acc) @@ -537,7 +730,9 @@ def _get_token_gather_autotune_configs(): if _AUTOTUNE_DISABLED: return [triton.Config({"BLOCK_H": 128, "BLOCK_K": 4}, num_warps=4, num_stages=4)] configs = [] - for bh in [64, 128, 256, 512]: + # BLOCK_H=1024 rows help on Blackwell (HBM3e rewards longer contiguous + # vectors; B300 sustains 6.4 TB/s vs H100's 3.3) and are harmless elsewhere. + for bh in [64, 128, 256, 512, 1024]: for bk in [1, 2, 4, 8, 16]: for nw in [4, 8]: if bk * bh <= 32768: @@ -565,35 +760,36 @@ def _token_gather_weighted_sum_kernel( BLOCK_K: tl.constexpr, w_is_None: tl.constexpr, # True → unweighted gather-sum (used for dx backward) ): - """One CTA per token. Gathers K expert outputs, reduces with routing weights - (forward) or without weights (backward dx via _token_broadcast_backward).""" + """Grid: (T, ceil(H/BLOCK_H)) — 2D so small-T launches still fill the GPU. + Each CTA gathers K expert rows for one (token, H-tile) and reduces with routing + weights (forward) or without weights (backward dx).""" # int64 prevents t * stride_out_T overflow at large T*H (see #1246). t = tl.program_id(0).to(tl.int64) + h_tile = tl.program_id(1) - for h_tile in tl.static_range(triton.cdiv(H_dim, BLOCK_H)): - h_idx = (h_tile * BLOCK_H + tl.arange(0, BLOCK_H)).to(tl.uint32) - h_mask = h_idx < H_dim - acc = tl.zeros([BLOCK_H], dtype=tl.float32) + h_idx = (h_tile * BLOCK_H + tl.arange(0, BLOCK_H)).to(tl.uint32) + h_mask = h_idx < H_dim + acc = tl.zeros([BLOCK_H], dtype=tl.float32) - for k_tile in tl.range(triton.cdiv(K_dim, BLOCK_K)): - k_offs = (k_tile * BLOCK_K + tl.arange(0, BLOCK_K)).to(tl.uint32) - k_mask = k_offs < K_dim + for k_tile in tl.range(triton.cdiv(K_dim, BLOCK_K)): + k_offs = (k_tile * BLOCK_K + tl.arange(0, BLOCK_K)).to(tl.uint32) + k_mask = k_offs < K_dim - flat_idx = t * K_dim + k_offs - # int64 prevents perm_idx * stride overflow when TK is large (see #1246). - perm_idx = tl.load(s_rev_ptr + flat_idx, mask=k_mask, other=0).to(tl.int64) + flat_idx = t * K_dim + k_offs + # int64 prevents perm_idx * stride overflow when TK is large (see #1246). + perm_idx = tl.load(s_rev_ptr + flat_idx, mask=k_mask, other=0).to(tl.int64) - y_ptrs = Y_ptr + perm_idx[:, None] * stride_Y_TK + h_idx[None, :] * stride_Y_H - y_vals = tl.load(y_ptrs, mask=k_mask[:, None] & h_mask[None, :], other=0.0).to(tl.float32) + y_ptrs = Y_ptr + perm_idx[:, None] * stride_Y_TK + h_idx[None, :] * stride_Y_H + y_vals = tl.load(y_ptrs, mask=k_mask[:, None] & h_mask[None, :], other=0.0).to(tl.float32) - if w_is_None: - acc += tl.sum(y_vals, axis=0) - else: - w_vals = tl.load(w_ptr + flat_idx, mask=k_mask, other=0.0).to(tl.float32) - acc += tl.sum(y_vals * w_vals[:, None], axis=0) + if w_is_None: + acc += tl.sum(y_vals, axis=0) + else: + w_vals = tl.load(w_ptr + flat_idx, mask=k_mask, other=0.0).to(tl.float32) + acc += tl.sum(y_vals * w_vals[:, None], axis=0) - out_ptrs = out_ptr + t * stride_out_T + h_idx * stride_out_H - tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=h_mask) + out_ptrs = out_ptr + t * stride_out_T + h_idx * stride_out_H + tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=h_mask) # --------------------------------------------------------------------------- @@ -604,8 +800,14 @@ def _token_gather_weighted_sum_kernel( @triton.autotune( configs=_get_gemm_autotune_configs(), - key=["H_dim", "I_dim"], + key=["H_dim", "I_dim", "BLOCK_M", "USE_TMA"], reset_to_zero=["dS_ptr"], # autotune runs multiple configs; atomic_add accumulates, so reset between runs + # Memory-efficient mode aliases d_pre_act onto pre_act (in-place), so tuning + # runs must restore pre_act between configs. In default mode the restore is + # skipped: its copy_() would version-bump a saved tensor on every tuning run + # and spuriously break retain_graph double-backward. + restore_value=["pre_act_ptr"] if _MEMORY_EFFICIENT else [], + prune_configs_by={"early_config_prune": _prune_gemm_n_is_I}, ) @triton.jit def _moe_bwd_down_proj_kernel( @@ -618,9 +820,11 @@ def _moe_bwd_down_proj_kernel( expert_start_ptr, # (E+1,) int32 tile_row_start_ptr, # (num_m_tiles,) int32 tile_expert_ptr, # (num_m_tiles,) int32 - d_pre_act_ptr, # (TK, 2I) — output: ∂L/∂z = [dgate, dup] - weighted_act_ptr, # (TK, I) — output: s_k * y1 (for dW2 kernel) - dS_ptr, # (TK,) — output: ∂L/∂s_k, indexed by flat (t,k) + total_tiles_ptr, # (1,) int32 — actual number of m-tiles (device scalar) + d_pre_act_ptr, # (TK, 2I) — output: ∂L/∂z = [dgate, dup]; MAY ALIAS pre_act_ptr + weighted_act_ptr, # (TK, I) — output: s_k * y1 (for dW2 kernel); unused if not WRITE_WACT + dS_ptr, # (TK,) fp32 — output: ∂L/∂s_k, indexed by flat (t,k) + w_rows, # E * H — rows of the flattened (E*H, I) weight view (TMA) H_dim: tl.constexpr, I_dim: tl.constexpr, stride_dO_T, @@ -637,12 +841,21 @@ def _moe_bwd_down_proj_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, + USE_TMA: tl.constexpr, + WRITE_WACT: tl.constexpr, # False in memory-efficient mode (dW2 recomputes s_k*y1) ): - """Grid: (num_m_tiles, ceil(I/BLOCK_N)). + """Grid: 1D (num_m_tiles_max * ceil(I/BLOCK_N),), L2-swizzled into (pid_m, pid_n). Accumulates dA' = dO @ W2^T (dO stays in registers), recomputes y1 from - pre_act, applies SwiGLU backward, writes d_pre_act, weighted_act, and dS.""" - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) + pre_act, applies SwiGLU backward, writes d_pre_act (in-place over pre_act — + each (row, n) element is read and written by exactly the same CTA), dS, and + optionally weighted_act.""" + pid = tl.program_id(0) + NUM_PID_N: tl.constexpr = (I_dim + BLOCK_N - 1) // BLOCK_N + num_pid_m = tl.load(total_tiles_ptr) + if pid >= num_pid_m * NUM_PID_N: + return + pid_m, pid_n = _grouped_pid_swizzle(pid, num_pid_m, NUM_PID_N, GROUP_M) row_start = tl.load(tile_row_start_ptr + pid_m) # int64 prevents expert_idx * stride_w_E overflow at large E*I*H (see #1246). @@ -668,6 +881,14 @@ def _moe_bwd_down_proj_kernel( weights = tl.load(topk_weights_ptr + flat_tk_idx, mask=row_mask, other=0.0).to(tl.float32) # K-loop: accumulate dA' = dO @ W2^T (unscaled; scale once after loop). + if USE_TMA: + w_desc = tl.make_tensor_descriptor( + down_proj_ptr, + shape=[w_rows, I_dim], + strides=[stride_w_H, 1], + block_shape=[BLOCK_K, BLOCK_N], + ) + w_row0 = (expert_idx * H_dim).to(tl.int32) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in tl.range(0, H_dim, BLOCK_K): k_idx = k + k_offs @@ -676,8 +897,11 @@ def _moe_bwd_down_proj_kernel( dO_ptrs = dO_ptr + token_idx[:, None] * stride_dO_T + k_idx[None, :] * stride_dO_H dO_tile = tl.load(dO_ptrs, mask=row_mask[:, None] & k_mask[None, :], other=0.0) - w_ptrs = down_proj_ptr + expert_idx * stride_w_E + k_idx[:, None] * stride_w_H + n_idx[None, :] * stride_w_I - w_tile = tl.load(w_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0) + if USE_TMA: + w_tile = w_desc.load([w_row0 + k, n_start]) + else: + w_ptrs = down_proj_ptr + expert_idx * stride_w_E + k_idx[:, None] * stride_w_H + n_idx[None, :] * stride_w_I + w_tile = tl.load(w_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0) acc = tl.dot(dO_tile, w_tile, acc=acc) # Epilogue: recompute y1 = silu(gate) * up from saved pre_act. @@ -690,14 +914,16 @@ def _moe_bwd_down_proj_kernel( silu_gate = gate * sig_gate y1 = silu_gate * up # (BLOCK_M, BLOCK_N) - # Write weighted_act = s_k * y1 for dW2. - wact_ptrs = weighted_act_ptr + row_offs[:, None] * stride_wact_TK + n_idx[None, :] * stride_wact_I - tl.store(wact_ptrs, (weights[:, None] * y1).to(weighted_act_ptr.dtype.element_ty), mask=out_mask) + if WRITE_WACT: + # Write weighted_act = s_k * y1 for dW2. + wact_ptrs = weighted_act_ptr + row_offs[:, None] * stride_wact_TK + n_idx[None, :] * stride_wact_I + tl.store(wact_ptrs, (weights[:, None] * y1).to(weighted_act_ptr.dtype.element_ty), mask=out_mask) # dS: ∂L/∂s_k = sum_I((dO @ W2^T) * y1) — accumulate across all N-tiles. # IMPORTANT: use atomic_add, not store — the grid has ceil(I/BLOCK_N) N-tiles per # M-tile, each contributing a partial sum over its I-chunk. tl.store would # overwrite previous tiles, leaving only the last chunk's contribution. + # dS buffer is fp32: bf16 atomics would round every partial contribution. dS_partial = tl.sum(acc * y1, axis=1) tl.atomic_add(dS_ptr + flat_tk_idx, dS_partial, mask=row_mask) @@ -721,12 +947,15 @@ def _moe_bwd_down_proj_kernel( @triton.autotune( configs=_get_dW_autotune_configs(), - key=["H_dim", "I_dim"], - reset_to_zero=["dW2_ptr"], + key=["H_dim", "I_dim", "RECOMPUTE_WACT", "TPE_BUCKET"], + prune_configs_by={"early_config_prune": _prune_dW2}, ) @triton.jit def _moe_bwd_dW2_kernel( - weighted_act_ptr, # (TK, I) — s_k * y1 from backward down-proj kernel + weighted_act_ptr, # (TK, I) — s_k * y1 from backward down-proj kernel (or dummy) + pre_act_ptr, # (TK, 2I) — saved [gate, up]; used when RECOMPUTE_WACT + s_scatter_idx_ptr, # (TK,) — sorted_pos → flat (t,k) index (RECOMPUTE_WACT only) + topk_weights_ptr, # (TK,) — s_k in flat (t,k) order (RECOMPUTE_WACT only) dout_ptr, # (T, H) — upstream gradient (gathered by x_gather_idx) x_gather_idx_ptr, # (TK,) — sorted_pos → original token index expert_start_ptr, # (E+1,) int32 @@ -735,6 +964,8 @@ def _moe_bwd_dW2_kernel( I_dim: tl.constexpr, stride_wact_TK, stride_wact_I: tl.constexpr, + stride_pre_TK, + stride_pre_N: tl.constexpr, stride_dout_T, stride_dout_H: tl.constexpr, stride_dW2_E, @@ -743,25 +974,33 @@ def _moe_bwd_dW2_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, + RECOMPUTE_WACT: tl.constexpr, # True → recompute s_k*y1 from pre_act (no wact buffer) + TPE_BUCKET: tl.constexpr = 0, # tokens-per-expert bucket; autotune key only (unused in body) ): - """dW2[e, h, i] += sum_t weighted_act[t, i] * dout[token(t), h] for tokens in e. - Grid: (E * ceil(I/BLOCK_M), ceil(H/BLOCK_N)). Early exit for empty experts.""" - pid0 = tl.program_id(0) - pid1 = tl.program_id(1) + """dW2[e, h, i] = sum_t (s_k*y1)[t, i] * dout[token(t), h] for tokens in e. + s_k*y1 either read from the materialized weighted_act buffer, or (memory- + efficient mode) recomputed on the fly from pre_act — in that case this kernel + MUST run before _moe_bwd_down_proj_kernel overwrites pre_act in place. + Grid: 1D (E * ceil(I/BLOCK_M) * ceil(H/BLOCK_N),), expert-major so all tiles of + one expert are resident together: its activation/dO rows stay in L2 instead of + being re-read from HBM per tile. Empty experts store zeros (no separate memset).""" + pid = tl.program_id(0) N_M_TILES: tl.constexpr = (I_dim + BLOCK_M - 1) // BLOCK_M + N_N_TILES: tl.constexpr = (H_dim + BLOCK_N - 1) // BLOCK_N + TILES_PER_E: tl.constexpr = N_M_TILES * N_N_TILES # int64 prevents expert_idx * stride_dW_E overflow at large E*I*H (see #1246). - expert_idx = (pid0 // N_M_TILES).to(tl.int64) - m_tile = pid0 % N_M_TILES + expert_idx = (pid // TILES_PER_E).to(tl.int64) + local = pid % TILES_PER_E + m_tile, n_tile = _grouped_pid_swizzle(local, N_M_TILES, N_N_TILES, GROUP_M) expert_start = tl.load(expert_start_ptr + expert_idx) expert_end = tl.load(expert_start_ptr + expert_idx + 1) M_e = expert_end - expert_start - if M_e == 0: - return m_start = m_tile * BLOCK_M - n_start = pid1 * BLOCK_N + n_start = n_tile * BLOCK_N m_offs = tl.arange(0, BLOCK_M) n_offs = tl.arange(0, BLOCK_N) @@ -779,8 +1018,21 @@ def _moe_bwd_dW2_kernel( k_mask = k_idx < M_e row_offs = (expert_start + k_idx).to(tl.int64) - wact_ptrs = weighted_act_ptr + row_offs[None, :] * stride_wact_TK + i_idx[:, None] * stride_wact_I - wact_tile = tl.load(wact_ptrs, mask=k_mask[None, :] & i_mask[:, None], other=0.0) + if RECOMPUTE_WACT: + # Recompute s_k * y1 = s_k * silu(gate) * up from saved pre_act. + load_mask = k_mask[None, :] & i_mask[:, None] + gate_ptrs = pre_act_ptr + row_offs[None, :] * stride_pre_TK + i_idx[:, None] * stride_pre_N + up_ptrs = gate_ptrs + I_dim * stride_pre_N + gate = tl.load(gate_ptrs, mask=load_mask, other=0.0).to(tl.float32) + up = tl.load(up_ptrs, mask=load_mask, other=0.0).to(tl.float32) + flat_tk_idx = tl.load(s_scatter_idx_ptr + row_offs, mask=k_mask, other=0) + weights = tl.load(topk_weights_ptr + flat_tk_idx, mask=k_mask, other=0.0).to(tl.float32) + y1w = gate * tl.sigmoid(gate) * up * weights[None, :] + # Cast to the activation dtype so numerics match the buffer path. + wact_tile = y1w.to(dout_ptr.dtype.element_ty) + else: + wact_ptrs = weighted_act_ptr + row_offs[None, :] * stride_wact_TK + i_idx[:, None] * stride_wact_I + wact_tile = tl.load(wact_ptrs, mask=k_mask[None, :] & i_mask[:, None], other=0.0) # int64 prevents token_idx * stride_T overflow at large T*H (see #1246). token_idx = tl.load(x_gather_idx_ptr + row_offs, mask=k_mask, other=0).to(tl.int64) @@ -801,7 +1053,8 @@ def _moe_bwd_dW2_kernel( @triton.autotune( configs=_get_gemm_autotune_configs(), - key=["H_dim", "I_dim"], + key=["H_dim", "I_dim", "BLOCK_M", "USE_TMA"], + prune_configs_by={"early_config_prune": _prune_gemm_n_is_H}, ) @triton.jit def _moe_bwd_dX_expanded_kernel( @@ -810,7 +1063,9 @@ def _moe_bwd_dX_expanded_kernel( expert_start_ptr, # (E+1,) int32 tile_row_start_ptr, # (num_m_tiles,) int32 tile_expert_ptr, # (num_m_tiles,) int32 + total_tiles_ptr, # (1,) int32 — actual number of m-tiles (device scalar) dx_expanded_ptr, # (TK, H) — output: clean write, indexed by sorted_pos + w_rows, # E * 2*I — rows of the flattened (E*2I, H) weight view (TMA) H_dim: tl.constexpr, I_dim: tl.constexpr, stride_d_pre_TK, @@ -823,12 +1078,19 @@ def _moe_bwd_dX_expanded_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, + USE_TMA: tl.constexpr, ): - """Grid: (num_m_tiles, ceil(H/BLOCK_N)). - dx_expanded[sorted_pos] = d_gate @ W1_gate^T + d_up @ W1_up^T. + """Grid: 1D (num_m_tiles_max * ceil(H/BLOCK_N),), L2-swizzled into (pid_m, pid_n). + dx_expanded[sorted_pos] = d_pre_act @ W1[e] — a single GEMM over K = 2I + (d_pre_act columns and W1[e] rows share the same [gate; up] ordering). No atomics — rows are unique per CTA in sorted space.""" - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) + pid = tl.program_id(0) + NUM_PID_N: tl.constexpr = (H_dim + BLOCK_N - 1) // BLOCK_N + num_pid_m = tl.load(total_tiles_ptr) + if pid >= num_pid_m * NUM_PID_N: + return + pid_m, pid_n = _grouped_pid_swizzle(pid, num_pid_m, NUM_PID_N, GROUP_M) row_start = tl.load(tile_row_start_ptr + pid_m) # int64 prevents expert_idx * stride_w_E overflow at large E*I*H (see #1246). @@ -848,31 +1110,38 @@ def _moe_bwd_dX_expanded_kernel( acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for k in tl.range(0, I_dim, BLOCK_K): - k_idx = k + k_offs - k_mask = k_idx < I_dim - - d_gate_ptrs = d_pre_act_ptr + row_offs[:, None] * stride_d_pre_TK + k_idx[None, :] * stride_d_pre_N - d_gate = tl.load(d_gate_ptrs, mask=row_mask[:, None] & k_mask[None, :], other=0.0) - - w_gate_ptrs = ( - gate_up_proj_ptr + expert_idx * stride_w_E + k_idx[:, None] * stride_w_N + h_idx[None, :] * stride_w_K + if USE_TMA: + w_desc = tl.make_tensor_descriptor( + gate_up_proj_ptr, + shape=[w_rows, H_dim], + strides=[stride_w_N, 1], + block_shape=[BLOCK_K, BLOCK_N], ) - w_gate = tl.load(w_gate_ptrs, mask=k_mask[:, None] & h_mask[None, :], other=0.0) - acc = tl.dot(d_gate, w_gate, acc=acc) + w_row0 = (expert_idx * (2 * I_dim)).to(tl.int32) + + # dx = d_pre_act @ W1[e] as ONE GEMM over K = 2I: d_pre_act columns and + # W1[e] rows share the same [gate; up] ordering, so the gate/up split the + # forward kernel needs (separate outputs) was never needed here. + # This also avoids a Triton 3.7.1 miscompile on sm103 (B300): two tl.dot + # calls chained through ONE accumulator in a K-loop produce wrong results + # under the tcgen05 MMA path for most tile configs (single-dot loops and + # two-accumulator loops are unaffected). + for k in tl.range(0, 2 * I_dim, BLOCK_K): + k_idx = k + k_offs + k_mask = k_idx < 2 * I_dim - d_up_ptrs = d_pre_act_ptr + row_offs[:, None] * stride_d_pre_TK + (I_dim + k_idx)[None, :] * stride_d_pre_N - d_up = tl.load(d_up_ptrs, mask=row_mask[:, None] & k_mask[None, :], other=0.0) + d_pre_ptrs = d_pre_act_ptr + row_offs[:, None] * stride_d_pre_TK + k_idx[None, :] * stride_d_pre_N + d_pre = tl.load(d_pre_ptrs, mask=row_mask[:, None] & k_mask[None, :], other=0.0) - w_up_ptrs = ( - gate_up_proj_ptr - + expert_idx * stride_w_E - + (I_dim + k_idx)[:, None] * stride_w_N - + h_idx[None, :] * stride_w_K - ) - w_up = tl.load(w_up_ptrs, mask=k_mask[:, None] & h_mask[None, :], other=0.0) + if USE_TMA: + w_tile = w_desc.load([w_row0 + k, n_start]) + else: + w_ptrs = ( + gate_up_proj_ptr + expert_idx * stride_w_E + k_idx[:, None] * stride_w_N + h_idx[None, :] * stride_w_K + ) + w_tile = tl.load(w_ptrs, mask=k_mask[:, None] & h_mask[None, :], other=0.0) - acc = tl.dot(d_up, w_up, acc=acc) + acc = tl.dot(d_pre, w_tile, acc=acc) dxe_ptrs = dx_expanded_ptr + row_offs[:, None] * stride_dxe_TK + h_idx[None, :] * stride_dxe_H tl.store(dxe_ptrs, acc.to(dx_expanded_ptr.dtype.element_ty), mask=row_mask[:, None] & h_mask[None, :]) @@ -886,8 +1155,8 @@ def _moe_bwd_dX_expanded_kernel( @triton.autotune( configs=_get_dW_autotune_configs(), - key=["H_dim", "I_dim"], - reset_to_zero=["dW1_ptr"], + key=["H_dim", "I_dim", "TPE_BUCKET"], + prune_configs_by={"early_config_prune": _prune_dW1}, ) @triton.jit def _moe_bwd_dW1_kernel( @@ -908,25 +1177,29 @@ def _moe_bwd_dW1_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, + TPE_BUCKET: tl.constexpr = 0, # tokens-per-expert bucket; autotune key only (unused in body) ): - """dW1[e, n, h] += sum_t X[token(t), h] * d_pre_act[t, n], where n in [0, 2I). - Grid: (E * ceil(H/BLOCK_M), ceil(2I/BLOCK_N)). Early exit for empty experts.""" - pid0 = tl.program_id(0) - pid1 = tl.program_id(1) + """dW1[e, n, h] = sum_t X[token(t), h] * d_pre_act[t, n], where n in [0, 2I). + Grid: 1D (E * ceil(H/BLOCK_M) * ceil(2I/BLOCK_N),), expert-major so all tiles of + one expert are resident together: its x/d_pre_act rows stay in L2 instead of + being re-read from HBM per tile. Empty experts store zeros (no separate memset).""" + pid = tl.program_id(0) N_M_TILES: tl.constexpr = (H_dim + BLOCK_M - 1) // BLOCK_M + N_N_TILES: tl.constexpr = (2 * I_dim + BLOCK_N - 1) // BLOCK_N + TILES_PER_E: tl.constexpr = N_M_TILES * N_N_TILES # int64 prevents expert_idx * stride_dW_E overflow at large E*I*H (see #1246). - expert_idx = (pid0 // N_M_TILES).to(tl.int64) - m_tile = pid0 % N_M_TILES + expert_idx = (pid // TILES_PER_E).to(tl.int64) + local = pid % TILES_PER_E + m_tile, n_tile = _grouped_pid_swizzle(local, N_M_TILES, N_N_TILES, GROUP_M) expert_start = tl.load(expert_start_ptr + expert_idx) expert_end = tl.load(expert_start_ptr + expert_idx + 1) M_e = expert_end - expert_start - if M_e == 0: - return m_start = m_tile * BLOCK_M - n_start = pid1 * BLOCK_N + n_start = n_tile * BLOCK_N m_offs = tl.arange(0, BLOCK_M) n_offs = tl.arange(0, BLOCK_N) diff --git a/src/liger_kernel/transformers/chunked_grpo_loss.py b/src/liger_kernel/transformers/chunked_grpo_loss.py new file mode 100644 index 000000000..022dcac8d --- /dev/null +++ b/src/liger_kernel/transformers/chunked_grpo_loss.py @@ -0,0 +1,138 @@ +"""Chunked Triton GRPO loss: fused lm_head + GRPO objective without materializing logits. + +Same objective math as :func:`liger_kernel.transformers.grpo_loss.triton_grpo_loss` +but takes (hidden_states, lm_head_weight) instead of logits, so the (N, V) +logits tensor never exists. The heavy work (selective log-softmax through the +lm_head, forward and backward) runs in the fused Triton kernels of +:mod:`liger_kernel.ops.chunked_grpo_loss`; the GRPO objective itself operates +on tiny (B, L) tensors in torch and is taken verbatim from the reference +implementations (TRL's torch path / the Triton kernel's loss math). +""" + +import torch + +from liger_kernel.ops.chunked_grpo_loss import chunked_selective_log_softmax +from liger_kernel.transformers.grpo_loss import _masked_mean +from liger_kernel.transformers.grpo_loss import _reduce_grpo_loss + +_PPO_CLIP_LOSS_TYPES = ("grpo", "bnpo", "dr_grpo", "dapo", "luspo") +_SUPPORTED_LOSS_TYPES = _PPO_CLIP_LOSS_TYPES + ("cispo", "sapo") + + +def chunked_triton_grpo_loss( + hidden, + weight, + old_logp, + ref_logp, + completion_ids, + advantages, + completion_mask=None, + temperature=0.9, + beta=0.04, + eps_low=0.2, + eps_high=0.4, + loss_type="dapo", + max_completion_length=None, + importance_sampling_level="token", + reduce=False, + sapo_temperature_pos=1.0, + sapo_temperature_neg=1.05, + vllm_is_ratio=None, + delta=None, + use_bias_correction_kl=False, + num_items_in_batch=None, +): + """Chunked Triton GRPO loss (fused linear, logits never materialized). + + Args: + hidden: Last hidden states (B, L, H) aligned to predict completion_ids + (i.e. already sliced to the L completion positions, unlike the + (B, L+1, V) logits that triton_grpo_loss takes). + weight: lm_head weight (V, H). + old_logp: Old policy log probabilities (B, L) or None. + ref_logp: Reference policy log probabilities (B, L) or None (required if beta != 0). + completion_ids: Token IDs for completions (B, L). + advantages: Per-sequence advantages (B,). + completion_mask: Mask for valid tokens (B, L) or None. + + Remaining arguments and the return convention match triton_grpo_loss: + If reduce=True: (loss, metrics) where metrics = [kl_mean, clip_ratio] or [clip_ratio] + If reduce=False: (per_token_loss, per_token_kl, is_clipped), all (B, L) + + loss_type "vespo" is not supported. + """ + if loss_type not in _SUPPORTED_LOSS_TYPES: + raise ValueError(f"Unsupported loss_type '{loss_type}' for chunked Triton GRPO loss.") + if importance_sampling_level not in ("token", "sequence"): + raise ValueError(f"importance_sampling_level must be 'token' or 'sequence', got {importance_sampling_level}") + if delta is not None and loss_type not in _PPO_CLIP_LOSS_TYPES: + raise ValueError(f"delta (two-sided clipping) is not supported for loss_type='{loss_type}'.") + if beta != 0.0 and ref_logp is None: + raise ValueError("ref_logp is required when beta != 0.") + + b, seq_len, h = hidden.shape + logp = chunked_selective_log_softmax( + hidden.reshape(-1, h), + weight, + completion_ids.reshape(-1), + temperature, + ).view(b, seq_len) + + mask = completion_mask + if mask is None: + mask = torch.ones_like(logp) + mask = mask.to(logp.dtype) + + old = old_logp.to(torch.float32) if old_logp is not None else logp.detach() + log_ratio = logp - old + if importance_sampling_level == "token": + log_importance_weights = log_ratio + else: + log_importance_weights = ((log_ratio * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).unsqueeze(-1) + + coef_1 = torch.exp(log_importance_weights) + adv = advantages.unsqueeze(1).to(torch.float32) + + if loss_type == "cispo": + clamped_ratios = torch.clamp(coef_1, max=eps_high).detach() + per_token_loss = -clamped_ratios * adv * logp # logp keeps this per-token + is_clipped = ((coef_1 > eps_high) & (adv > 0)).to(logp.dtype) + elif loss_type == "sapo": + temperatures = torch.where(adv > 0, sapo_temperature_pos, sapo_temperature_neg) + soft_coef_1 = torch.sigmoid(temperatures * (coef_1 - 1)) * 4 / temperatures + per_token_loss = -soft_coef_1 * adv + is_clipped = torch.zeros_like(coef_1) + else: # standard PPO clipping: grpo / bnpo / dr_grpo / dapo / luspo + coef_2 = torch.clamp(coef_1, 1 - eps_low, 1 + eps_high) + is_clipped = (((coef_1 < 1 - eps_low) & (adv < 0)) | ((coef_1 > 1 + eps_high) & (adv > 0))).to(logp.dtype) + if delta is not None: + coef_1 = torch.clamp(coef_1, max=delta) + per_token_loss = -torch.min(coef_1 * adv, coef_2 * adv) + + if vllm_is_ratio is not None: + per_token_loss = per_token_loss * vllm_is_ratio + + per_token_kl = None + if beta != 0.0: + ref = ref_logp.to(torch.float32) + per_token_kl = torch.exp(ref - logp) - (ref - logp) - 1 + if use_bias_correction_kl: + per_token_kl = per_token_kl * coef_1 + per_token_loss = per_token_loss + beta * per_token_kl + + # Sequence-level tensors are (B, 1); expand to (B, L) to match the Triton + # kernel's reduce=False convention (same value for all tokens in a sequence). + if per_token_loss.shape[1] == 1: + per_token_loss = per_token_loss.expand(b, seq_len) + if is_clipped.shape[1] == 1: + is_clipped = is_clipped.expand(b, seq_len) + + if not reduce: + return per_token_loss, per_token_kl, is_clipped + + loss = _reduce_grpo_loss(per_token_loss, mask, loss_type, max_completion_length, num_items_in_batch) + metrics = [] + if beta != 0.0: + metrics.append(_masked_mean(per_token_kl, mask)) + metrics.append(_masked_mean(is_clipped, mask)) + return loss, metrics diff --git a/test/transformers/test_chunked_grpo_loss.py b/test/transformers/test_chunked_grpo_loss.py new file mode 100644 index 000000000..4ab9edbc6 --- /dev/null +++ b/test/transformers/test_chunked_grpo_loss.py @@ -0,0 +1,589 @@ +"""Tests for the chunked Triton GRPO loss (fused lm_head, logits never materialized). + +Checks, at multiple context lengths and across loss configurations: + 1. Intermediates: per-token logp/lse vs an fp32 ground truth and vs the + non-chunked Triton path's fused_selective_log_softmax. + 2. Per-token results (reduce=False): per_token_loss / per_token_kl / + is_clipped vs the non-chunked triton_grpo_loss. + 3. End results (reduce=True): loss, metrics, grad_hidden and grad_weight vs + three references: a plain torch implementation from logits (TRL-style, + "non-chunked torch"), the non-chunked Triton kernel, and the chunked + torch LigerFusedLinearGRPOLoss. + 4. Bitwise determinism across reruns (no atomics, fixed launch configs). +""" + +import pytest +import torch + +from test.utils import infer_device +from test.utils import set_seed + +from liger_kernel.chunked_loss import LigerFusedLinearGRPOLoss +from liger_kernel.ops.chunked_grpo_loss import chunked_selective_log_softmax_with_lse +from liger_kernel.ops.grpo_loss import fused_selective_log_softmax +from liger_kernel.transformers.chunked_grpo_loss import chunked_triton_grpo_loss +from liger_kernel.transformers.grpo_loss import _reduce_grpo_loss +from liger_kernel.transformers.grpo_loss import triton_grpo_loss + +device = infer_device() + +HIDDEN_SIZE = 2048 +VOCAB_SIZE = 248320 # Qwen3.5-MoE vocab; divisible by the kernel's BN +ODD_VOCAB_SIZE = 50257 # gpt2 vocab; exercises the vocab-tail masking path + + +# --------------------------------------------------------------------------- +# references +# --------------------------------------------------------------------------- + + +@torch.no_grad() +def fp32_logp_lse(hidden, weight, targets, temperature): + """fp32 ground-truth selective log-softmax, row-chunked to bound memory.""" + n = hidden.shape[0] + logp = torch.empty(n, dtype=torch.float32, device=hidden.device) + lse = torch.empty(n, dtype=torch.float32, device=hidden.device) + chunk = 8192 + w = weight.float() + for i in range(0, n, chunk): + logits = (hidden[i : i + chunk].float() @ w.t()) / temperature + row_lse = torch.logsumexp(logits, dim=-1) + tgt = logits.gather(-1, targets[i : i + chunk].unsqueeze(-1)).squeeze(-1) + logp[i : i + chunk] = tgt - row_lse + lse[i : i + chunk] = row_lse + return logp, lse + + +def torch_grpo_loss_from_logits( + logits, # (B, L, V), already sliced to completion positions + old_logp, + ref_logp, + completion_ids, + advantages, + mask, + *, + temperature, + beta, + eps_low, + eps_high, + loss_type, + max_completion_length, + importance_sampling_level, + sapo_temperature_pos=1.0, + sapo_temperature_neg=1.05, + vllm_is_ratio=None, + delta=None, + use_bias_correction_kl=False, + num_items_in_batch=None, +): + """Plain torch reference (TRL _compute_loss math), written independently of + the implementation under test.""" + logp = torch.nn.functional.log_softmax(logits.float() / temperature, dim=-1) + logp = logp.gather(-1, completion_ids.unsqueeze(-1)).squeeze(-1) + + old = old_logp.float() if old_logp is not None else logp.detach() + log_ratio = logp - old + if importance_sampling_level == "sequence": + log_ratio = ((log_ratio * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).unsqueeze(-1) + coef_1 = torch.exp(log_ratio) + adv = advantages.unsqueeze(1).float() + + if loss_type == "cispo": + per_token_loss = -torch.clamp(coef_1, max=eps_high).detach() * adv * logp + elif loss_type == "sapo": + temps = torch.where(adv > 0, sapo_temperature_pos, sapo_temperature_neg) + per_token_loss = -torch.sigmoid(temps * (coef_1 - 1)) * 4 / temps * adv + else: + coef_2 = torch.clamp(coef_1, 1 - eps_low, 1 + eps_high) + c1 = torch.clamp(coef_1, max=delta) if delta is not None else coef_1 + per_token_loss = -torch.min(c1 * adv, coef_2 * adv) + + if vllm_is_ratio is not None: + per_token_loss = per_token_loss * vllm_is_ratio + if beta != 0.0: + kl = torch.exp(ref_logp.float() - logp) - (ref_logp.float() - logp) - 1 + if use_bias_correction_kl: + kl = kl * coef_1 + per_token_loss = per_token_loss + beta * kl + + if per_token_loss.shape[1] == 1: + per_token_loss = per_token_loss.expand_as(logp) + return _reduce_grpo_loss(per_token_loss, mask, loss_type, max_completion_length, num_items_in_batch) + + +# --------------------------------------------------------------------------- +# fixtures +# --------------------------------------------------------------------------- + + +def make_inputs(batch, seq_len, vocab, *, seed=0, dtype=torch.bfloat16): + set_seed(seed) + # (B, L+1, H): logits-based paths use all L+1 positions and slice + # internally; hidden-based paths take [:, :-1, :]. + hidden = (torch.randn(batch, seq_len + 1, HIDDEN_SIZE, device=device) * 0.02).to(dtype) + weight = (torch.randn(vocab, HIDDEN_SIZE, device=device) * 0.02).to(dtype) + completion_ids = torch.randint(0, vocab, (batch, seq_len), device=device) + lengths = torch.randint(seq_len // 2, seq_len + 1, (batch,), device=device) + mask = (torch.arange(seq_len, device=device).unsqueeze(0) < lengths.unsqueeze(1)).float() + advantages = torch.randn(batch, device=device) + + logp_true, _ = fp32_logp_lse( + hidden[:, :-1, :].reshape(-1, HIDDEN_SIZE).contiguous(), weight, completion_ids.reshape(-1), 1.0 + ) + logp_true = logp_true.view(batch, seq_len) + # old: sizeable perturbation so clipping actually triggers; ref: small one + old_logp = logp_true + torch.randn_like(logp_true) * 0.3 + ref_logp = logp_true + torch.randn_like(logp_true) * 0.1 + return { + "hidden": hidden, + "weight": weight, + "completion_ids": completion_ids, + "mask": mask, + "advantages": advantages, + "old_logp": old_logp, + "ref_logp": ref_logp, + } + + +# their production config + variations exercising every supported branch +CONFIGS = { + "dapo_seq_onpolicy": dict(loss_type="dapo", importance_sampling_level="sequence", beta=0.0, use_old=False), + "dapo_seq_offpolicy_niib": dict( + loss_type="dapo", importance_sampling_level="sequence", beta=0.0, use_old=True, use_niib=True + ), + "dapo_token_kl_biascorr": dict( + loss_type="dapo", + importance_sampling_level="token", + beta=0.04, + use_old=True, + use_bias_correction_kl=True, + ), + "grpo_token_kl": dict(loss_type="grpo", importance_sampling_level="token", beta=0.04, use_old=True), + "bnpo_token": dict(loss_type="bnpo", importance_sampling_level="token", beta=0.0, use_old=True), + "dr_grpo_token": dict( + loss_type="dr_grpo", importance_sampling_level="token", beta=0.0, use_old=True, use_max_len=True + ), + "cispo_token": dict(loss_type="cispo", importance_sampling_level="token", beta=0.0, use_old=True, eps_high=4.0), + "sapo_token": dict(loss_type="sapo", importance_sampling_level="token", beta=0.0, use_old=True), + "dapo_token_delta": dict(loss_type="dapo", importance_sampling_level="token", beta=0.0, use_old=True, delta=4.0), + "dapo_seq_vllm_ratio": dict( + loss_type="dapo", importance_sampling_level="sequence", beta=0.0, use_old=True, use_vllm_ratio=True + ), +} + + +def config_kwargs(cfg, inputs, seq_len): + kwargs = dict( + temperature=1.0, + beta=cfg["beta"], + eps_low=cfg.get("eps_low", 0.2), + eps_high=cfg.get("eps_high", 0.2), + loss_type=cfg["loss_type"], + max_completion_length=seq_len if cfg.get("use_max_len") else None, + importance_sampling_level=cfg["importance_sampling_level"], + delta=cfg.get("delta"), + use_bias_correction_kl=cfg.get("use_bias_correction_kl", False), + num_items_in_batch=inputs["mask"].sum() if cfg.get("use_niib") else None, + vllm_is_ratio=( + (1 + 0.05 * torch.randn(inputs["mask"].shape[0], 1, device=device)).clamp(0.8, 1.2) + if cfg.get("use_vllm_ratio") + else None + ), + ) + old = inputs["old_logp"] if cfg["use_old"] else None + ref = inputs["ref_logp"] if cfg["beta"] != 0.0 else None + return kwargs, old, ref + + +# --------------------------------------------------------------------------- +# 1. intermediates: logp / lse +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "batch, seq_len, vocab", + [ + (4, 128, VOCAB_SIZE), + (4, 1024, VOCAB_SIZE), + (4, 4096, VOCAB_SIZE), + (2, 16384, VOCAB_SIZE), + (4, 512, ODD_VOCAB_SIZE), + (3, 300, ODD_VOCAB_SIZE), # ragged batch/seq/vocab all at once + ], +) +@pytest.mark.parametrize("temperature", [1.0, 0.9]) +def test_logp_lse_intermediates(batch, seq_len, vocab, temperature): + inputs = make_inputs(batch, seq_len, vocab) + hidden2d = inputs["hidden"][:, :-1, :].reshape(-1, HIDDEN_SIZE).contiguous() + ids = inputs["completion_ids"].reshape(-1) + + logp, lse = chunked_selective_log_softmax_with_lse(hidden2d, inputs["weight"], ids, temperature) + logp_ref, lse_ref = fp32_logp_lse(hidden2d, inputs["weight"], ids, temperature) + + # fp32 ground truth: only bf16-GEMM accumulation-order noise separates them + assert torch.allclose(logp, logp_ref, atol=1e-2, rtol=1e-3), ( + f"logp vs fp32 ref: max diff {(logp - logp_ref).abs().max().item():.2e}" + ) + assert torch.allclose(lse, lse_ref, atol=1e-2, rtol=1e-3), ( + f"lse vs fp32 ref: max diff {(lse - lse_ref).abs().max().item():.2e}" + ) + + # non-chunked Triton intermediate (from materialized bf16 logits) + logits = inputs["hidden"] @ inputs["weight"].t() # (B, L+1, V) + logp_triton = fused_selective_log_softmax(logits, inputs["completion_ids"], temperature) + assert torch.allclose(logp.view(batch, seq_len), logp_triton, atol=3e-2, rtol=1e-3), ( + f"logp vs non-chunked triton: max diff {(logp.view(batch, seq_len) - logp_triton).abs().max().item():.2e}" + ) + + +# --------------------------------------------------------------------------- +# 2. per-token results (reduce=False) vs non-chunked Triton +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("config_name", sorted(CONFIGS)) +def test_per_token_results_vs_triton(config_name): + cfg = CONFIGS[config_name] + batch, seq_len = 4, 1024 + inputs = make_inputs(batch, seq_len, VOCAB_SIZE) + kwargs, old, ref = config_kwargs(cfg, inputs, seq_len) + + ptl_c, kl_c, clipped_c = chunked_triton_grpo_loss( + inputs["hidden"][:, :-1, :].contiguous(), + inputs["weight"], + old, + ref, + inputs["completion_ids"], + inputs["advantages"], + inputs["mask"], + reduce=False, + **kwargs, + ) + logits = inputs["hidden"] @ inputs["weight"].t() + ptl_t, kl_t, clipped_t = triton_grpo_loss( + logits, + old, + ref, + inputs["completion_ids"], + inputs["advantages"], + inputs["mask"], + inplace=False, + reduce=False, + **kwargs, + ) + + mask = inputs["mask"].bool() + assert torch.allclose(ptl_c[mask], ptl_t[mask], atol=2e-2, rtol=2e-2), ( + f"per_token_loss: max diff {(ptl_c[mask] - ptl_t[mask]).abs().max().item():.2e}" + ) + if cfg["beta"] != 0.0: + assert torch.allclose(kl_c[mask], kl_t[mask], atol=2e-2, rtol=2e-2), ( + f"per_token_kl: max diff {(kl_c[mask] - kl_t[mask]).abs().max().item():.2e}" + ) + # clipping indicators may flip at ratio~threshold boundaries due to bf16 + # rounding differences; require near-total agreement instead of equality + mismatch = (clipped_c[mask].bool() != clipped_t[mask].bool()).float().mean().item() + assert mismatch < 0.01, f"is_clipped mismatch fraction {mismatch:.4f}" + + +# --------------------------------------------------------------------------- +# 3. end results: loss, metrics, grads vs all three references +# --------------------------------------------------------------------------- + + +def run_chunked_triton(inputs, old, ref, kwargs): + hidden, weight = inputs["hidden"], inputs["weight"] + loss, metrics = chunked_triton_grpo_loss( + hidden[:, :-1, :].contiguous(), + weight, + old, + ref, + inputs["completion_ids"], + inputs["advantages"], + inputs["mask"], + reduce=True, + **kwargs, + ) + return loss, metrics + + +def run_triton(inputs, old, ref, kwargs): + logits = inputs["hidden"] @ inputs["weight"].t() + loss, metrics = triton_grpo_loss( + logits, + old, + ref, + inputs["completion_ids"], + inputs["advantages"], + inputs["mask"], + inplace=False, + reduce=True, + **kwargs, + ) + return loss, metrics + + +def run_chunked_torch(inputs, old, ref, kwargs): + module = LigerFusedLinearGRPOLoss( + beta=kwargs["beta"], + compiled=False, + use_ref_model=kwargs["beta"] != 0.0, + epsilon_low=kwargs["eps_low"], + epsilon_high=kwargs["eps_high"], + loss_type=kwargs["loss_type"], + max_completion_length=kwargs["max_completion_length"], + importance_sampling_level=kwargs["importance_sampling_level"], + temperature=kwargs["temperature"], + delta=kwargs["delta"], + use_bias_correction_kl=kwargs["use_bias_correction_kl"], + ) + loss, metrics = module( + inputs["hidden"][:, :-1, :], + inputs["weight"], + inputs["completion_ids"], + inputs["mask"], + inputs["advantages"], + ref_per_token_logps=ref, + old_per_token_logps=old, + vllm_is_ratio=kwargs["vllm_is_ratio"], + num_items_in_batch=kwargs["num_items_in_batch"], + ) + return loss, metrics + + +def run_torch_reference(inputs, old, ref, kwargs): + logits = (inputs["hidden"] @ inputs["weight"].t())[:, :-1, :] + loss = torch_grpo_loss_from_logits( + logits, + old, + ref, + inputs["completion_ids"], + inputs["advantages"], + inputs["mask"], + **kwargs, + ) + return loss, None + + +IMPLEMENTATIONS = { + "torch_ref": run_torch_reference, + "triton": run_triton, + "chunked_torch": run_chunked_torch, + "chunked_triton": run_chunked_triton, +} + + +def compute_all(inputs, cfg, seq_len): + kwargs, old, ref = config_kwargs(cfg, inputs, seq_len) + results = {} + for name, fn in IMPLEMENTATIONS.items(): + inputs["hidden"].grad = None + inputs["weight"].grad = None + inputs["hidden"].requires_grad_(True) + inputs["weight"].requires_grad_(True) + loss, metrics = fn(inputs, old, ref, kwargs) + loss.backward() + results[name] = { + "loss": loss.item(), + "metrics": [m.item() for m in metrics] if metrics is not None else None, + "grad_hidden": inputs["hidden"].grad.float().flatten().clone(), + "grad_weight": inputs["weight"].grad.float().flatten().clone(), + } + inputs["hidden"].requires_grad_(False) + inputs["weight"].requires_grad_(False) + return results + + +def assert_end_results_match(results): + ref = results["torch_ref"] + for name, res in results.items(): + if name == "torch_ref": + continue + rel = abs(res["loss"] - ref["loss"]) / max(abs(ref["loss"]), 1e-6) + assert rel < 2e-2, f"{name} loss {res['loss']:.6f} vs torch_ref {ref['loss']:.6f} (rel {rel:.2e})" + for key in ("grad_hidden", "grad_weight"): + cos = torch.nn.functional.cosine_similarity(res[key], ref[key], dim=0).item() + norm_ratio = (res[key].norm() / ref[key].norm().clamp(min=1e-12)).item() + assert cos > 0.999, f"{name} {key} cosine {cos:.6f}" + assert abs(norm_ratio - 1) < 2e-2, f"{name} {key} norm ratio {norm_ratio:.4f}" + # metrics parity between the two Triton-based implementations + m_ct, m_t = results["chunked_triton"]["metrics"], results["triton"]["metrics"] + assert len(m_ct) == len(m_t) + for a, b in zip(m_ct, m_t): + assert abs(a - b) < 2e-2, f"metrics mismatch: chunked_triton {m_ct} vs triton {m_t}" + + +@pytest.mark.parametrize("config_name", sorted(CONFIGS)) +def test_end_results_all_configs(config_name): + seq_len = 1024 + inputs = make_inputs(4, seq_len, VOCAB_SIZE, seed=1) + assert_end_results_match(compute_all(inputs, CONFIGS[config_name], seq_len)) + + +@pytest.mark.parametrize( + "batch, seq_len", + [(4, 128), (4, 1024), (4, 4096), (2, 16384)], +) +def test_end_results_context_lengths(batch, seq_len): + """Production config (dapo, sequence-level IS, beta=0) across context lengths.""" + inputs = make_inputs(batch, seq_len, VOCAB_SIZE, seed=2) + cfg = CONFIGS["dapo_seq_offpolicy_niib"] + assert_end_results_match(compute_all(inputs, cfg, seq_len)) + + +def test_end_results_odd_vocab(): + seq_len = 512 + inputs = make_inputs(4, seq_len, ODD_VOCAB_SIZE, seed=3) + cfg = CONFIGS["dapo_seq_offpolicy_niib"] + assert_end_results_match(compute_all(inputs, cfg, seq_len)) + + +# --------------------------------------------------------------------------- +# 4. edge shapes, dtypes, and extreme values +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "batch, seq_len, vocab, dtype", + [ + (1, 1, 256, torch.bfloat16), # single token + (1, 100, 8192, torch.bfloat16), # N below one row tile (BM=128) + (1, 129, 8192, torch.bfloat16), # one over a row tile + (1, 4097, 8192, torch.bfloat16), # one over the backward chunk (4096) + (3, 2731, 8192, torch.bfloat16), # ragged multi-chunk (N=8193) + (3, 77, 250, torch.bfloat16), # vocab smaller than the vocab tile (BN=256) + (3, 77, 251, torch.bfloat16), # prime vocab + (2, 33, 2, torch.bfloat16), # degenerate two-token vocab + (64, 3, 8192, torch.bfloat16), # wide batch, tiny sequences + (4, 512, 50257, torch.float16), + (4, 512, 50257, torch.float32), # needs the 2-stage fp32 SMEM config + (3, 77, 250, torch.float32), # fp32 + sub-tile vocab + ], +) +def test_edge_shapes_and_dtypes(batch, seq_len, vocab, dtype): + set_seed(7) + hidden = (torch.randn(batch, seq_len + 1, HIDDEN_SIZE, device=device) * 0.02).to(dtype) + weight = (torch.randn(vocab, HIDDEN_SIZE, device=device) * 0.02).to(dtype) + completion_ids = torch.randint(0, vocab, (batch, seq_len), device=device) + # pin some targets to the vocab boundaries (exercises tail masking) + completion_ids[:, 0] = 0 + completion_ids[:, -1] = vocab - 1 + lengths = torch.randint(max(seq_len // 2, 1), seq_len + 1, (batch,), device=device) + mask = (torch.arange(seq_len, device=device).unsqueeze(0) < lengths.unsqueeze(1)).float() + advantages = torch.randn(batch, device=device) + + hidden2d = hidden[:, :-1, :].reshape(-1, HIDDEN_SIZE).contiguous() + ids_flat = completion_ids.reshape(-1) + logp, lse = chunked_selective_log_softmax_with_lse(hidden2d, weight, ids_flat, 1.0) + logp_ref, lse_ref = fp32_logp_lse(hidden2d, weight, ids_flat, 1.0) + atol = 1e-2 if dtype != torch.float32 else 1e-3 + assert torch.allclose(logp, logp_ref, atol=atol, rtol=1e-3) + assert torch.allclose(lse, lse_ref, atol=atol, rtol=1e-3) + + old = logp_ref.view(batch, seq_len) + torch.randn(batch, seq_len, device=device) * 0.3 + kwargs = dict( + temperature=1.0, + beta=0.0, + eps_low=0.2, + eps_high=0.2, + loss_type="dapo", + max_completion_length=None, + importance_sampling_level="sequence", + delta=None, + use_bias_correction_kl=False, + num_items_in_batch=mask.sum(), + vllm_is_ratio=None, + ) + hidden.requires_grad_(True) + weight.requires_grad_(True) + loss, _ = chunked_triton_grpo_loss( + hidden[:, :-1, :].contiguous(), + weight, + old, + None, + completion_ids, + advantages, + mask, + reduce=True, + **kwargs, + ) + loss.backward() + gh, gw = hidden.grad.float().flatten().clone(), weight.grad.float().flatten().clone() + hidden.grad = None + weight.grad = None + logits = (hidden @ weight.t())[:, :-1, :] + loss_ref = torch_grpo_loss_from_logits(logits, old, None, completion_ids, advantages, mask, **kwargs) + loss_ref.backward() + rel = abs(loss.item() - loss_ref.item()) / max(abs(loss_ref.item()), 1e-6) + assert rel < 2e-2, f"loss {loss.item():.6f} vs ref {loss_ref.item():.6f}" + for got, ref in ((gh, hidden.grad), (gw, weight.grad)): + ref = ref.float().flatten() + if got.norm() == 0 and ref.norm() == 0: + continue + cos = torch.nn.functional.cosine_similarity(got, ref, dim=0).item() + assert cos > 0.999, f"grad cosine {cos:.6f}" + hidden.requires_grad_(False) + weight.requires_grad_(False) + + +def test_zero_mask_and_extreme_scale(): + set_seed(8) + batch, seq_len, vocab = 4, 256, 8192 + for mask_mode, scale in (("zeros", 0.02), ("ones", 5.0)): + hidden = (torch.randn(batch, seq_len, HIDDEN_SIZE, device=device) * scale).to(torch.bfloat16) + weight = (torch.randn(vocab, HIDDEN_SIZE, device=device) * scale).to(torch.bfloat16) + ids = torch.randint(0, vocab, (batch, seq_len), device=device) + mask = ( + torch.zeros(batch, seq_len, device=device) + if mask_mode == "zeros" + else torch.ones(batch, seq_len, device=device) + ) + adv = torch.randn(batch, device=device) + loss, _ = chunked_triton_grpo_loss( + hidden, + weight, + None, + None, + ids, + adv, + mask, + temperature=1.0, + beta=0.0, + eps_low=0.2, + eps_high=0.2, + loss_type="dapo", + importance_sampling_level="sequence", + reduce=True, + num_items_in_batch=mask.sum(), + ) + assert torch.isfinite(loss), f"non-finite loss with mask={mask_mode}, scale={scale}" + if mask_mode == "zeros": + assert loss.item() == 0.0 + + +# --------------------------------------------------------------------------- +# 5. bitwise determinism +# --------------------------------------------------------------------------- + + +def test_bitwise_determinism(): + seq_len = 2048 + inputs = make_inputs(4, seq_len, VOCAB_SIZE, seed=4) + cfg = CONFIGS["dapo_seq_offpolicy_niib"] + kwargs, old, ref = config_kwargs(cfg, inputs, seq_len) + + runs = [] + for _ in range(2): + inputs["hidden"].grad = None + inputs["weight"].grad = None + inputs["hidden"].requires_grad_(True) + inputs["weight"].requires_grad_(True) + loss, _ = run_chunked_triton(inputs, old, ref, kwargs) + loss.backward() + runs.append((loss.item(), inputs["hidden"].grad.clone(), inputs["weight"].grad.clone())) + inputs["hidden"].requires_grad_(False) + inputs["weight"].requires_grad_(False) + + assert runs[0][0] == runs[1][0], "loss not bitwise deterministic" + assert torch.equal(runs[0][1], runs[1][1]), "grad_hidden not bitwise deterministic" + assert torch.equal(runs[0][2], runs[1][2]), "grad_weight not bitwise deterministic" diff --git a/test/transformers/test_fused_moe.py b/test/transformers/test_fused_moe.py index 9877efffb..fd8cd0df3 100644 --- a/test/transformers/test_fused_moe.py +++ b/test/transformers/test_fused_moe.py @@ -85,7 +85,7 @@ def test_routing_metadata_invariants(T, E, K): logits = torch.randn(T, E, device=device) top_k_index = torch.topk(logits, K, dim=-1).indices.to(torch.int32) - expert_freq, expert_freq_offset, x_gather_idx, s_scatter_idx, s_rev_scatter_idx, _, _ = compute_routing_metadata( + expert_freq, expert_freq_offset, x_gather_idx, s_scatter_idx, s_rev_scatter_idx, *_ = compute_routing_metadata( top_k_index, E ) @@ -208,3 +208,90 @@ def test_K_equals_E(): out = LigerFusedMoEFunction.apply(x, gate_up_proj, down_proj, top_k_index, top_k_weights) ref = _reference_moe_forward(x, gate_up_proj, down_proj, top_k_index, top_k_weights) torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-4) + + +def test_inference_no_grad_path(): + """No input requires grad → the fast inference path (no pre_act store) must + produce the same output as the training path.""" + T, E, H, intermediate_dim, K = 256, 8, 128, 64, 2 + dtype = torch.float32 + x, gate_up_proj, down_proj, top_k_index, top_k_weights = _make_inputs(T, E, H, intermediate_dim, K, dtype, device) + + with torch.no_grad(): + out_infer = LigerFusedMoEFunction.apply(x, gate_up_proj, down_proj, top_k_index, top_k_weights) + + x2 = x.clone().requires_grad_(True) + out_train = LigerFusedMoEFunction.apply(x2, gate_up_proj, down_proj, top_k_index, top_k_weights) + torch.testing.assert_close(out_infer, out_train, atol=1e-5, rtol=1e-5) + + +def test_retain_graph_second_backward(): + """Default mode: backward twice over the same graph (retain_graph) must work + and produce identical gradients (no in-place clobbering of saved tensors).""" + T, E, H, intermediate_dim, K = 64, 4, 64, 32, 2 + dtype = torch.float32 + x, gate_up_proj, down_proj, top_k_index, top_k_weights = _make_inputs(T, E, H, intermediate_dim, K, dtype, device) + x = x.requires_grad_(True) + out = LigerFusedMoEFunction.apply(x, gate_up_proj, down_proj, top_k_index, top_k_weights) + out.sum().backward(retain_graph=True) + g1 = x.grad.clone() + x.grad = None + out.sum().backward() + torch.testing.assert_close(x.grad, g1) + + +_MEM_EFFICIENT_SUBPROCESS_SCRIPT = """ +import torch +import torch.nn as nn + +from liger_kernel.ops import LigerFusedMoEFunction + +from test.transformers.test_fused_moe import _make_inputs, _reference_moe_forward + +device = "cuda" +T, E, H, intermediate_dim, K = 128, 4, 64, 32, 2 +x, gup, dn, idx, wts = _make_inputs(T, E, H, intermediate_dim, K, torch.float32, device) + +x1, gup1, dn1, wts1 = (t.detach().clone().requires_grad_(True) for t in (x, gup, dn, wts)) +x2, gup2, dn2, wts2 = (t.detach().clone().requires_grad_(True) for t in (x, gup, dn, wts)) + +ref = _reference_moe_forward(x1, gup1, dn1, idx, wts1) +ref.sum().backward() +out = LigerFusedMoEFunction.apply(x2, gup2, dn2, idx, wts2) +out.sum().backward(retain_graph=True) + +torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-4) +for a, b in [(x2, x1), (gup2, gup1), (dn2, dn1), (wts2, wts1)]: + torch.testing.assert_close(a.grad, b.grad, atol=3e-3, rtol=1e-2) + +# in-place mode: a second backward over the same graph must raise, not corrupt +try: + out.sum().backward() +except RuntimeError as e: + assert "modified by an inplace operation" in str(e), e +else: + raise AssertionError("second backward did not raise in memory-efficient mode") +print("MEM_EFFICIENT_OK") +""" + + +@pytest.mark.skipif(device != "cuda", reason="subprocess script assumes cuda") +def test_memory_efficient_mode(): + """LIGER_FUSED_MOE_MEMORY_EFFICIENT=1 (import-time flag → subprocess): + gradients must match the reference, and a second backward must raise (SwiGLU + backward runs in place over the saved pre-activations, guarded by a version + bump).""" + import os + import subprocess + import sys + + env = dict(os.environ, LIGER_FUSED_MOE_MEMORY_EFFICIENT="1") + result = subprocess.run( + [sys.executable, "-c", _MEM_EFFICIENT_SUBPROCESS_SCRIPT], + env=env, + capture_output=True, + text=True, + timeout=600, + ) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert "MEM_EFFICIENT_OK" in result.stdout