refactor: simplify inference engine and backend dispatch
- merge _generate_streaming/_generate_non_streaming into single _generate() with stream flag - delete dead GenerationRequest class and generate_with_request method - inline _next_token helper into generate_async - replace flash-attn double-checked locking with functools.lru_cache - extract _write_and_gather_kv helper shared by TorchNative/FlashAttn backends - inline _kv_cache_is_contiguous into its sole call site in FlashAttnBackend - change default backend priority from flash>cuda>torch to cuda>flash>torch - add ASTR_BACKEND env var to override default backend at resolve time - add supports_graph() static method to AttentionBackend ABC, override in CudaBackend - replace isinstance(get_backend(), CudaBackend) with get_backend().supports_graph() in executor - add torch.cuda.is_available() guard to CudaBackend.supports()
This commit is contained in:
@@ -179,9 +179,7 @@ class Executor:
|
||||
max_q_heads = config.num_attention_heads
|
||||
head_dim = config.hidden_size // config.num_attention_heads
|
||||
self._head_dim = head_dim
|
||||
self._graph_supported = CudaBackend.supports(
|
||||
head_dim=head_dim
|
||||
) and "cuda" in str(self.device)
|
||||
self._graph_supported = CudaBackend.supports(head_dim=head_dim)
|
||||
self._workspace = InferenceWorkspace(
|
||||
max_batch_size=kv_cache.max_batch_size,
|
||||
max_seq_len=kv_cache.max_seq_len,
|
||||
@@ -367,7 +365,7 @@ class Executor:
|
||||
use_graph = (
|
||||
self._graph_ctx.enabled
|
||||
and self._graph_supported
|
||||
and isinstance(get_backend(), CudaBackend)
|
||||
and get_backend().supports_graph()
|
||||
)
|
||||
key = (b,)
|
||||
if use_graph:
|
||||
|
||||
Reference in New Issue
Block a user