fix: support HybridLinearAttnBackend in TboAttnBackend (#20114)
This commit is contained in:
@@ -701,12 +701,20 @@ class TboForwardBatchPreparer:
|
|||||||
"spec_algorithm",
|
"spec_algorithm",
|
||||||
"capture_hidden_mode",
|
"capture_hidden_mode",
|
||||||
"padded_static_len",
|
"padded_static_len",
|
||||||
"mrope_positions", # only used by qwen2-vl, thus not care
|
|
||||||
"split_index", # for split prefill
|
"split_index", # for split prefill
|
||||||
"orig_seq_lens", # only used by qwen-1m, thus not care
|
"orig_seq_lens", # only used by qwen-1m, thus not care
|
||||||
"return_pooled_hidden_states",
|
"return_pooled_hidden_states",
|
||||||
]:
|
]:
|
||||||
output_dict[key] = getattr(batch, key)
|
output_dict[key] = getattr(batch, key)
|
||||||
|
|
||||||
|
mrope_positions = getattr(batch, "mrope_positions")
|
||||||
|
if mrope_positions is not None:
|
||||||
|
output_dict["mrope_positions"] = mrope_positions[
|
||||||
|
:, start_token_index:end_token_index
|
||||||
|
]
|
||||||
|
else:
|
||||||
|
output_dict["mrope_positions"] = None
|
||||||
|
|
||||||
if not batch.forward_mode.is_target_verify():
|
if not batch.forward_mode.is_target_verify():
|
||||||
assert (
|
assert (
|
||||||
_compute_extend_num_tokens(batch.input_ids, batch.forward_mode)
|
_compute_extend_num_tokens(batch.input_ids, batch.forward_mode)
|
||||||
|
|||||||
@@ -179,6 +179,9 @@ class TboAttnBackend(AttentionBackend):
|
|||||||
assert ans == child.get_cuda_graph_seq_len_fill_value()
|
assert ans == child.get_cuda_graph_seq_len_fill_value()
|
||||||
return ans
|
return ans
|
||||||
|
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
return self.primary.forward(*args, **kwargs)
|
||||||
|
|
||||||
def forward_extend(self, *args, **kwargs):
|
def forward_extend(self, *args, **kwargs):
|
||||||
return self.primary.forward_extend(*args, **kwargs)
|
return self.primary.forward_extend(*args, **kwargs)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user