Skip to content

vllm_omni.diffusion.attention.backends.trtllm_attn

HAS_FLASHINFER module-attribute 

HAS_FLASHINFER = True

logger module-attribute 

logger = init_logger(__name__)

QuantConfig dataclass

dtype_qk class-attribute instance-attribute 

dtype_qk: str | None = None

enabled property 

enabled: bool

k_block_size class-attribute instance-attribute 

k_block_size: int = 16

q_block_size class-attribute instance-attribute 

q_block_size: int = 1

from_backend_kwargs classmethod 

from_backend_kwargs(
    backend_kwargs: dict | None,
) -> QuantConfig

quantize 

quantize(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    quantize_fn,
    cu_seq_lens_q: Tensor,
    cu_seq_lens_kv: Tensor,
)

SkipSoftmaxConfig dataclass

a class-attribute instance-attribute 

a: float | None = None

b class-attribute instance-attribute 

b: float | None = None

configured property 

configured: bool

disabled_until_timestep class-attribute instance-attribute 

disabled_until_timestep: float = 0.0

enabled property 

enabled: bool

gated property 

gated: bool

target_sparsity class-attribute instance-attribute 

target_sparsity: float | None = None

threshold class-attribute instance-attribute 

threshold: float | None = None

from_backend_kwargs classmethod 

from_backend_kwargs(
    backend_kwargs: dict | None,
) -> SkipSoftmaxConfig

resolve_factor 

resolve_factor(
    seqlen: int, timestep: float | None
) -> float | None

TrtllmAttentionBackend

Bases: AttentionBackend

accept_output_buffer class-attribute instance-attribute 

accept_output_buffer: bool = True

get_impl_cls staticmethod 

get_impl_cls() -> type[TrtllmAttentionImpl]

get_name staticmethod 

get_name() -> str

get_supported_head_sizes staticmethod 

get_supported_head_sizes() -> list[int]

supports_multi_doc_packed_varlen classmethod 

supports_multi_doc_packed_varlen() -> bool

supports_packed_mask_free classmethod 

supports_packed_mask_free() -> bool

TrtllmAttentionImpl

Bases: AttentionImpl

causal instance-attribute 

causal = causal

head_size instance-attribute 

head_size = head_size

num_heads instance-attribute 

num_heads = num_heads

num_kv_heads instance-attribute 

num_kv_heads = (
    num_kv_heads if num_kv_heads is not None else num_heads
)

quant instance-attribute 

quant = QuantConfig(
    dtype_qk=None, q_block_size=0, k_block_size=0
)

role instance-attribute 

role = role

skip instance-attribute 

skip = SkipSoftmaxConfig.from_backend_kwargs(backend_kwargs)

softmax_scale instance-attribute 

softmax_scale = softmax_scale

forward_cuda 

forward_cuda(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    attn_metadata: AttentionMetadata | None = None,
) -> Tensor

set_layer_calibration 

set_layer_calibration(a: float, b: float) -> None