[Fix] Wait for PDL before reading DeepSeek V4 K cache locations (#38409)
Co-authored-by: Aurick Qiao <6137920+aurickq@users.noreply.github.com>
This commit is contained in:
co-authored by
Aurick Qiao
parent
87db743021
commit
2123aca87e
@@ -280,10 +280,10 @@ K_KERNEL void fused_k_norm_rope_flashmla(const __grid_constant__ FusedKNormRopeF
|
||||
|
||||
const auto input_ptr = static_cast<const DType*>(params.kv) + work_id * params.kv_stride_batch;
|
||||
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[work_id]);
|
||||
const auto out_loc = params.out_loc[work_id];
|
||||
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
const auto out_loc = params.out_loc[work_id];
|
||||
Float2 data, freq;
|
||||
|
||||
// part 1: norm. Each thread owns one 2-elem pack (the `tx`-th).
|
||||
|
||||
Reference in New Issue
Block a user