Skip to content

Commit e74f074

Browse files
Ceng23333Ceng23333
committed
reduce contiguous
Signed-off-by: Ceng23333 <441651826@qq.com>
1 parent f855d39 commit e74f074

2 files changed

Lines changed: 30 additions & 103 deletions

File tree

‎src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,10 +55,10 @@ void run(void *planned_meta) {
5555
Tensor out_work = out_need_copy_back ? p->out->contiguous() : Tensor(p->out);
5656
auto out_tensor = infinicore::adaptor::to_aten_tensor(out_work);
5757
auto q = infinicore::adaptor::to_aten_tensor(p->q);
58-
#if defined(ENABLE_NVIDIA_API)
58+
#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API)
5959
auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache);
6060
auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache);
61-
#elif defined(ENABLE_QY_API) || defined(ENABLE_METAX_API)
61+
#elif defined(ENABLE_QY_API)
6262
Tensor k_cache_work = p->k_cache->contiguous();
6363
Tensor v_cache_work = p->v_cache->contiguous();
6464
auto k_cache = infinicore::adaptor::to_aten_tensor(k_cache_work);

‎src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc‎

Lines changed: 28 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -48,128 +48,47 @@ void *plan(Tensor out,
4848
namespace {
4949

5050
#ifdef ENABLE_FLASH_ATTN
51-
struct VarlenFlashPrepared {
52-
Tensor k_work;
53-
Tensor v_work;
54-
Tensor out_work_ic;
55-
at::Tensor q;
56-
at::Tensor k;
57-
at::Tensor v;
58-
bool out_need_copy_back;
59-
std::optional<at::Tensor> out_opt;
60-
at::Tensor cu_seqlens_q;
61-
at::Tensor cu_seqlens_kv;
62-
std::optional<at::Tensor> block_table;
63-
std::optional<at::Tensor> alibi_slopes;
64-
int max_seqlen_q;
65-
int max_seqlen_k;
66-
float scale;
67-
};
68-
69-
VarlenFlashPrepared prepare_varlen_flash_tensors(PlannedMeta *p) {
70-
VarlenFlashPrepared t;
71-
// Varlen flash-attn: keep k/v contiguous for dense/paged layout; avoid extra copies for q/metadata when already dense.
72-
t.q = infinicore::adaptor::to_aten_tensor(p->q);
73-
t.k_work = p->k->contiguous();
74-
t.v_work = p->v->contiguous();
75-
t.k = infinicore::adaptor::to_aten_tensor(t.k_work);
76-
t.v = infinicore::adaptor::to_aten_tensor(t.v_work);
77-
78-
t.out_need_copy_back = !p->out->is_contiguous();
79-
t.out_work_ic = t.out_need_copy_back ? p->out->contiguous() : Tensor(p->out);
80-
auto out_work = infinicore::adaptor::to_aten_tensor(t.out_work_ic);
81-
t.out_opt = std::optional<at::Tensor>(out_work);
82-
t.cu_seqlens_q = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_q);
83-
t.cu_seqlens_kv = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_k);
84-
t.block_table = std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(p->block_table));
85-
t.max_seqlen_q = p->max_seqlen_q;
86-
t.max_seqlen_k = p->max_seqlen_k;
87-
t.alibi_slopes = p->alibi_slopes
88-
? std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes))
89-
: std::nullopt;
90-
t.scale = p->scale;
91-
return t;
92-
}
93-
94-
void copy_varlen_flash_output_back(PlannedMeta *p, VarlenFlashPrepared &t) {
95-
if (t.out_need_copy_back) {
96-
p->out->copy_from(t.out_work_ic);
97-
}
98-
}
99-
51+
// MetaX/hpcc pip `flash_attn_2_cuda` exports `mha_varlen_fwd` at global scope (no namespace),
52+
// while NVIDIA `flash-attn-nvidia.so` uses `flash::mha_varlen_fwd`.
10053
#if defined(ENABLE_METAX_API)
101-
/**
102-
* MetaX / hpcc: pip `flash_attn_2_cuda` exposes `mha_varlen_fwd` in the global namespace and
103-
* relies on the current CUDA stream. Kept separate from the NVIDIA `run()` body to reduce merge conflicts.
104-
*/
105-
void run_flashattn_varlen_metax(PlannedMeta *p) {
106-
c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream());
107-
auto t = prepare_varlen_flash_tensors(p);
108-
std::optional<at::Tensor> seqused_k = std::nullopt;
109-
std::optional<const at::Tensor> leftpad_k = std::nullopt;
110-
// `flash_attn_2_cuda` (pip / MetaX) may append an extra trailing argument (`flash_attn_mars_ext_`)
111-
// depending on the HPCC/MetaX stack version.
112-
#if defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
113-
std::optional<at::Tensor> flash_attn_mars_ext = std::nullopt;
114-
#endif
115-
::mha_varlen_fwd(
116-
t.q,
117-
t.k,
118-
t.v,
119-
t.out_opt,
120-
t.cu_seqlens_q,
121-
t.cu_seqlens_kv,
122-
seqused_k,
123-
leftpad_k,
124-
t.block_table,
125-
t.alibi_slopes,
126-
t.max_seqlen_q,
127-
t.max_seqlen_k,
128-
0.0,
129-
t.scale,
130-
false,
131-
true,
132-
-1,
133-
-1,
134-
0.0,
135-
false,
136-
std::nullopt
137-
#if defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
138-
,
139-
flash_attn_mars_ext
140-
#endif
141-
);
142-
copy_varlen_flash_output_back(p, t);
143-
}
54+
#define INFINICORE_FLASH_OP(name) ::name
55+
#else
56+
#define INFINICORE_FLASH_OP(name) flash::name
14457
#endif
14558

