update cu_seqlens_repeat

This commit is contained in:
jhqxxx
2025-11-04 12:01:41 +08:00
parent b42c5c718c
commit 082bf544d8
2 changed files with 18 additions and 18 deletions
+7 -13
View File
@@ -516,19 +516,11 @@ impl Qwen3VLVisionModel {
let sin = emb.sin()?;
let cu_seqlens = grid_thw.i((.., 1))?.mul(&grid_thw.i((.., 2))?)?;
let grid_t = grid_thw.i((.., 0))?.to_vec1::<u32>()?;
let cu_seqlens_full = match cu_seqlens.rank() {
1 => cu_seqlens.repeat(grid_t[0] as usize)?,
2 => {
let mut cu_seqlens_repeat = Vec::new();
for (index, t) in grid_t.iter().enumerate() {
cu_seqlens_repeat.push(cu_seqlens.i(index)?.repeat(*t as usize)?);
}
Tensor::cat(&cu_seqlens_repeat, 0)?.flatten_all()?
}
_ => {
return Err(anyhow!(format!("create cu_seqlens error")));
}
};
let mut cu_seqlens_repeat = Vec::new();
for (index, t) in grid_t.iter().enumerate() {
cu_seqlens_repeat.push(cu_seqlens.i(index)?.repeat(*t as usize)?);
}
let cu_seqlens_full = Tensor::cat(&cu_seqlens_repeat, 0)?.flatten_all()?;
let cu_seqlens = cu_seqlens_full
.to_dtype(DType::F64)?
.cumsum(0)?
@@ -895,6 +887,8 @@ impl Qwen3VLModel {
let mut v_thw_vec = Vec::new();
for (index, t) in grid_t.iter().enumerate() {
let mut thw_i = thw.i(index)?.to_vec1::<u32>()?;
// [12, 30, 50]
// [1, 30, 50]*t
thw_i[0] = 1;
v_thw_vec.push(
Tensor::new(thw_i, thw.device())?