VSA refactoring

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>

Author: o-stoner

tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py

@@ -20,7 +20,14 @@
 simplified metadata that doesn't require KV caching.
 """
 
-from .cute_dsl import CuTeDSLAttention
+from .cute_dsl import (
+    VSA_TILE_SIZE,
+    CuTeDSLAttention,
+    VSAMetadata,
+    VSAMetadataBuilder,
+    get_vsa_forward_context,
+    set_vsa_forward_context,
+)
 from .flash_attn4 import FlashAttn4Attention
 from .interface import AttentionBackend, AttentionTensorLayout
 from .parallel import Attention2DAttention, RingAttention, UlyssesAttention
@@ -41,4 +48,9 @@
     "UlyssesAttention",
     "VanillaAttention",
     "RingAttention",
+    "VSAMetadata",
+    "VSAMetadataBuilder",
+    "VSA_TILE_SIZE",
+    "get_vsa_forward_context",
+    "set_vsa_forward_context",
 ]