14659
#endif // ENABLE_FLASH_ATTN
14760
} // namespace
14861

14962
void run(void *planned_meta) {
15063
#ifdef ENABLE_FLASH_ATTN
151-
auto *p = reinterpret_cast<PlannedMeta *>(planned_meta);
152-
153-
#if defined(ENABLE_METAX_API)
154-
run_flashattn_varlen_metax(p);
155-
#else
15664
c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream());
65+
auto *p = reinterpret_cast<PlannedMeta *>(planned_meta);
15766

15867
auto q = infinicore::adaptor::to_aten_tensor(p->q);
15968
auto k = infinicore::adaptor::to_aten_tensor(p->k);
16069
auto v = infinicore::adaptor::to_aten_tensor(p->v);
161-
auto out = std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(p->out));
70+
71+
const bool out_need_copy_back = !p->out->is_contiguous();
72+
Tensor out_work_ic = out_need_copy_back ? p->out->contiguous() : Tensor(p->out);
73+
auto out_work = infinicore::adaptor::to_aten_tensor(out_work_ic);
74+
auto out = std::optional<at::Tensor>(out_work);
75+
16276
auto cu_seqlens_q = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_q);
16377
auto cu_seqlens_kv = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_k);
16478
std::optional<at::Tensor> seqused_k = std::nullopt;
16579
std::optional<const at::Tensor> leftpad_k = std::nullopt;
16680
auto block_table = std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(p->block_table));
16781
auto max_seqlen_q = p->max_seqlen_q;
16882
auto max_seqlen_k = p->max_seqlen_k;
169-
auto alibi_slopes = p->alibi_slopes ? std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt;
83+
auto alibi_slopes =
84+
p->alibi_slopes ? std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt;
17085
auto scale = p->scale;
17186

172-
flash::mha_varlen_fwd(
87+
#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
88+
std::optional<at::Tensor> flash_attn_mars_ext = std::nullopt;
89+
#endif
90+
91+
INFINICORE_FLASH_OP(mha_varlen_fwd)(
17392
q,
17493
k,
17594
v,
@@ -190,8 +109,16 @@ void run(void *planned_meta) {
190109
-1,
191110
0.0,
192111
false,
193-
std::nullopt);
112+
std::nullopt
113+
#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3)
114+
,
115+
flash_attn_mars_ext
194116
#endif
117+
);
118+
119+
if (out_need_copy_back) {
120+
p->out->copy_from(out_work_ic);
121+
}
195122

196123
#else
197124
throw std::runtime_error("FlashAttention is not enabled in this build");

0 commit comments

Comments
 (0)