Skip to content

Commit 0dd1cd1

Browse files
yihwang-nvyihwang-nv
committed
[None][refactor] Split thop.attention into phased fmha
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
1 parent a2e0fba commit 0dd1cd1

17 files changed

Lines changed: 2923 additions & 2977 deletions

File tree

‎cpp/tensorrt_llm/CMakeLists.txt‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
1+
# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION &
22
# AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
33
#
44
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
@@ -22,6 +22,26 @@ set(API_INCLUDE_DIR ${PROJECT_SOURCE_DIR}/include)
2222
include_directories(${CMAKE_CURRENT_SOURCE_DIR}/cutlass_extensions/include
2323
${API_INCLUDE_DIR})
2424

25+
get_filename_component(TRTLLM_REPO_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../.."
26+
ABSOLUTE)
27+
set(TRTLLM_FMHA_PARAMS_GENERATOR
28+
"${TRTLLM_REPO_ROOT_DIR}/scripts/generate_fmha_params.py")
29+
set(TRTLLM_FMHA_PARAMS_GENERATED_INCLUDE_DIR
30+
"${CMAKE_CURRENT_BINARY_DIR}/generated")
31+
set(TRTLLM_FMHA_PARAMS_FIELDS_INC
32+
"${TRTLLM_FMHA_PARAMS_GENERATED_INCLUDE_DIR}/tensorrt_llm/thop/fmha_params_fields.inc"
33+
)
34+
set(TRTLLM_FMHA_PARAMS_GENERATED_TARGET trtllm_generate_fmha_params)
35+
add_custom_command(
36+
OUTPUT ${TRTLLM_FMHA_PARAMS_FIELDS_INC}
37+
COMMAND "${Python3_EXECUTABLE}" "${TRTLLM_FMHA_PARAMS_GENERATOR}" --out-dir
38+
"${TRTLLM_FMHA_PARAMS_GENERATED_INCLUDE_DIR}/tensorrt_llm/thop"
39+
DEPENDS "${TRTLLM_FMHA_PARAMS_GENERATOR}"
40+
COMMENT "Generating FmhaParams field include"
41+
VERBATIM)
42+
add_custom_target(${TRTLLM_FMHA_PARAMS_GENERATED_TARGET}
43+
DEPENDS ${TRTLLM_FMHA_PARAMS_FIELDS_INC})
44+
2545
set(TARGET_ARCH "unknown")
2646

2747
message(STATUS "CMAKE_SYSTEM_PROCESSOR: ${CMAKE_SYSTEM_PROCESSOR}")

‎cpp/tensorrt_llm/common/CMakeLists.txt‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#
2-
# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
2+
# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION &
33
# AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
44
#
55
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
@@ -36,6 +36,9 @@ add_library(common_src OBJECT ${SRCS} ${CU_SRCS})
3636
add_cuda_architectures(common_src 89)
3737
set_property(TARGET common_src PROPERTY POSITION_INDEPENDENT_CODE ON)
3838
set_property(TARGET common_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
39+
target_include_directories(common_src
40+
PRIVATE ${TRTLLM_FMHA_PARAMS_GENERATED_INCLUDE_DIR})
41+
add_dependencies(common_src ${TRTLLM_FMHA_PARAMS_GENERATED_TARGET})
3942
target_link_libraries(common_src PUBLIC trtllm_gen_fmha_interface)
4043

4144
if(ENABLE_CUBLASLT_FP4_GEMM)

‎cpp/tensorrt_llm/common/attentionOp.cpp‎

Lines changed: 574 additions & 580 deletions
Large diffs are not rendered by default.

‎cpp/tensorrt_llm/common/attentionOp.h‎

Lines changed: 73 additions & 307 deletions
Large diffs are not rendered by default.

‎cpp/tensorrt_llm/nanobind/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,9 @@ endif()
5151
# compilation of kvCacheManagerV2.cpp.
5252
target_include_directories(
5353
${TRTLLM_NB_MODULE} PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/sha256)
54+
target_include_directories(${TRTLLM_NB_MODULE}
55+
PRIVATE ${TRTLLM_FMHA_PARAMS_GENERATED_INCLUDE_DIR})
56+
add_dependencies(${TRTLLM_NB_MODULE} ${TRTLLM_FMHA_PARAMS_GENERATED_TARGET})
5457

5558
set_property(TARGET ${TRTLLM_NB_MODULE} PROPERTY POSITION_INDEPENDENT_CODE ON)
5659

‎cpp/tensorrt_llm/nanobind/thop/bindings.cpp‎

Lines changed: 105 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
#include <tensorrt_llm/thop/outputTensor.h>
2727
#include <tensorrt_llm/thop/trtllmGenFusedOps.h>
2828
#include <torch/extension.h>
29+
#include <type_traits>
2930

3031
namespace nb = nanobind;
3132

@@ -45,6 +46,93 @@ nb::object optionalToObject(std::optional<T> const& value)
4546
return nb::none();
4647
}
4748

