[kernel] Tune LingBot-Video MoE TMA configs for H100 (#36543)

This commit is contained in:
Xiaoyu Zhang
2026-08-27 02:47:32 -07:00
committed by GitHub
parent 08315c56df
commit ad911a5ec0
3 changed files with 313 additions and 0 deletions
@@ -22,6 +22,25 @@ def test_h200_bf16_config_is_available_for_current_triton_runtime():
assert json.loads(config_path.read_text())["128"]["BLOCK_SIZE_M"] > 0
def test_h100_lingbot_video_configs_enable_tma_only_for_the_tuned_shape():
config_root = (
Path(fused_moe_triton_config.__file__).parent / "configs" / "triton_3_7_1"
)
for suffix in ("", "_down"):
config_path = (
config_root / f"E=128,N=768,device_name=NVIDIA_H100_80GB_HBM3{suffix}.json"
)
configs = json.loads(config_path.read_text())
assert configs["4096"]["USE_TMA"] is True
assert all(
"USE_TMA" not in config
for num_tokens, config in configs.items()
if num_tokens != "4096"
)
def test_down_moe_reuses_tuned_up_config_when_separate_config_is_absent(
monkeypatch, tmp_path
):