[Feature] Support EP prefill with num_worst_tokens (#6574)

* support num worst tokens * support num worst tokens * fix build error * support num worst tokens: fix errors * support num worst tokens: fix feild * support num worst tokens: delete requiements * replace permute and depermute op by pure cuda * replace permute and depermute op by pure cuda * fix ci * fix op * fix nan * fix code style --------- Co-authored-by: YuBaoku <49938469+EmmonsCurse@users.noreply.github.com>
2026-04-23 00:17:25 +08:00 · 2026-03-11 17:09:07 +08:00
parent 0466c7e8a8
commit 9f0778f991
21 changed files with 1775 additions and 166 deletions
@@ -1090,6 +1090,12 @@ def parse_args():
        help="Enable overlap schedule",
    )

+    parser.add_argument(
+        "--ep_prefill_use_worst_num_tokens",
+        action="store_true",
+        help="enable to avoid cpu sync",
+    )
+
    args = parser.parse_args()
    return args