anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Lance] Restore deploy YAML + patch post-main-merge API drift

Browse Source 
After merging main into ``lance-model-support`` (commit c15f51a2),
Bagel's internal API was aggressively trimmed: ``forward_cache_update_text``,
``forward_cache_update_vit``, ``generate_image``, ``generate_text``, and
``Qwen2MoTForCausalLM.forward`` no longer accept ``packed_text_indexes`` /
``packed_key_value_indexes`` / ``key_values_lens`` / ``packed_query_indexes``
— they're now derived from ``past_key_values``.  Lance's pipeline was
still passing those kwargs and crashed on every task.

Changes:

- ``vllm_omni/deploy/lance.yaml`` (restored): Lance's HF ``config.json``
  is descriptive metadata with no ``model_type`` field, so the loader
  can't auto-detect the pipeline.  This YAML is the minimal hook that
  tells the loader "use the Lance pipeline" — one stage, no multi-stage
  orchestration.
- ``examples/offline_inference/lance/gradio_demo.py``: pass
  ``deploy_config`` again so ``Omni`` picks up the YAML.
- ``vllm_omni/diffusion/models/bagel/bagel_transformer.py``: in the
  batched CFG setup, concat ``branches_pid`` along ``dim=-1`` (the
  sequence axis) instead of the default ``dim=0`` so Lance's 3-D mRoPE
  position ids ``(3, S)`` aren't collapsed into ``(6, S)`` —
  ``BagelRotaryEmbedding`` enforces exactly 3 axes for mRoPE.
- ``vllm_omni/diffusion/models/bagel/pipeline_bagel.py``: mirror the
  generated PIL image into ``custom_output["image"]`` so callers reading
  through the orchestrator IPC boundary (which strips the bare
  ``output`` field) can still recover it.  ``video_frames`` already used
  this pattern.
- ``vllm_omni/diffusion/models/lance/lance_transformer.py``: drop
  ``packed_query_indexes`` / ``key_values_lens`` /
  ``packed_key_value_indexes`` from the ``language_model.forward`` call
  inside ``forward_cache_update_vae`` — derived now.
- ``vllm_omni/diffusion/models/lance/pipeline_lance.py``:
    * drop the same 3 fields from the inline ``inp`` dicts that build
      ``forward_cache_update_text`` arguments in ``_forward_image_edit``,
      ``_forward_i2v``, ``_forward_video_edit``, and
      ``_x2t_raw_text_prefill``.
    * drop the 6 ``cfg_text_packed_query_indexes`` /
      ``cfg_text_key_values_lens`` / ``cfg_text_packed_key_value_indexes``
      (+ ``cfg_img_*`` triple) kwargs from the four ``generate_image``
      call sites.
    * filter ``key_values_lens`` / ``packed_indexes`` /
      ``packed_key_value_indexes`` out of the ``gen_input_lat`` /
      ``gen_input_vit`` dicts before unpacking into ``generate_image``
      and ``forward_cache_update_vit`` (the ``prepare_*_latent`` /
      ``prepare_vit_*`` helpers still return them for the
      backwards-compatible parent-class path).
    * drop ``key_values_lens`` / ``packed_key_value_indexes`` from
      ``_x2t_prepare_assistant_start``'s returned dict so
      ``Bagel.generate_text`` doesn't reject them.

Verification (local online inference, single GPU pair, CFG=4, 10 steps):

| Task          | Status                                          |
|---------------|-------------------------------------------------|
| t2i           | OK  (~2 s)                                      |
| image_edit    | OK  (~2.5 s)                                    |
| i2v           | OK  (~8 s @ 25 frames)                          |
| t2v           | OK  (~7 s @ 25 frames)                          |
| video_edit    | OK  (~60 s)                                     |
| x2t_image     | RUNS but generates degenerate token repetition  |
| x2t_video     | RUNS but generates degenerate token repetition  |

x2t degeneration is a deeper semantic mismatch with main's new
``past_key_values``-derived cache layout: Lance's x2t prefill stitches
the chat template across multiple ``forward_cache_update_*`` calls
(system → ViT → question → ``assistant\\n``), but the new API expects a
single contiguous KV layout.  Fixing requires reworking the segmented
prefill to match main's new cache contract — left as a follow-up since
the 5 generative tasks now work.