49+
template <typename T>
50+
struct OptionalInner
51+
{
52+
};
53+
54+
template <typename T>
55+
struct OptionalInner<std::optional<T>>
56+
{
57+
using Type = T;
58+
};
59+
60+
template <typename T>
61+
inline constexpr bool kIsOptional = false;
62+
63+
template <typename T>
64+
inline constexpr bool kIsOptional<std::optional<T>> = true;
65+
66+
template <typename T>
67+
nb::object fmhaFieldToObject(T const& value)
68+
{
69+
if constexpr (std::is_enum_v<T>)
70+
{
71+
return nb::cast(static_cast<int64_t>(static_cast<std::underlying_type_t<T>>(value)));
72+
}
73+
else if constexpr (std::is_same_v<T, tensorrt_llm::common::QuantMode>)
74+
{
75+
return nb::cast(static_cast<int64_t>(value.value()));
76+
}
77+
else
78+
{
79+
return nb::cast(value);
80+
}
81+
}
82+
83+
template <typename T>
84+
nb::object fmhaFieldToObject(std::optional<T> const& value)
85+
{
86+
if (value.has_value())
87+
{
88+
return fmhaFieldToObject(*value);
89+
}
90+
return nb::none();
91+
}
92+
93+
template <typename T>
94+
T castFmhaField(nb::object value)
95+
{
96+
if constexpr (std::is_enum_v<T>)
97+
{
98+
return static_cast<T>(static_cast<std::underlying_type_t<T>>(nb::cast<int64_t>(value)));
99+
}
100+
else if constexpr (std::is_same_v<T, tensorrt_llm::common::QuantMode>)
101+
{
102+
return T(static_cast<typename T::BaseType>(nb::cast<int64_t>(value)));
103+
}
104+
else
105+
{
106+
return nb::cast<T>(value);
107+
}
108+
}
109+
110+
template <typename FieldT>
111+
void bindFmhaParamField(nb::class_<torch_ext::FmhaParams>& cls, char const* name, FieldT torch_ext::FmhaParams::*member)
112+
{
113+
cls.def_prop_rw(
114+
name, [member](torch_ext::FmhaParams const& self) { return fmhaFieldToObject(self.*member); },
115+
[member](torch_ext::FmhaParams& self, nb::object value)
116+
{
117+
if constexpr (kIsOptional<FieldT>)
118+
{
119+
if (value.is_none())
120+
{
121+
self.*member = std::nullopt;
122+
}
123+
else
124+
{
125+
using InnerT = typename OptionalInner<FieldT>::Type;
126+
self.*member = castFmhaField<InnerT>(value);
127+
}
128+
}
129+
else
130+
{
131+
self.*member = castFmhaField<FieldT>(value);
132+
}
133+
});
134+
}
135+
48136
nb::tuple trtllmGenContextPreprocessBinding(torch::Tensor qkv_input, torch::Tensor workspace,
49137
torch::Tensor sequence_lengths, torch::Tensor context_lengths, std::optional<torch::Tensor> kv_cache_block_offsets,
50138
std::optional<torch::Tensor> host_kv_cache_pool_pointers, std::optional<torch::Tensor> host_kv_cache_pool_mapping,
@@ -133,55 +221,23 @@ void initBindings(nb::module_& m)
133221
m.attr(kv.first) = kv.second;
134222
}
135223

