__post_init__() — pytorch Function Reference

Architecture documentation for the __post_init__() function in generate_kernels.py from the pytorch codebase.

Function python CoreTensor Registration calls 1 called by 1

Entity Profile

CoreTensor→ Registration→ __post_init__() — pytorch Function Reference

Dependency Diagram

graph TD
  cd417a58_19c8_e2b0_9578_fd8e295187a9["__post_init__()"]
  4ef88021_ab7d_96c9_bc9d_caab0cbeafee["__post_init__()"]
  4ef88021_ab7d_96c9_bc9d_caab0cbeafee -->|calls| cd417a58_19c8_e2b0_9578_fd8e295187a9
  4ef88021_ab7d_96c9_bc9d_caab0cbeafee["__post_init__()"]
  cd417a58_19c8_e2b0_9578_fd8e295187a9 -->|calls| 4ef88021_ab7d_96c9_bc9d_caab0cbeafee
  style cd417a58_19c8_e2b0_9578_fd8e295187a9 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/generate_kernels.py lines 160–175

    def __post_init__(self) -> None:
        # Set kernel selection priority
        # The lowest value that matches inputs
        # will be selected
        self.sort_index = (
            # First select aligned kernel
            0 if self.aligned else 1,
            # Take a kernel without dropout if possible
            1 if self.apply_dropout else 0,
            # Then take the smallest maxK
            self.max_k,
            # .. and the highest block_i
            -self.block_i,
            # and finally avoid bounds-checks if possible
            0 if self.keys_queries_aligned_to_blocksizes else 1,
        )