Skip to main content

Ctrl+K

ejkernel documentation

Getting Started:

Installation

Architecture & Design:

ejKernel Project Overview
Kernel Registry System Analysis
Ops System Architecture Analysis
MaskInfo Guide: Comprehensive Attention Mask Management
Kernel Implementations Analysis
Module Operations Layer Analysis
Utilities and Helper Functions Analysis
Test Suite and Examples Analysis
ejKernel Comprehensive Architecture Report

User Guides:

ejkernel API Documentation
ejkernel.modules - High-Level Operations API
ejkernel.types - Masking and Type Utilities
ejkernel.ops - Kernel Framework and Configuration Management

API Reference:

ejkernel API Reference

.rst

ejkernel.kernels._pallas.tpu.flash_attention._pallas_impl_fwd

ejkernel.kernels._pallas.tpu.flash_attention._pallas_impl_fwd#

Flash Attention TPU kernel.

previous

ejkernel.kernels._pallas.tpu.flash_attention._pallas_impl_bwd

next

ejkernel.kernels._pallas.tpu.flash_attention._utils

By Erfan Zare Chavoshi

© Copyright 2023, Erfan Zare Chavoshi - ejkernel.