136-
m.def("attention", &torch_ext::attention,
137-
// Parameters with default values: using std::nullopt for trailing optional arguments (omittable)
138-
// and .none() for optional arguments followed by parameters without default values (not omittable).
139-
nb::arg("q"), nb::arg("k").none(), nb::arg("v").none(), nb::arg("output"), nb::arg("output_sf").none(),
140-
nb::arg("workspace_").none(), nb::arg("sequence_length"), nb::arg("host_past_key_value_lengths"),
141-
nb::arg("host_total_kv_lens"), nb::arg("context_lengths"), nb::arg("host_context_lengths"),
142-
nb::arg("host_request_types"), nb::arg("max_context_q_len_override").none(),
143-
nb::arg("kv_cache_block_offsets").none(), nb::arg("host_kv_cache_pool_pointers").none(),
144-
nb::arg("host_kv_cache_pool_mapping").none(), nb::arg("cache_indirection").none(),
145-
nb::arg("kv_scale_orig_quant").none(), nb::arg("kv_scale_quant_orig").none(), nb::arg("out_scale").none(),
146-
nb::arg("rotary_inv_freq").none(), nb::arg("rotary_cos_sin").none(), nb::arg("latent_cache").none(),
147-
nb::arg("q_pe").none(), nb::arg("block_ids_per_seq").none(), nb::arg("attention_sinks").none(),
148-
nb::arg("is_fused_qkv"), nb::arg("update_kv_cache"), nb::arg("predicted_tokens_per_seq"),
149-
nb::arg("local_layer_idx"), nb::arg("num_heads"), nb::arg("num_kv_heads"), nb::arg("head_size"),
150-
nb::arg("tokens_per_block").none(), nb::arg("max_num_requests"), nb::arg("max_context_length"),
151-
nb::arg("max_seq_len"), nb::arg("attention_window_size"), nb::arg("beam_width"), nb::arg("mask_type"),
152-
nb::arg("quant_mode"), nb::arg("q_scaling"), nb::arg("position_embedding_type"), nb::arg("rope_dim"),
153-
nb::arg("rope_base"), nb::arg("rope_scale_type"), nb::arg("rope_scale"), nb::arg("rope_short_m_scale"),
154-
nb::arg("rope_long_m_scale"), nb::arg("rope_max_positions"), nb::arg("rope_original_max_positions"),
155-
nb::arg("use_paged_context_fmha"), nb::arg("attention_input_type").none(), nb::arg("is_mla_enable"),
156-
nb::arg("chunked_prefill_buffer_batch_size").none(), nb::arg("q_lora_rank").none(),
157-
nb::arg("kv_lora_rank").none(), nb::arg("qk_nope_head_dim").none(), nb::arg("qk_rope_head_dim").none(),
158-
nb::arg("v_head_dim").none(), nb::arg("rope_append").none(), nb::arg("mrope_rotary_cos_sin").none(),
159-
nb::arg("mrope_position_deltas").none(), nb::arg("helix_position_offsets").none(),
160-
nb::arg("helix_is_inactive_rank").none(), nb::arg("attention_chunk_size").none(),
161-
nb::arg("softmax_stats_tensor").none(), nb::arg("is_spec_decoding_enabled"), nb::arg("use_spec_decoding"),
162-
nb::arg("is_spec_dec_tree"), nb::arg("spec_decoding_generation_lengths").none(),
163-
nb::arg("spec_decoding_position_offsets_for_cpp").none(), nb::arg("spec_decoding_packed_mask").none(),
164-
nb::arg("spec_decoding_bl_tree_mask_offset").none(), nb::arg("spec_decoding_bl_tree_mask").none(),
165-
nb::arg("spec_bl_tree_first_sparse_mask_offset_kv").none(), nb::arg("sparse_kv_indices").none(),
166-
nb::arg("sparse_kv_offsets").none(), nb::arg("sparse_attn_indices").none(),
167-
nb::arg("sparse_attn_offsets").none(), nb::arg("sparse_attn_indices_block_size"),
168-
nb::arg("num_sparse_topk") = std::nullopt, nb::arg("sparse_attn_kv_lens") = std::nullopt,
169-
nb::arg("skip_softmax_threshold_scale_factor_prefill") = std::nullopt,
170-
nb::arg("skip_softmax_threshold_scale_factor_decode") = std::nullopt,
171-
nb::arg("skip_softmax_stat") = std::nullopt, nb::arg("cu_q_seqlens") = std::nullopt,
172-
nb::arg("cu_kv_seqlens") = std::nullopt, nb::arg("fmha_scheduler_counter") = std::nullopt,
173-
nb::arg("mla_bmm1_scale") = std::nullopt, nb::arg("mla_bmm2_scale") = std::nullopt,
174-
nb::arg("quant_q_buffer") = std::nullopt, nb::arg("flash_mla_tile_scheduler_metadata") = std::nullopt,
175-
nb::arg("flash_mla_num_splits") = std::nullopt, nb::arg("sage_attn_num_elts_per_blk_q") = 0,
176-
nb::arg("sage_attn_num_elts_per_blk_k") = 0, nb::arg("sage_attn_num_elts_per_blk_v") = 0,
177-
nb::arg("sage_attn_qk_int8") = false, nb::arg("num_contexts") = 0, nb::arg("num_ctx_tokens") = 0,
178-
nb::arg("trtllm_gen_jit_warmup") = false, nb::arg("aux_kv_cache_pool_ptr") = std::nullopt,
179-
nb::arg("is_cross") = false, nb::arg("cross_kv") = std::nullopt,
180-
nb::arg("relative_attention_bias") = std::nullopt, nb::arg("relative_attention_max_distance") = 0,
181-
nb::arg("spec_decoding_target_max_draft_tokens") = std::nullopt, nb::arg("quant_scale_qkv") = std::nullopt,
182-
nb::arg("dsv4_inv_rope_cos_sin_cache") = std::nullopt, nb::arg("enable_dsv4_epilogue_fusion") = false,
183-
nb::arg("force_prepare_spec_dec_tree_mask") = false, nb::arg("max_num_sequences") = std::nullopt,
184-
"Multi-head attention operation", nb::call_guard<nb::gil_scoped_release>());
224+
// ---- Phased attention API ----
225+
auto fmhaParams = nb::class_<torch_ext::FmhaParams>(m, "FmhaParams").def(nb::init<>());
226+
#define TRTLLM_FMHA_PARAM_FIELD(name, cpp_type, py_type, cpp_default, py_default) \
227+
bindFmhaParamField(fmhaParams, #name, &torch_ext::FmhaParams::name);
228+
#include "tensorrt_llm/thop/fmha_params_fields.inc"
229+
#undef TRTLLM_FMHA_PARAM_FIELD
230+
231+
m.def("run_context", &torch_ext::run_context, nb::arg("params"), "Phased attention run_context.",
232+
nb::call_guard<nb::gil_scoped_release>());
233+
m.def("run_generation", &torch_ext::run_generation, nb::arg("params"), "Phased attention run_generation.",
234+
nb::call_guard<nb::gil_scoped_release>());
235+
m.def("run_mla_generation", &torch_ext::run_mla_generation, nb::arg("params"),
236+
"Phased attention run_mla_generation.", nb::call_guard<nb::gil_scoped_release>());
237+
m.def("get_attention_workspace_size", &torch_ext::get_attention_workspace_size, nb::arg("params"),
238+
nb::arg("num_tokens"), nb::arg("max_attention_window_size"), nb::arg("num_gen_tokens"),
239+
nb::arg("max_blocks_per_sequence"), nb::arg("ctx_total_kv_len") = 0,
240+
"Max of the context/generation workspace byte requirements for sizing FmhaParams.workspace.");
185241

186242
m.def(
187243
"get_helix_workspace_size_per_rank",

‎cpp/tensorrt_llm/thop/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,9 @@ add_library(
146146
mhcOp.cpp
147147
compressorOp.cpp)
148148
set_property(TARGET th_common PROPERTY POSITION_INDEPENDENT_CODE ON)
149+
target_include_directories(th_common
150+
PRIVATE ${TRTLLM_FMHA_PARAMS_GENERATED_INCLUDE_DIR})
151+
add_dependencies(th_common ${TRTLLM_FMHA_PARAMS_GENERATED_TARGET})
149152
target_link_libraries(
150153
th_common PRIVATE ${TORCH_LIBRARIES} th_utils ${Python3_LIBRARIES}
151154
${SHARED_TARGET} pg_utils)

0 commit comments

Comments
 (0)