Signed-off-by: lishunyang12 <lishunyang12@163.com>
This commit is contained in:
lishunyang12 committed 2026-05-31 13:28:11 +00:00
1 parent c15f51a2f6
commit a326cd0d57
6 files changed
+127 -60
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

examples/offline_inference/lance/gradio_demo.py
+5 -12
View File
@@ -288,18 +288,11 @@ def _init_models(
f" replicas={replicas_per_omni}...",
flush=True,
)
# Lance is single-stage diffusion — no deploy YAML. The defaults
# below match the (now-removed) ``vllm_omni/deploy/lance.yaml`` so
# ``create_default_diffusion`` produces the same stage config.
lance_engine_defaults = dict(
pipeline="lance",
max_num_batched_tokens=32768,
max_num_seqs=1,
enforce_eager=True,
trust_remote_code=True,
enable_prefix_caching=False,
async_chunk=False,
)
# Lance needs ``vllm_omni/deploy/lance.yaml`` as the pipeline-selector
# (Lance HF config has no ``model_type`` field for auto-detect). Pass
# via ``deploy_config`` — flat kwargs alone don't trigger pipeline
# resolution.
lance_engine_defaults = dict(deploy_config="vllm_omni/deploy/lance.yaml")
img_kwargs: dict = {"model": img_ckpt, **lance_engine_defaults}
if img_devices != "0":
vllm_omni/deploy/lance.yaml
+20
View File
@@ -0,0 +1,20 @@
# Lance (ByteDance) — single-stage diffusion pipeline selector.
#
# Lance's HF config.json is descriptive metadata (no ``model_type``),
# so the loader can't auto-detect from the model dir alone. This YAML
# only tells the loader "use the Lance pipeline" — one stage, no multi-
# stage orchestration.
pipeline: lance
async_chunk: false
stages:
- stage_id: 0
max_num_batched_tokens: 32768
max_num_seqs: 1
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
devices: "0"
default_sampling_params:
seed: 42
vllm_omni/diffusion/models/bagel/bagel_transformer.py
+5 -1
View File
@@ -1852,7 +1852,11 @@ class Bagel(nn.Module):
"num_branches": num_branches,
"seq_len": seq_len,
"batched_query_lens": packed_seqlens.repeat(num_branches),
"batched_position_ids": torch.cat(branches_pid),
# Concat along the *sequence* (last) dim — works for both
# scalar 1-D position ids ``(S,)`` and Lance's 3-D mRoPE
# ``(3, S)``. Default ``dim=0`` collapses ``(3, S) + (3, S)``
# to ``(6, S)`` and breaks ``BagelRotaryEmbedding``.
"batched_position_ids": torch.cat(branches_pid, dim=-1),
"batched_text_indexes": torch.cat(
[packed_text_indexes + b_idx * seq_len for b_idx in range(num_branches)]
),
vllm_omni/diffusion/models/bagel/pipeline_bagel.py
+5
View File
@@ -820,6 +820,11 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
custom = {}
if think_text is not None:
custom["think_text"] = think_text
# Mirror the PIL image into ``custom_output`` so callers reading via
# the orchestrator IPC boundary (which strips the bare ``output``
# field) can still recover the result. ``video_frames`` already
# uses this pattern.
custom["image"] = img
return DiffusionOutput(
output=img,
vllm_omni/diffusion/models/lance/lance_transformer.py
+3 -3
View File
@@ -1535,14 +1535,14 @@ class LanceBagel(Bagel):
"packed_text_indexes": packed_text_indexes,
}
# ``Qwen2MoTForCausalLM.forward`` post-main-merge no longer accepts
# ``packed_query_indexes`` / ``key_values_lens`` /
# ``packed_key_value_indexes`` — derived internally now.
output = self.language_model.forward(
packed_query_sequence=packed_query_sequence,
query_lens=packed_seqlens,
packed_query_position_ids=packed_position_ids,
packed_query_indexes=packed_indexes,
past_key_values=past_key_values,
key_values_lens=key_values_lens,
packed_key_value_indexes=packed_key_value_indexes,
update_past_key_values=True,
is_causal=False,
**extra_inputs,
vllm_omni/diffusion/models/lance/pipeline_lance.py
+89 -44
View File
@@ -669,6 +669,12 @@ class LancePipeline(BagelPipeline):
enabled=self.device.type != "cpu",
dtype=self.od_config.dtype,
):
# ``prepare_video_latent`` / ``prepare_vae_latent`` still return
# ``key_values_lens`` / ``packed_indexes`` / ``packed_key_value_indexes``
# but post-main-merge ``generate_image`` doesn't accept them — drop
# before unpacking via ``**gen_input_lat``.
for _drop in ("key_values_lens", "packed_indexes", "packed_key_value_indexes"):
gen_input_lat.pop(_drop, None)
latents, *_ = self.bagel.generate_image(
past_key_values=gen_context["past_key_values"],
cfg_text_past_key_values=cfg_text_context["past_key_values"],
@@ -682,13 +688,11 @@ class LancePipeline(BagelPipeline):
cfg_renorm_min=cfg_renorm_min,
**gen_input_lat,
cfg_text_packed_position_ids=cfg_text_lat["cfg_packed_position_ids"],
cfg_text_packed_query_indexes=cfg_text_lat["cfg_packed_query_indexes"],
cfg_text_key_values_lens=cfg_text_lat["cfg_key_values_lens"],
cfg_text_packed_key_value_indexes=cfg_text_lat["cfg_packed_key_value_indexes"],
# ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
# / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
# (derived from ``cfg_text_past_key_values``).
cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
cfg_img_packed_query_indexes=cfg_img_lat["cfg_packed_query_indexes"],
cfg_img_key_values_lens=cfg_img_lat["cfg_key_values_lens"],
cfg_img_packed_key_value_indexes=cfg_img_lat["cfg_packed_key_value_indexes"],
# ``cfg_img_*`` index/lens kwargs removed — same as above.
)
frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape)
@@ -836,17 +840,17 @@ class LancePipeline(BagelPipeline):
curr_kvlen = ctx["kv_lens"][0]
curr_rope = ctx["ropes"][0]
seq_len = len(text_ids)
# ``forward_cache_update_text`` post-main-merge derives the
# remaining positions from ``past_key_values`` + the query
# length; the old ``packed_text_indexes`` /
# ``packed_key_value_indexes`` / ``key_values_lens`` kwargs
# were dropped from its signature.
inp = {
"text_token_lens": torch.tensor([seq_len], dtype=torch.int, device=self.device),
"packed_text_ids": torch.tensor(text_ids, dtype=torch.long, device=self.device),
"packed_text_position_ids": torch.arange(
curr_rope, curr_rope + seq_len, dtype=torch.long, device=self.device
),
"packed_text_indexes": torch.arange(
curr_kvlen, curr_kvlen + seq_len, dtype=torch.long, device=self.device
),
"packed_key_value_indexes": torch.arange(0, curr_kvlen, dtype=torch.long, device=self.device),
"key_values_lens": torch.tensor([curr_kvlen], dtype=torch.int, device=self.device),
}
with torch.autocast(**autocast_kwargs):
ctx["past_key_values"] = self.bagel.forward_cache_update_text(ctx["past_key_values"], **inp)
@@ -864,6 +868,12 @@ class LancePipeline(BagelPipeline):
for k, v in inp.items():
if torch.is_tensor(v):
inp[k] = v.to(self.device)
# ``prepare_vit_images`` still returns ``packed_indexes`` /
# ``packed_key_value_indexes`` / ``key_values_lens`` for
# historical reasons; ``forward_cache_update_vit`` post-main
# no longer accepts them — pop before unpacking.
for _drop in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"):
inp.pop(_drop, None)
with torch.autocast(**autocast_kwargs):
ctx["past_key_values"] = self.bagel.forward_cache_update_vit(ctx["past_key_values"], **inp)
ctx["kv_lens"] = new_kvlens
@@ -979,6 +989,12 @@ class LancePipeline(BagelPipeline):
cfg_img_lat = cfg_text_lat
with torch.autocast(**autocast_kwargs):
# ``prepare_video_latent`` / ``prepare_vae_latent`` still return
# ``key_values_lens`` / ``packed_indexes`` / ``packed_key_value_indexes``
# but post-main-merge ``generate_image`` doesn't accept them — drop
# before unpacking via ``**gen_input_lat``.
for _drop in ("key_values_lens", "packed_indexes", "packed_key_value_indexes"):
gen_input_lat.pop(_drop, None)
latents, *_ = self.bagel.generate_image(
past_key_values=gen_context["past_key_values"],
cfg_text_past_key_values=cfg_text_context["past_key_values"]
@@ -994,13 +1010,11 @@ class LancePipeline(BagelPipeline):
cfg_renorm_min=cfg_renorm_min,
**gen_input_lat,
cfg_text_packed_position_ids=cfg_text_lat["cfg_packed_position_ids"],
cfg_text_packed_query_indexes=cfg_text_lat["cfg_packed_query_indexes"],
cfg_text_key_values_lens=cfg_text_lat["cfg_key_values_lens"],
cfg_text_packed_key_value_indexes=cfg_text_lat["cfg_packed_key_value_indexes"],
# ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
# / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
# (derived from ``cfg_text_past_key_values``).
cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
cfg_img_packed_query_indexes=cfg_img_lat["cfg_packed_query_indexes"],
cfg_img_key_values_lens=cfg_img_lat["cfg_key_values_lens"],
cfg_img_packed_key_value_indexes=cfg_img_lat["cfg_packed_key_value_indexes"],
# ``cfg_img_*`` index/lens kwargs removed — same as above.
)
img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape)
@@ -1162,17 +1176,17 @@ class LancePipeline(BagelPipeline):
curr_kvlen = ctx["kv_lens"][0]
curr_rope = ctx["ropes"][0]
seq_len = len(text_ids)
# ``forward_cache_update_text`` post-main-merge derives the
# remaining positions from ``past_key_values`` + the query
# length; the old ``packed_text_indexes`` /
# ``packed_key_value_indexes`` / ``key_values_lens`` kwargs
# were dropped from its signature.
inp = {
"text_token_lens": torch.tensor([seq_len], dtype=torch.int, device=self.device),
"packed_text_ids": torch.tensor(text_ids, dtype=torch.long, device=self.device),
"packed_text_position_ids": torch.arange(
curr_rope, curr_rope + seq_len, dtype=torch.long, device=self.device
),
"packed_text_indexes": torch.arange(
curr_kvlen, curr_kvlen + seq_len, dtype=torch.long, device=self.device
),
"packed_key_value_indexes": torch.arange(0, curr_kvlen, dtype=torch.long, device=self.device),
"key_values_lens": torch.tensor([curr_kvlen], dtype=torch.int, device=self.device),
}
with torch.autocast(**autocast_kwargs):
ctx["past_key_values"] = self.bagel.forward_cache_update_text(ctx["past_key_values"], **inp)
@@ -1190,6 +1204,12 @@ class LancePipeline(BagelPipeline):
for k, v in inp.items():
if torch.is_tensor(v):
inp[k] = v.to(self.device)
# ``prepare_vit_images`` still returns ``packed_indexes`` /
# ``packed_key_value_indexes`` / ``key_values_lens`` for
# historical reasons; ``forward_cache_update_vit`` post-main
# no longer accepts them — pop before unpacking.
for _drop in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"):
inp.pop(_drop, None)
with torch.autocast(**autocast_kwargs):
ctx["past_key_values"] = self.bagel.forward_cache_update_vit(ctx["past_key_values"], **inp)
ctx["kv_lens"] = new_kvlens
@@ -1351,6 +1371,12 @@ class LancePipeline(BagelPipeline):
)
with torch.autocast(**autocast_kwargs):
# ``prepare_video_latent`` / ``prepare_vae_latent`` still return
# ``key_values_lens`` / ``packed_indexes`` / ``packed_key_value_indexes``
# but post-main-merge ``generate_image`` doesn't accept them — drop
# before unpacking via ``**gen_input_lat``.
for _drop in ("key_values_lens", "packed_indexes", "packed_key_value_indexes"):
gen_input_lat.pop(_drop, None)
latents, *_ = self.bagel.generate_image(
past_key_values=gen_context["past_key_values"],
cfg_text_past_key_values=cfg_text_context["past_key_values"]
@@ -1366,13 +1392,11 @@ class LancePipeline(BagelPipeline):
cfg_renorm_min=cfg_renorm_min,
**gen_input_lat,
cfg_text_packed_position_ids=cfg_text_lat["cfg_packed_position_ids"],
cfg_text_packed_query_indexes=cfg_text_lat["cfg_packed_query_indexes"],
cfg_text_key_values_lens=cfg_text_lat["cfg_key_values_lens"],
cfg_text_packed_key_value_indexes=cfg_text_lat["cfg_packed_key_value_indexes"],
# ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
# / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
# (derived from ``cfg_text_past_key_values``).
cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
cfg_img_packed_query_indexes=cfg_img_lat["cfg_packed_query_indexes"],
cfg_img_key_values_lens=cfg_img_lat["cfg_key_values_lens"],
cfg_img_packed_key_value_indexes=cfg_img_lat["cfg_packed_key_value_indexes"],
# ``cfg_img_*`` index/lens kwargs removed — same as above.
frame_condition_token_indexes=frame_condition_token_indexes,
)
@@ -1551,17 +1575,17 @@ class LancePipeline(BagelPipeline):
curr_kvlen = ctx["kv_lens"][0]
curr_rope = ctx["ropes"][0]
seq_len = len(text_ids)
# ``forward_cache_update_text`` post-main-merge derives the
# remaining positions from ``past_key_values`` + the query
# length; the old ``packed_text_indexes`` /
# ``packed_key_value_indexes`` / ``key_values_lens`` kwargs
# were dropped from its signature.
inp = {
"text_token_lens": torch.tensor([seq_len], dtype=torch.int, device=self.device),
"packed_text_ids": torch.tensor(text_ids, dtype=torch.long, device=self.device),
"packed_text_position_ids": torch.arange(
curr_rope, curr_rope + seq_len, dtype=torch.long, device=self.device
),
"packed_text_indexes": torch.arange(
curr_kvlen, curr_kvlen + seq_len, dtype=torch.long, device=self.device
),
"packed_key_value_indexes": torch.arange(0, curr_kvlen, dtype=torch.long, device=self.device),
"key_values_lens": torch.tensor([curr_kvlen], dtype=torch.int, device=self.device),
}
with torch.autocast(**autocast_kwargs):
ctx["past_key_values"] = self.bagel.forward_cache_update_text(ctx["past_key_values"], **inp)
@@ -1581,6 +1605,12 @@ class LancePipeline(BagelPipeline):
for k, v in inp.items():
if torch.is_tensor(v):
inp[k] = v.to(self.device)
# ``prepare_vit_images`` still returns ``packed_indexes`` /
# ``packed_key_value_indexes`` / ``key_values_lens`` for
# historical reasons; ``forward_cache_update_vit`` post-main
# no longer accepts them — pop before unpacking.
for _drop in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"):
inp.pop(_drop, None)
with torch.autocast(**autocast_kwargs):
ctx["past_key_values"] = self.bagel.forward_cache_update_vit(ctx["past_key_values"], **inp)
ctx["kv_lens"] = new_kvlens
@@ -1676,6 +1706,12 @@ class LancePipeline(BagelPipeline):
cfg_img_lat = cfg_text_lat
with torch.autocast(**autocast_kwargs):
# ``prepare_video_latent`` / ``prepare_vae_latent`` still return
# ``key_values_lens`` / ``packed_indexes`` / ``packed_key_value_indexes``
# but post-main-merge ``generate_image`` doesn't accept them — drop
# before unpacking via ``**gen_input_lat``.
for _drop in ("key_values_lens", "packed_indexes", "packed_key_value_indexes"):
gen_input_lat.pop(_drop, None)
latents, *_ = self.bagel.generate_image(
past_key_values=gen_context["past_key_values"],
cfg_text_past_key_values=cfg_text_context["past_key_values"]
@@ -1691,13 +1727,11 @@ class LancePipeline(BagelPipeline):
cfg_renorm_min=cfg_renorm_min,
**gen_input_lat,
cfg_text_packed_position_ids=cfg_text_lat["cfg_packed_position_ids"],
cfg_text_packed_query_indexes=cfg_text_lat["cfg_packed_query_indexes"],
cfg_text_key_values_lens=cfg_text_lat["cfg_key_values_lens"],
cfg_text_packed_key_value_indexes=cfg_text_lat["cfg_packed_key_value_indexes"],
# ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
# / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
# (derived from ``cfg_text_past_key_values``).
cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
cfg_img_packed_query_indexes=cfg_img_lat["cfg_packed_query_indexes"],
cfg_img_key_values_lens=cfg_img_lat["cfg_key_values_lens"],
cfg_img_packed_key_value_indexes=cfg_img_lat["cfg_packed_key_value_indexes"],
# ``cfg_img_*`` index/lens kwargs removed — same as above.
)
frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape)
@@ -1777,6 +1811,11 @@ class LancePipeline(BagelPipeline):
for k, v in gen_input_vit.items():
if torch.is_tensor(v):
gen_input_vit[k] = v.to(self.device)
# Post-main-merge ``forward_cache_update_vit`` drops the
# legacy ``packed_indexes`` / ``packed_key_value_indexes`` /
# ``key_values_lens`` kwargs.
for _drop in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"):
gen_input_vit.pop(_drop, None)
with torch.autocast(**autocast_kwargs):
gen_context["past_key_values"] = self.bagel.forward_cache_update_vit(
gen_context["past_key_values"], **gen_input_vit
@@ -1895,6 +1934,11 @@ class LancePipeline(BagelPipeline):
for k, v in gen_input_vit.items():
if torch.is_tensor(v):
gen_input_vit[k] = v.to(self.device)
# Post-main-merge ``forward_cache_update_vit`` drops the
# legacy ``packed_indexes`` / ``packed_key_value_indexes`` /
# ``key_values_lens`` kwargs.
for _drop in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"):
gen_input_vit.pop(_drop, None)
with torch.autocast(**autocast_kwargs):
gen_context["past_key_values"] = self.bagel.forward_cache_update_vit(
gen_context["past_key_values"], **gen_input_vit
@@ -2003,11 +2047,12 @@ class LancePipeline(BagelPipeline):
packed_start_tokens.append(nl_id)
packed_query_position_ids.append(curr_position_id)
curr += curr_kvlen
# ``Bagel.generate_text`` post-main-merge no longer accepts
# ``key_values_lens`` / ``packed_key_value_indexes`` — they're
# derived from ``past_key_values``.
return {
"packed_start_tokens": torch.tensor(packed_start_tokens, dtype=torch.long),
"packed_query_position_ids": torch.tensor(packed_query_position_ids, dtype=torch.long),
"key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int),
"packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long),
}
def _x2t_raw_text_prefill(self, gen_context, text_str, autocast_kwargs):
@@ -2029,9 +2074,9 @@ class LancePipeline(BagelPipeline):
"packed_text_position_ids": torch.arange(
curr_rope, curr_rope + seq_len, dtype=torch.long, device=self.device
),
"packed_text_indexes": torch.arange(curr_kvlen, curr_kvlen + seq_len, dtype=torch.long, device=self.device),
"packed_key_value_indexes": torch.arange(0, curr_kvlen, dtype=torch.long, device=self.device),
"key_values_lens": torch.tensor([curr_kvlen], dtype=torch.int, device=self.device),
# ``packed_text_indexes`` / ``packed_key_value_indexes`` /
# ``key_values_lens`` removed from ``forward_cache_update_text``
# post-main-merge (derived from ``past_key_values``).
}
with torch.autocast(**autocast_kwargs):
gen_context["past_key_values"] = self.bagel.forward_cache_update_text(gen_context["past_key_values"], **inp)