hyx21 commited on
Commit
20e9f30
·
verified ·
1 Parent(s): 97665a1

Update modeling_llama_long_infllmv2.py

Browse files
Files changed (1) hide show
  1. modeling_llama_long_infllmv2.py +0 -3
modeling_llama_long_infllmv2.py CHANGED
@@ -50,10 +50,8 @@ from transformers.utils import (
50
  replace_return_docstrings,
51
  )
52
  from transformers.models.llama.configuration_llama import LlamaConfig
53
- from moba import moba_attn_varlen
54
  from functools import lru_cache
55
  from .cis_pooling import nosa_mean_pooling
56
- from native_sparse_attention.ops.triton.topk_sparse_attention import topk_sparse_attention
57
 
58
  logger = logging.get_logger(__name__)
59
 
@@ -680,7 +678,6 @@ class LlamaFlashAttention2(LlamaAttention):
680
 
681
  return attn_output, attn_weights, past_key_value
682
 
683
- from native_sparse_attention.ops.triton.topk_sparse_attention import topk_sparse_attention
684
  try:
685
  from flash_attn import flash_attn_func, flash_attn_varlen_func
686
  from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
 
50
  replace_return_docstrings,
51
  )
52
  from transformers.models.llama.configuration_llama import LlamaConfig
 
53
  from functools import lru_cache
54
  from .cis_pooling import nosa_mean_pooling
 
55
 
56
  logger = logging.get_logger(__name__)
57
 
 
678
 
679
  return attn_output, attn_weights, past_key_value
680
 
 
681
  try:
682
  from flash_attn import flash_attn_func, flash_attn_varlen_func
683
  from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa