[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:
Aurick Qiao
2026-09-14 21:04:28 +08:00
committed by GitHub
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).