Skip to content

Commit a3302f9

Browse files
authored
[Feature] support new sampler backend with triton (#7639)
* support new sampler backend with triton * test case * fix the test case
1 parent 3faf93e commit a3302f9

5 files changed

Lines changed: 1553 additions & 30 deletions

File tree

.flake8

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,3 +5,4 @@ max-line-length = 119
55
# E402: module level import not at top of file
66
per-file-ignores =
77
__init__.py:F401,F403,E402
8+
fastdeploy/model_executor/layers/sample/ops/top_k_top_p_triton.py:E241,E121,E131,E266

fastdeploy/envs.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ def _validate_split_kv_size(value: int) -> int:
6666
# Set attention backend. "NATIVE_ATTN", "APPEND_ATTN"
6767
# and "MLA_ATTN" can be set currently.
6868
"FD_ATTENTION_BACKEND": lambda: os.getenv("FD_ATTENTION_BACKEND", "APPEND_ATTN"),
69-
# Set sampling class. "base", "base_non_truncated", "air" and "rejection" can be set currently.
69+
# Set sampling class. "base", "base_non_truncated", "air", "rejection" and "triton" can be set currently.
7070
"FD_SAMPLING_CLASS": lambda: os.getenv("FD_SAMPLING_CLASS", "base"),
7171
# Set moe backend."cutlass","marlin", "triton", "flashinfer-cutlass", "flashinfer-cutedsl" and "flashinfer-trtllm" can be set currently.
7272
"FD_MOE_BACKEND": lambda: os.getenv("FD_MOE_BACKEND", "cutlass"),

0 commit comments

Comments
 (0)