Triton TLX-based block attention kernels for variable-length attention workloads.
triton/- Triton TLX kernel implementationtests/- project tests
from ads_mkl.ops.oss.block_attention.triton.tlx_block_attention import (
block_attention_api,
)buck2 test fbcode//ads_mkl/ops/oss/block_attention:test_tlx_block_attentionThis test target includes:
fbcode/ads_mkl/ops/oss/block_attention/triton/tests/test_tlx_block_attention.pyfbcode/ads_mkl/ops/oss/block_attention/tests/block_attention_test.py