Skip to content

GitLab

Unverified Commit ac0bb2c3 authored by Lukas Geiger's avatar Lukas Geiger Committed by GitHub
Browse files

[Core] Cache `vllm_is_batch_invariant` (#28304) 


Signed-off-by: default avatarLukas Geiger <lukas.geiger94@gmail.com>
parent f31419ed
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
vllm/model_executor/layers/batch_invariant.py
View file @ ac0bb2c3
...@@ -4,6 +4,7 @@ import contextlib ...@@ -4,6 +4,7 @@ import contextlib
import os import os
from collections import namedtuple from collections import namedtuple
from collections.abc import Callable from collections.abc import Callable
from functools import cache
from typing import Any from typing import Any
import torch import torch
...@@ -857,6 +858,7 @@ def get_batch_invariant_attention_block_size() -> AttentionBlockSize: ...@@ -857,6 +858,7 @@ def get_batch_invariant_attention_block_size() -> AttentionBlockSize:
return AttentionBlockSize(block_m=16, block_n=16) return AttentionBlockSize(block_m=16, block_n=16)
@cache
def vllm_is_batch_invariant(): def vllm_is_batch_invariant():
env_key = "VLLM_BATCH_INVARIANT" env_key = "VLLM_BATCH_INVARIANT"
is_overridden = False is_overridden = False
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment