[Model Runner] Prepare token count and move FA3 initialization into the graph (#6170)

* prepare for token num and put FA3 init in graph
This commit is contained in:
sunxin
2026-01-26 12:16:57 +08:00
committed by GitHub
parent 0966df78dc
commit adc69c15d0
10 changed files with 64 additions and 42 deletions
@@ -2037,6 +2037,7 @@ __global__ void merge_multi_chunks_kernel(
const int vid = threadIdx.x, hid = threadIdx.y;
const int qid = blockIdx.x;
const uint32_t bid = batch_id_per_token[qid];
if (bid == -1) return;
if (seq_lens_q[bid] <= 0 || seq_lens_kv[bid] <= 0) {
return;
}