vllm.v1.attention.backends.tpu_v6_adaptive_pallas

class TPUArchitectureDetector:
    """
    Detects TPU version and provides optimization configuration.
    Falls back gracefully when running on CPU/GPU for development.
    """

    # Known TPU configurations based on public documentation
    TPU_CONFIGS = {
        6:
        TPUConfig(version=6,
                  name="TPU v6e (Trillium)",
                  mxu_size=256,
                  memory_bandwidth_gbps=3584,
                  ici_bandwidth_gbps=3584,
                  sparse_cores=2,
                  head_size_multiple=256,
                  optimal_block_q=512,
                  optimal_block_kv=1024,
                  memory_pipeline_stages=4),
        5:
        TPUConfig(version=5,
                  name="TPU v5e",
                  mxu_size=128,
                  memory_bandwidth_gbps=1600,
                  ici_bandwidth_gbps=1600,
                  sparse_cores=4,
                  head_size_multiple=128,
                  optimal_block_q=256,
                  optimal_block_kv=512,
                  memory_pipeline_stages=2),
        4:
        TPUConfig(version=4,
                  name="TPU v4",
                  mxu_size=128,
                  memory_bandwidth_gbps=1200,
                  ici_bandwidth_gbps=1200,
                  sparse_cores=0,
                  head_size_multiple=128,
                  optimal_block_q=256,
                  optimal_block_kv=512,
                  memory_pipeline_stages=2)
    }

    def __init__(self):
        self.tpu_version = self._detect_tpu_version()
        self.config = self._get_config()
        self.is_simulated = self.tpu_version == -1

        if self.is_simulated:
            logger.warning("Running in simulation mode - no TPU detected")
        else:
            logger.info(f"Detected {self.config.name}")

    def _detect_tpu_version(self) -> int:
        """Detect TPU version from environment"""
        # Method 1: PyTorch XLA
        try:
            import torch_xla
            version = torch_xla.tpu.version()
            logger.info(f"Detected TPU v{version} via torch_xla")
            return version
        except (ImportError, AttributeError):
            pass

        # Method 2: JAX
        try:
            import jax
            devices = jax.devices()
            if devices and 'TPU' in str(devices[0]):
                # Parse version from device string
                device_str = str(devices[0])
                if 'v6' in device_str:
                    return 6
                elif 'v5' in device_str:
                    return 5
                elif 'v4' in device_str:
                    return 4
        except (ImportError, AttributeError, IndexError):
            pass

        # Method 3: Environment variable (for testing)
        env_version = os.environ.get('TPU_VERSION', None)
        if env_version:
            logger.info(f"Using TPU v{env_version} from environment")
            return int(env_version)

        # No TPU detected - simulation mode
        return -1

    def _get_config(self) -> TPUConfig:
        """Get configuration for detected TPU version"""
        if self.tpu_version in self.TPU_CONFIGS:
            return self.TPU_CONFIGS[self.tpu_version]
        elif self.tpu_version == -1:
            # Simulation mode - default to v6 config
            logger.info("Using TPU v6e configuration for simulation")
            return self.TPU_CONFIGS[6]
        else:
            # Unknown version - use v5 as safe default
            logger.warning(f"Unknown TPU v{self.tpu_version}, using v5 config")
            return self.TPU_CONFIGS[5]

    def optimize_head_dimension(self, head_dim: int) -> int:
        """Optimize head dimension for MXU alignment"""
        multiple = self.config.head_size_multiple
        optimized = ((head_dim + multiple - 1) // multiple) * multiple

        if optimized != head_dim:
            logger.info(
                f"Optimizing head dimension: {head_dim} -> {optimized}")

        return optimized

    def get_attention_config(self, seq_len: int) -> Dict[str, Any]:
        """Get optimized attention configuration"""
        return {
            "block_q": min(self.config.optimal_block_q, seq_len),
            "block_kv": min(self.config.optimal_block_kv, seq_len),
            "memory_pipeline_stages": self.config.memory_pipeline_stages,
            "mxu_size": self.config.mxu_size,
            "is_v6_optimized": self.config.version >= 6
        }

@dataclass
class TPUConfig:
    """TPU architecture configuration for optimization"""
    version: int
    name: str
    mxu_size: int
    memory_bandwidth_gbps: float
    ici_bandwidth_gbps: float
    sparse_cores: int
    head_size_multiple: int
    optimal_block_q: int
    optimal_block_kv: int
    memory_pipeline_stages: int

class TPUv6AdaptiveAttentionBackend(PallasAttentionBackend):
    """
    TPU v6e adaptive attention backend that extends the base PallasAttentionBackend
    with architecture-specific optimizations.
    """

    @staticmethod
    def get_name() -> str:
        return "TPU_V6E_ADAPTIVE_PALLAS_VLLM_V1"

    @staticmethod
    def get_impl_cls() -> type["TPUv6AdaptiveAttentionBackendImpl"]:
        return TPUv6AdaptiveAttentionBackendImpl

    @staticmethod
    def get_kv_cache_shape(
        num_blocks: int,
        block_size: int,
        num_kv_heads: int,
        head_size: int,
    ) -> tuple[int, ...]:
        # Use architecture-adaptive head size alignment
        alignment = tpu_detector.config.head_size_multiple
        padded_head_size = cdiv(head_size, alignment) * alignment
        return (num_blocks, block_size, num_kv_heads * 2, padded_head_size)

    @staticmethod
    def get_page_size(vllm_config) -> int:
        """Get optimized page size for TPU architecture"""
        # For TPU v6e with larger memory bandwidth, we can use larger page sizes
        if tpu_detector.config.version >= 6:
            # Use larger page sizes for better memory pipeline utilization
            if vllm_config.model_config.max_model_len > 8192:
                return 32  # Doubled from original 16
            page_size = next_power_of_2(
                vllm_config.model_config.max_model_len) // 8  # Reduced divisor
            if page_size <= 32:
                return 32
            if page_size >= 512:
                return 512
            return page_size
        else:
            # Use original logic for v5e and earlier
            return super().get_page_size(vllm_config)

class TPUv6AdaptiveAttentionBackendImpl(AttentionImpl):
    """
    TPU v6e adaptive attention implementation with architecture-specific optimizations.
    """

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: int,
        alibi_slopes: Optional[list[float]],
        sliding_window: Optional[int],
        kv_cache_dtype: str,
        logits_soft_cap: Optional[float] = None,
        attn_type: str = AttentionType.DECODER,
        kv_sharing_target_layer_name: Optional[int] = None,
    ) -> None:

        # Store original parameters
        self.num_heads = num_heads
        self.original_head_size = head_size
        self.scale = float(scale)
        self.num_kv_heads = num_kv_heads
        self.sliding_window = sliding_window
        self.logits_soft_cap = logits_soft_cap
        self.kv_sharing_target_layer_name = kv_sharing_target_layer_name

        # Optimize head size for TPU architecture
        self.head_size = tpu_detector.optimize_head_dimension(head_size)
        self.attention_config = tpu_detector.get_attention_config(
            4096)  # Default seq len

        # Performance tracking
        self.call_count = 0
        self.total_optimization_time = 0.0

        self.num_queries_per_kv = self.num_heads // self.num_kv_heads
        if alibi_slopes is not None:
            raise NotImplementedError("Alibi slopes is not supported.")

        if attn_type != AttentionType.DECODER:
            raise NotImplementedError("Encoder self-attention and "
                                      "encoder/decoder cross-attention "
                                      "are not implemented for "
                                      "TPUv6AdaptiveAttentionBackendImpl")

        self.kv_cache_quantized_dtype = None
        if kv_cache_dtype != "auto":
            self.kv_cache_quantized_dtype = TPU_STR_DTYPE_TO_TORCH_DTYPE.get(
                kv_cache_dtype.lower().strip())

        # Log optimization information
        logger.info(f"Initialized TPU v6e Adaptive Attention Backend")
        logger.info(f"  Architecture: {tpu_detector.config.name}")
        logger.info(
            f"  Head size optimization: {self.original_head_size} -> {self.head_size}"
        )
        logger.info(
            f"  MXU target: {tpu_detector.config.mxu_size}x{tpu_detector.config.mxu_size}"
        )
        logger.info(
            f"  Memory pipeline: {self.attention_config['memory_pipeline_stages']} stages"
        )

    def forward(
        self,
        layer: AttentionLayer,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: PallasMetadata,
        output: Optional[torch.Tensor] = None,
        output_scale: Optional[torch.Tensor] = None,
        output_block_scale: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Forward pass with TPU v6e optimizations."""

        import time
        start_time = time.perf_counter()

        if output_scale is not None or output_block_scale is not None:
            raise NotImplementedError(
                "fused output quantization is not yet supported"
                " for TPUv6AdaptiveAttentionBackendImpl")

        # For determine_available_memory case.
        if kv_cache.numel() == 0:
            if output is None:
                output = torch.ones_like(query)
            return output

        num_tokens, hidden_size = query.shape
        query = query.view(num_tokens, self.num_heads, self.head_size)
        key = key.view(-1, self.num_kv_heads, self.head_size)
        value = value.view(-1, self.num_kv_heads, self.head_size)

        # TPU v6e adaptive padding with architecture-specific alignment
        alignment = tpu_detector.config.head_size_multiple
        if self.head_size % alignment != 0:
            padded_head_size = cdiv(self.head_size, alignment) * alignment
            query = torch.nn.functional.pad(
                query, (0, padded_head_size - self.head_size), value=0.0)
            key = torch.nn.functional.pad(
                key, (0, padded_head_size - self.head_size), value=0.0)
            value = torch.nn.functional.pad(
                value, (0, padded_head_size - self.head_size), value=0.0)

        if self.kv_sharing_target_layer_name is None and kv_cache.numel() > 0:
            # Write input keys and values to the KV cache with v6e optimization
            slot_mapping = attn_metadata.slot_mapping
            write_to_kv_cache(
                key,
                value,
                kv_cache,
                slot_mapping,
                attn_metadata.num_slices_per_kv_cache_update_block,
                attn_metadata.num_kv_update_slices,
                self.kv_cache_quantized_dtype,
                layer._k_scale_float,
                layer._v_scale_float,
            )

        if self.kv_cache_quantized_dtype is not None and (
                layer._k_scale_float == 0.0 or layer._v_scale_float == 0.0):
            raise ValueError(
                "k_scale_float and v_scale_float must be non-zero")

        # TPU v6e optimized attention with architecture-adaptive parameters
        if tpu_detector.config.version >= 6:
            # Use v6e optimizations - larger blocks and memory pipeline depth
            num_kv_pages_per_block = min(
                4, max(1, self.attention_config["block_kv"] // 128))
            num_queries_per_block = min(
                8, max(1, self.attention_config["block_q"] // 64))
            # Increased vmem limit for v6e's larger memory bandwidth
            vmem_limit_bytes = min(1024 * 1024, 768 * 1024)  # 768KB for v6e
        else:
            # Use v5e defaults
            num_kv_pages_per_block = None
            num_queries_per_block = None
            vmem_limit_bytes = None

        output = torch.ops.xla.ragged_paged_attention(
            query,
            kv_cache,
            attn_metadata.context_lens,
            attn_metadata.block_tables,
            attn_metadata.query_start_loc,
            attn_metadata.num_seqs,
            num_kv_pages_per_block=num_kv_pages_per_block,
            num_queries_per_block=num_queries_per_block,
            vmem_limit_bytes=vmem_limit_bytes,
            use_kernel=True,
            sm_scale=self.scale,
            sliding_window=self.sliding_window,
            soft_cap=self.logits_soft_cap,
            k_scale=layer._k_scale_float,
            v_scale=layer._v_scale_float,
        )

        # Remove padding for output
        if self.head_size % alignment != 0:
            output = output[:, :, :self.head_size]

        # Performance tracking
        end_time = time.perf_counter()
        self.call_count += 1
        self.total_optimization_time += (end_time - start_time)

        # Log performance periodically
        if self.call_count % 100 == 0:
            avg_time = self.total_optimization_time / self.call_count * 1000
            logger.info(f"TPU v6e Adaptive: {self.call_count} calls, "
                        f"avg time: {avg_time:.2f}ms, "
                        f"architecture: {tpu_detector.config.name}")

        return output.reshape(num_tokens, hidden_size)

    def get_performance_report(self) -> Dict[str, Any]:
        """Generate performance report for monitoring"""
        return {
            "backend":
            "TPUv6AdaptiveAttentionBackend",
            "architecture":
            tpu_detector.config.name,
            "tpu_version":
            tpu_detector.config.version,
            "calls":
            self.call_count,
            "mxu_size":
            f"{tpu_detector.config.mxu_size}x{tpu_detector.config.mxu_size}",
            "head_size_optimization":
            f"{self.original_head_size} -> {self.head_size}",
            "memory_pipeline_stages":
            self.attention_config["memory_pipeline_stages"],
            "is_v6_optimized":
            self.attention_config["is_v6_optimized"],
            "average_call_time_ms":
            (self.total_optimization_time / max(1, self.call_count)) * 1000,
            "optimizations_applied":
            self._get_applied_optimizations()
        }

    def _get_applied_optimizations(self) -> list[str]:
        """Get list of applied optimizations"""
        optimizations = []
        if tpu_detector.config.version >= 6:
            optimizations.extend([
                "mxu_256x256_alignment", "4_stage_memory_pipeline",
                "enhanced_vmem_limits", "optimized_block_sizing"
            ])
        else:
            optimizations.extend([
                "mxu_128x128_alignment", "2_stage_memory_pipeline",
                "standard_block_sizing"
            ])

        if self.head_size != self.original_head_size:
            optimizations.append("head_dimension_padding")

        return optimizations