Compare commits

...
2 changed files with 9 additions and 1 deletions
@@ -60,7 +60,12 @@ def test_async_tp_pass_correctness(
"mode": CompilationMode.VLLM_COMPILE,
"compile_sizes": [2, 4, 8],
"splitting_ops": [],
"pass_config": {"fuse_gemm_comms": async_tp_enabled},
"pass_config": {
"fuse_gemm_comms": async_tp_enabled,
# Override threshold to ensure SP (required by async TP) is
# applied even for small models.
"sp_min_token_num": 0,
},
}
async_tp_args = [
@@ -258,6 +258,9 @@ def _compare_sp(
"fuse_norm_quant": fuse_norm_quant,
"fuse_act_quant": fuse_act_quant,
"eliminate_noops": True,
# Override threshold to ensure SP is applied even for small
# models (e.g. tiny-random-LlamaForCausalLM).
"sp_min_token_num": 0,
},
"use_inductor_graph_partition": use_inductor_graph_partition,
}