Add files using upload-large-folder tool

lib/python3.12/site-packages/matplotlib/_docstring.pyi

@@ -0,0 +1,34 @@
+from collections.abc import Callable
+from typing import Any, TypeVar, overload
+
+
+_T = TypeVar('_T')
+
+
+def kwarg_doc(text: str) -> Callable[[_T], _T]: ...
+
+
+class Substitution:
+    @overload
+    def __init__(self, *args: str): ...
+    @overload
+    def __init__(self, **kwargs: str): ...
+    def __call__(self, func: _T) -> _T: ...
+    def update(self, *args, **kwargs): ...  # type: ignore[no-untyped-def]
+
+
+class _ArtistKwdocLoader(dict[str, str]):
+    def __missing__(self, key: str) -> str: ...
+
+
+class _ArtistPropertiesSubstitution:
+    def __init__(self) -> None: ...
+    def register(self, **kwargs) -> None: ...
+    def __call__(self, obj: _T) -> _T: ...
+
+
+def copy(source: Any) -> Callable[[_T], _T]: ...
+
+
+dedent_interpd: _ArtistPropertiesSubstitution
+interpd: _ArtistPropertiesSubstitution

lib/python3.12/site-packages/matplotlib/_enums.py

@@ -0,0 +1,177 @@
+"""
+Enums representing sets of strings that Matplotlib uses as input parameters.
+
+Matplotlib often uses simple data types like strings or tuples to define a
+concept; e.g. the line capstyle can be specified as one of 'butt', 'round',
+or 'projecting'. The classes in this module are used internally and serve to
+document these concepts formally.
+
+As an end-user you will not use these classes directly, but only the values
+they define.
+"""
+
+from enum import Enum
+from matplotlib import _docstring
+
+
+class JoinStyle(str, Enum):
+    """
+    Define how the connection between two line segments is drawn.
+
+    For a visual impression of each *JoinStyle*, `view these docs online
+    <JoinStyle>`, or run `JoinStyle.demo`.
+
+    Lines in Matplotlib are typically defined by a 1D `~.path.Path` and a
+    finite ``linewidth``, where the underlying 1D `~.path.Path` represents the
+    center of the stroked line.
+
+    By default, `~.backend_bases.GraphicsContextBase` defines the boundaries of
+    a stroked line to simply be every point within some radius,
+    ``linewidth/2``, away from any point of the center line. However, this
+    results in corners appearing "rounded", which may not be the desired
+    behavior if you are drawing, for example, a polygon or pointed star.
+
+    **Supported values:**
+
+    .. rst-class:: value-list
+
+        'miter'
+            the "arrow-tip" style. Each boundary of the filled-in area will
+            extend in a straight line parallel to the tangent vector of the
+            centerline at the point it meets the corner, until they meet in a
+            sharp point.
+        'round'
+            stokes every point within a radius of ``linewidth/2`` of the center
+            lines.
+        'bevel'
+            the "squared-off" style. It can be thought of as a rounded corner
+            where the "circular" part of the corner has been cut off.
+
+    .. note::
+
+        Very long miter tips are cut off (to form a *bevel*) after a
+        backend-dependent limit called the "miter limit", which specifies the
+        maximum allowed ratio of miter length to line width. For example, the
+        PDF backend uses the default value of 10 specified by the PDF standard,
+        while the SVG backend does not even specify the miter limit, resulting
+        in a default value of 4 per the SVG specification. Matplotlib does not
+        currently allow the user to adjust this parameter.
+
+        A more detailed description of the effect of a miter limit can be found
+        in the `Mozilla Developer Docs
+        <https://developer.mozilla.org/en-US/docs/Web/SVG/Attribute/stroke-miterlimit>`_
+
+    .. plot::
+        :alt: Demo of possible JoinStyle's
+
+        from matplotlib._enums import JoinStyle
+        JoinStyle.demo()
+
+    """
+
+    miter = "miter"
+    round = "round"
+    bevel = "bevel"
+
+    @staticmethod
+    def demo():
+        """Demonstrate how each JoinStyle looks for various join angles."""
+        import numpy as np
+        import matplotlib.pyplot as plt
+
+        def plot_angle(ax, x, y, angle, style):
+            phi = np.radians(angle)
+            xx = [x + .5, x, x + .5*np.cos(phi)]
+            yy = [y, y, y + .5*np.sin(phi)]
+            ax.plot(xx, yy, lw=12, color='tab:blue', solid_joinstyle=style)
+            ax.plot(xx, yy, lw=1, color='black')
+            ax.plot(xx[1], yy[1], 'o', color='tab:red', markersize=3)
+
+        fig, ax = plt.subplots(figsize=(5, 4), constrained_layout=True)
+        ax.set_title('Join style')
+        for x, style in enumerate(['miter', 'round', 'bevel']):
+            ax.text(x, 5, style)
+            for y, angle in enumerate([20, 45, 60, 90, 120]):
+                plot_angle(ax, x, y, angle, style)
+                if x == 0:
+                    ax.text(-1.3, y, f'{angle} degrees')
+        ax.set_xlim(-1.5, 2.75)
+        ax.set_ylim(-.5, 5.5)
+        ax.set_axis_off()
+        fig.show()
+
+
+JoinStyle.input_description = "{" \
+        + ", ".join([f"'{js.name}'" for js in JoinStyle]) \
+        + "}"
+
+
+class CapStyle(str, Enum):
+    r"""
+    Define how the two endpoints (caps) of an unclosed line are drawn.
+
+    How to draw the start and end points of lines that represent a closed curve
+    (i.e. that end in a `~.path.Path.CLOSEPOLY`) is controlled by the line's
+    `JoinStyle`. For all other lines, how the start and end points are drawn is
+    controlled by the *CapStyle*.
+
+    For a visual impression of each *CapStyle*, `view these docs online
+    <CapStyle>` or run `CapStyle.demo`.
+
+    By default, `~.backend_bases.GraphicsContextBase` draws a stroked line as
+    squared off at its endpoints.
+
+    **Supported values:**
+
+    .. rst-class:: value-list
+
+        'butt'
+            the line is squared off at its endpoint.
+        'projecting'
+            the line is squared off as in *butt*, but the filled in area
+            extends beyond the endpoint a distance of ``linewidth/2``.
+        'round'
+            like *butt*, but a semicircular cap is added to the end of the
+            line, of radius ``linewidth/2``.
+
+    .. plot::
+        :alt: Demo of possible CapStyle's
+
+        from matplotlib._enums import CapStyle
+        CapStyle.demo()
+
+    """
+    butt = "butt"
+    projecting = "projecting"
+    round = "round"
+
+    @staticmethod
+    def demo():
+        """Demonstrate how each CapStyle looks for a thick line segment."""
+        import matplotlib.pyplot as plt
+
+        fig = plt.figure(figsize=(4, 1.2))
+        ax = fig.add_axes([0, 0, 1, 0.8])
+        ax.set_title('Cap style')
+
+        for x, style in enumerate(['butt', 'round', 'projecting']):
+            ax.text(x+0.25, 0.85, style, ha='center')
+            xx = [x, x+0.5]
+            yy = [0, 0]
+            ax.plot(xx, yy, lw=12, color='tab:blue', solid_capstyle=style)
+            ax.plot(xx, yy, lw=1, color='black')
+            ax.plot(xx, yy, 'o', color='tab:red', markersize=3)
+
+        ax.set_ylim(-.5, 1.5)
+        ax.set_axis_off()
+        fig.show()
+
+
+CapStyle.input_description = "{" \
+        + ", ".join([f"'{cs.name}'" for cs in CapStyle]) \
+        + "}"
+
+_docstring.interpd.register(
+    JoinStyle=JoinStyle.input_description,
+    CapStyle=CapStyle.input_description,
+)

lib/python3.12/site-packages/matplotlib/_type1font.py

@@ -0,0 +1,879 @@
+"""
+A class representing a Type 1 font.
+
+This version reads pfa and pfb files and splits them for embedding in
+pdf files. It also supports SlantFont and ExtendFont transformations,
+similarly to pdfTeX and friends. There is no support yet for subsetting.
+
+Usage::
+
+    font = Type1Font(filename)
+    clear_part, encrypted_part, finale = font.parts
+    slanted_font = font.transform({'slant': 0.167})
+    extended_font = font.transform({'extend': 1.2})
+
+Sources:
+
+* Adobe Technical Note #5040, Supporting Downloadable PostScript
+  Language Fonts.
+
+* Adobe Type 1 Font Format, Adobe Systems Incorporated, third printing,
+  v1.1, 1993. ISBN 0-201-57044-0.
+"""
+
+from __future__ import annotations
+
+import binascii
+import functools
+import logging
+import re
+import string
+import struct
+import typing as T
+
+import numpy as np
+
+from matplotlib.cbook import _format_approx
+from . import _api
+
+_log = logging.getLogger(__name__)
+
+
+class _Token:
+    """
+    A token in a PostScript stream.
+
+    Attributes
+    ----------
+    pos : int
+        Position, i.e. offset from the beginning of the data.
+    raw : str
+        Raw text of the token.
+    kind : str
+        Description of the token (for debugging or testing).
+    """
+    __slots__ = ('pos', 'raw')
+    kind = '?'
+
+    def __init__(self, pos, raw):
+        _log.debug('type1font._Token %s at %d: %r', self.kind, pos, raw)
+        self.pos = pos
+        self.raw = raw
+
+    def __str__(self):
+        return f"<{self.kind} {self.raw} @{self.pos}>"
+
+    def endpos(self):
+        """Position one past the end of the token"""
+        return self.pos + len(self.raw)
+
+    def is_keyword(self, *names):
+        """Is this a name token with one of the names?"""
+        return False
+
+    def is_slash_name(self):
+        """Is this a name token that starts with a slash?"""
+        return False
+
+    def is_delim(self):
+        """Is this a delimiter token?"""
+        return False
+
+    def is_number(self):
+        """Is this a number token?"""
+        return False
+
+    def value(self):
+        return self.raw
+
+
+class _NameToken(_Token):
+    kind = 'name'
+
+    def is_slash_name(self):
+        return self.raw.startswith('/')
+
+    def value(self):
+        return self.raw[1:]
+
+
+class _BooleanToken(_Token):
+    kind = 'boolean'
+
+    def value(self):
+        return self.raw == 'true'
+
+
+class _KeywordToken(_Token):
+    kind = 'keyword'
+
+    def is_keyword(self, *names):
+        return self.raw in names
+
+
+class _DelimiterToken(_Token):
+    kind = 'delimiter'
+
+    def is_delim(self):
+        return True
+
+    def opposite(self):
+        return {'[': ']', ']': '[',
+                '{': '}', '}': '{',
+                '<<': '>>', '>>': '<<'
+                }[self.raw]
+
+
+class _WhitespaceToken(_Token):
+    kind = 'whitespace'
+
+
+class _StringToken(_Token):
+    kind = 'string'
+    _escapes_re = re.compile(r'\\([\\()nrtbf]|[0-7]{1,3})')
+    _replacements = {'\\': '\\', '(': '(', ')': ')', 'n': '\n',
+                     'r': '\r', 't': '\t', 'b': '\b', 'f': '\f'}
+    _ws_re = re.compile('[\0\t\r\f\n ]')
+
+    @classmethod
+    def _escape(cls, match):
+        group = match.group(1)
+        try:
+            return cls._replacements[group]
+        except KeyError:
+            return chr(int(group, 8))
+
+    @functools.lru_cache
+    def value(self):
+        if self.raw[0] == '(':
+            return self._escapes_re.sub(self._escape, self.raw[1:-1])
+        else:
+            data = self._ws_re.sub('', self.raw[1:-1])
+            if len(data) % 2 == 1:
+                data += '0'
+            return binascii.unhexlify(data)
+
+
+class _BinaryToken(_Token):
+    kind = 'binary'
+
+    def value(self):
+        return self.raw[1:]
+
+
+class _NumberToken(_Token):
+    kind = 'number'
+
+    def is_number(self):
+        return True
+
+    def value(self):
+        if '.' not in self.raw:
+            return int(self.raw)
+        else:
+            return float(self.raw)
+
+
+def _tokenize(data: bytes, skip_ws: bool) -> T.Generator[_Token, int, None]:
+    """
+    A generator that produces _Token instances from Type-1 font code.
+
+    The consumer of the generator may send an integer to the tokenizer to
+    indicate that the next token should be _BinaryToken of the given length.
+
+    Parameters
+    ----------
+    data : bytes
+        The data of the font to tokenize.
+
+    skip_ws : bool
+        If true, the generator will drop any _WhitespaceTokens from the output.
+    """
+
+    text = data.decode('ascii', 'replace')
+    whitespace_or_comment_re = re.compile(r'[\0\t\r\f\n ]+|%[^\r\n]*')
+    token_re = re.compile(r'/{0,2}[^]\0\t\r\f\n ()<>{}/%[]+')
+    instring_re = re.compile(r'[()\\]')
+    hex_re = re.compile(r'^<[0-9a-fA-F\0\t\r\f\n ]*>$')
+    oct_re = re.compile(r'[0-7]{1,3}')
+    pos = 0
+    next_binary: int | None = None
+
+    while pos < len(text):
+        if next_binary is not None:
+            n = next_binary
+            next_binary = (yield _BinaryToken(pos, data[pos:pos+n]))
+            pos += n
+            continue
+        match = whitespace_or_comment_re.match(text, pos)
+        if match:
+            if not skip_ws:
+                next_binary = (yield _WhitespaceToken(pos, match.group()))
+            pos = match.end()
+        elif text[pos] == '(':
+            # PostScript string rules:
+            # - parentheses must be balanced
+            # - backslashes escape backslashes and parens
+            # - also codes \n\r\t\b\f and octal escapes are recognized
+            # - other backslashes do not escape anything
+            start = pos
+            pos += 1
+            depth = 1
+            while depth:
+                match = instring_re.search(text, pos)
+                if match is None:
+                    raise ValueError(
+                        f'Unterminated string starting at {start}')
+                pos = match.end()
+                if match.group() == '(':
+                    depth += 1
+                elif match.group() == ')':
+                    depth -= 1
+                else:  # a backslash
+                    char = text[pos]
+                    if char in r'\()nrtbf':
+                        pos += 1
+                    else:
+                        octal = oct_re.match(text, pos)
+                        if octal:
+                            pos = octal.end()
+                        else:
+                            pass  # non-escaping backslash
+            next_binary = (yield _StringToken(start, text[start:pos]))
+        elif text[pos:pos + 2] in ('<<', '>>'):
+            next_binary = (yield _DelimiterToken(pos, text[pos:pos + 2]))
+            pos += 2
+        elif text[pos] == '<':
+            start = pos
+            try:
+                pos = text.index('>', pos) + 1
+            except ValueError as e:
+                raise ValueError(f'Unterminated hex string starting at {start}'
+                                 ) from e
+            if not hex_re.match(text[start:pos]):
+                raise ValueError(f'Malformed hex string starting at {start}')
+            next_binary = (yield _StringToken(pos, text[start:pos]))
+        else:
+            match = token_re.match(text, pos)
+            if match:
+                raw = match.group()
+                if raw.startswith('/'):
+                    next_binary = (yield _NameToken(pos, raw))
+                elif match.group() in ('true', 'false'):
+                    next_binary = (yield _BooleanToken(pos, raw))
+                else:
+                    try:
+                        float(raw)
+                        next_binary = (yield _NumberToken(pos, raw))
+                    except ValueError:
+                        next_binary = (yield _KeywordToken(pos, raw))
+                pos = match.end()
+            else:
+                next_binary = (yield _DelimiterToken(pos, text[pos]))
+                pos += 1
+
+
+class _BalancedExpression(_Token):
+    pass
+
+
+def _expression(initial, tokens, data):
+    """
+    Consume some number of tokens and return a balanced PostScript expression.
+
+    Parameters
+    ----------
+    initial : _Token
+        The token that triggered parsing a balanced expression.
+    tokens : iterator of _Token
+        Following tokens.
+    data : bytes
+        Underlying data that the token positions point to.
+
+    Returns
+    -------
+    _BalancedExpression
+    """
+    delim_stack = []
+    token = initial
+    while True:
+        if token.is_delim():
+            if token.raw in ('[', '{'):
+                delim_stack.append(token)
+            elif token.raw in (']', '}'):
+                if not delim_stack:
+                    raise RuntimeError(f"unmatched closing token {token}")
+                match = delim_stack.pop()
+                if match.raw != token.opposite():
+                    raise RuntimeError(
+                        f"opening token {match} closed by {token}"
+                    )
+                if not delim_stack:
+                    break
+            else:
+                raise RuntimeError(f'unknown delimiter {token}')
+        elif not delim_stack:
+            break
+        token = next(tokens)
+    return _BalancedExpression(
+        initial.pos,
+        data[initial.pos:token.endpos()].decode('ascii', 'replace')
+    )
+
+
+class Type1Font:
+    """
+    A class representing a Type-1 font, for use by backends.
+
+    Attributes
+    ----------
+    parts : tuple
+        A 3-tuple of the cleartext part, the encrypted part, and the finale of
+        zeros.
+
+    decrypted : bytes
+        The decrypted form of ``parts[1]``.
+
+    prop : dict[str, Any]
+        A dictionary of font properties. Noteworthy keys include:
+
+        - FontName: PostScript name of the font
+        - Encoding: dict from numeric codes to glyph names
+        - FontMatrix: bytes object encoding a matrix
+        - UniqueID: optional font identifier, dropped when modifying the font
+        - CharStrings: dict from glyph names to byte code
+        - Subrs: array of byte code subroutines
+        - OtherSubrs: bytes object encoding some PostScript code
+    """
+    __slots__ = ('parts', 'decrypted', 'prop', '_pos', '_abbr')
+    # the _pos dict contains (begin, end) indices to parts[0] + decrypted
+    # so that they can be replaced when transforming the font;
+    # but since sometimes a definition appears in both parts[0] and decrypted,
+    # _pos[name] is an array of such pairs
+    #
+    # _abbr maps three standard abbreviations to their particular names in
+    # this font (e.g. 'RD' is named '-|' in some fonts)
+
+    def __init__(self, input):
+        """
+        Initialize a Type-1 font.
+
+        Parameters
+        ----------
+        input : str or 3-tuple
+            Either a pfb file name, or a 3-tuple of already-decoded Type-1
+            font `~.Type1Font.parts`.
+        """
+        if isinstance(input, tuple) and len(input) == 3:
+            self.parts = input
+        else:
+            with open(input, 'rb') as file:
+                data = self._read(file)
+            self.parts = self._split(data)
+
+        self.decrypted = self._decrypt(self.parts[1], 'eexec')
+        self._abbr = {'RD': 'RD', 'ND': 'ND', 'NP': 'NP'}
+        self._parse()
+
+    def _read(self, file):
+        """Read the font from a file, decoding into usable parts."""
+        rawdata = file.read()
+        if not rawdata.startswith(b'\x80'):
+            return rawdata
+
+        data = b''
+        while rawdata:
+            if not rawdata.startswith(b'\x80'):
+                raise RuntimeError('Broken pfb file (expected byte 128, '
+                                   'got %d)' % rawdata[0])
+            type = rawdata[1]
+            if type in (1, 2):
+                length, = struct.unpack('<i', rawdata[2:6])
+                segment = rawdata[6:6 + length]
+                rawdata = rawdata[6 + length:]
+
+            if type == 1:       # ASCII text: include verbatim
+                data += segment
+            elif type == 2:     # binary data: encode in hexadecimal
+                data += binascii.hexlify(segment)
+            elif type == 3:     # end of file
+                break
+            else:
+                raise RuntimeError('Unknown segment type %d in pfb file' % type)
+
+        return data
+
+    def _split(self, data):
+        """
+        Split the Type 1 font into its three main parts.
+
+        The three parts are: (1) the cleartext part, which ends in a
+        eexec operator; (2) the encrypted part; (3) the fixed part,
+        which contains 512 ASCII zeros possibly divided on various
+        lines, a cleartomark operator, and possibly something else.
+        """
+
+        # Cleartext part: just find the eexec and skip whitespace
+        idx = data.index(b'eexec')
+        idx += len(b'eexec')
+        while data[idx] in b' \t\r\n':
+            idx += 1
+        len1 = idx
+
+        # Encrypted part: find the cleartomark operator and count
+        # zeros backward
+        idx = data.rindex(b'cleartomark') - 1
+        zeros = 512
+        while zeros and data[idx] in b'0' or data[idx] in b'\r\n':
+            if data[idx] in b'0':
+                zeros -= 1
+            idx -= 1
+        if zeros:
+            # this may have been a problem on old implementations that
+            # used the zeros as necessary padding
+            _log.info('Insufficiently many zeros in Type 1 font')
+
+        # Convert encrypted part to binary (if we read a pfb file, we may end
+        # up converting binary to hexadecimal to binary again; but if we read
+        # a pfa file, this part is already in hex, and I am not quite sure if
+        # even the pfb format guarantees that it will be in binary).
+        idx1 = len1 + ((idx - len1 + 2) & ~1)  # ensure an even number of bytes
+        binary = binascii.unhexlify(data[len1:idx1])
+
+        return data[:len1], binary, data[idx+1:]
+
+    @staticmethod
+    def _decrypt(ciphertext, key, ndiscard=4):
+        """
+        Decrypt ciphertext using the Type-1 font algorithm.
+
+        The algorithm is described in Adobe's "Adobe Type 1 Font Format".
+        The key argument can be an integer, or one of the strings
+        'eexec' and 'charstring', which map to the key specified for the
+        corresponding part of Type-1 fonts.
+
+        The ndiscard argument should be an integer, usually 4.
+        That number of bytes is discarded from the beginning of plaintext.
+        """
+
+        key = _api.check_getitem({'eexec': 55665, 'charstring': 4330}, key=key)
+        plaintext = []
+        for byte in ciphertext:
+            plaintext.append(byte ^ (key >> 8))
+            key = ((key+byte) * 52845 + 22719) & 0xffff
+
+        return bytes(plaintext[ndiscard:])
+
+    @staticmethod
+    def _encrypt(plaintext, key, ndiscard=4):
+        """
+        Encrypt plaintext using the Type-1 font algorithm.
+
+        The algorithm is described in Adobe's "Adobe Type 1 Font Format".
+        The key argument can be an integer, or one of the strings
+        'eexec' and 'charstring', which map to the key specified for the
+        corresponding part of Type-1 fonts.
+
+        The ndiscard argument should be an integer, usually 4. That
+        number of bytes is prepended to the plaintext before encryption.
+        This function prepends NUL bytes for reproducibility, even though
+        the original algorithm uses random bytes, presumably to avoid
+        cryptanalysis.
+        """
+
+        key = _api.check_getitem({'eexec': 55665, 'charstring': 4330}, key=key)
+        ciphertext = []
+        for byte in b'\0' * ndiscard + plaintext:
+            c = byte ^ (key >> 8)
+            ciphertext.append(c)
+            key = ((key + c) * 52845 + 22719) & 0xffff
+
+        return bytes(ciphertext)
+
+    def _parse(self):
+        """
+        Find the values of various font properties. This limited kind
+        of parsing is described in Chapter 10 "Adobe Type Manager
+        Compatibility" of the Type-1 spec.
+        """
+        # Start with reasonable defaults
+        prop = {'Weight': 'Regular', 'ItalicAngle': 0.0, 'isFixedPitch': False,
+                'UnderlinePosition': -100, 'UnderlineThickness': 50}
+        pos = {}
+        data = self.parts[0] + self.decrypted
+
+        source = _tokenize(data, True)
+        while True:
+            # See if there is a key to be assigned a value
+            # e.g. /FontName in /FontName /Helvetica def
+            try:
+                token = next(source)
+            except StopIteration:
+                break
+            if token.is_delim():
+                # skip over this - we want top-level keys only
+                _expression(token, source, data)
+            if token.is_slash_name():
+                key = token.value()
+                keypos = token.pos
+            else:
+                continue
+
+            # Some values need special parsing
+            if key in ('Subrs', 'CharStrings', 'Encoding', 'OtherSubrs'):
+                prop[key], endpos = {
+                    'Subrs': self._parse_subrs,
+                    'CharStrings': self._parse_charstrings,
+                    'Encoding': self._parse_encoding,
+                    'OtherSubrs': self._parse_othersubrs
+                }[key](source, data)
+                pos.setdefault(key, []).append((keypos, endpos))
+                continue
+
+            try:
+                token = next(source)
+            except StopIteration:
+                break
+
+            if isinstance(token, _KeywordToken):
+                # constructs like
+                # FontDirectory /Helvetica known {...} {...} ifelse
+                # mean the key was not really a key
+                continue
+
+            if token.is_delim():
+                value = _expression(token, source, data).raw
+            else:
+                value = token.value()
+
+            # look for a 'def' possibly preceded by access modifiers
+            try:
+                kw = next(
+                    kw for kw in source
+                    if not kw.is_keyword('readonly', 'noaccess', 'executeonly')
+                )
+            except StopIteration:
+                break
+
+            # sometimes noaccess def and readonly def are abbreviated
+            if kw.is_keyword('def', self._abbr['ND'], self._abbr['NP']):
+                prop[key] = value
+                pos.setdefault(key, []).append((keypos, kw.endpos()))
+
+            # detect the standard abbreviations
+            if value == '{noaccess def}':
+                self._abbr['ND'] = key
+            elif value == '{noaccess put}':
+                self._abbr['NP'] = key
+            elif value == '{string currentfile exch readstring pop}':
+                self._abbr['RD'] = key
+
+        # Fill in the various *Name properties
+        if 'FontName' not in prop:
+            prop['FontName'] = (prop.get('FullName') or
+                                prop.get('FamilyName') or
+                                'Unknown')
+        if 'FullName' not in prop:
+            prop['FullName'] = prop['FontName']
+        if 'FamilyName' not in prop:
+            extras = ('(?i)([ -](regular|plain|italic|oblique|(semi)?bold|'
+                      '(ultra)?light|extra|condensed))+$')
+            prop['FamilyName'] = re.sub(extras, '', prop['FullName'])
+        # Decrypt the encrypted parts
+        ndiscard = prop.get('lenIV', 4)
+        cs = prop['CharStrings']
+        for key, value in cs.items():
+            cs[key] = self._decrypt(value, 'charstring', ndiscard)
+        if 'Subrs' in prop:
+            prop['Subrs'] = [
+                self._decrypt(value, 'charstring', ndiscard)
+                for value in prop['Subrs']
+            ]
+
+        self.prop = prop
+        self._pos = pos
+
+    def _parse_subrs(self, tokens, _data):
+        count_token = next(tokens)
+        if not count_token.is_number():
+            raise RuntimeError(
+                f"Token following /Subrs must be a number, was {count_token}"
+            )
+        count = count_token.value()
+        array = [None] * count
+        next(t for t in tokens if t.is_keyword('array'))
+        for _ in range(count):
+            next(t for t in tokens if t.is_keyword('dup'))
+            index_token = next(tokens)
+            if not index_token.is_number():
+                raise RuntimeError(
+                    "Token following dup in Subrs definition must be a "
+                    f"number, was {index_token}"
+                )
+            nbytes_token = next(tokens)
+            if not nbytes_token.is_number():
+                raise RuntimeError(
+                    "Second token following dup in Subrs definition must "
+                    f"be a number, was {nbytes_token}"
+                )
+            token = next(tokens)
+            if not token.is_keyword(self._abbr['RD']):
+                raise RuntimeError(
+                    f"Token preceding subr must be {self._abbr['RD']}, "
+                    f"was {token}"
+                )
+            binary_token = tokens.send(1+nbytes_token.value())
+            array[index_token.value()] = binary_token.value()
+
+        return array, next(tokens).endpos()
+
+    @staticmethod
+    def _parse_charstrings(tokens, _data):
+        count_token = next(tokens)
+        if not count_token.is_number():
+            raise RuntimeError(
+                "Token following /CharStrings must be a number, "
+                f"was {count_token}"
+            )
+        count = count_token.value()
+        charstrings = {}
+        next(t for t in tokens if t.is_keyword('begin'))
+        while True:
+            token = next(t for t in tokens
+                         if t.is_keyword('end') or t.is_slash_name())
+            if token.raw == 'end':
+                return charstrings, token.endpos()
+            glyphname = token.value()
+            nbytes_token = next(tokens)
+            if not nbytes_token.is_number():
+                raise RuntimeError(
+                    f"Token following /{glyphname} in CharStrings definition "
+                    f"must be a number, was {nbytes_token}"
+                )
+            next(tokens)  # usually RD or |-
+            binary_token = tokens.send(1+nbytes_token.value())
+            charstrings[glyphname] = binary_token.value()
+
+    @staticmethod
+    def _parse_encoding(tokens, _data):
+        # this only works for encodings that follow the Adobe manual
+        # but some old fonts include non-compliant data - we log a warning
+        # and return a possibly incomplete encoding
+        encoding = {}
+        while True:
+            token = next(t for t in tokens
+                         if t.is_keyword('StandardEncoding', 'dup', 'def'))
+            if token.is_keyword('StandardEncoding'):
+                return _StandardEncoding, token.endpos()
+            if token.is_keyword('def'):
+                return encoding, token.endpos()
+            index_token = next(tokens)
+            if not index_token.is_number():
+                _log.warning(
+                    f"Parsing encoding: expected number, got {index_token}"
+                )
+                continue
+            name_token = next(tokens)
+            if not name_token.is_slash_name():
+                _log.warning(
+                    f"Parsing encoding: expected slash-name, got {name_token}"
+                )
+                continue
+            encoding[index_token.value()] = name_token.value()
+
+    @staticmethod
+    def _parse_othersubrs(tokens, data):
+        init_pos = None
+        while True:
+            token = next(tokens)
+            if init_pos is None:
+                init_pos = token.pos
+            if token.is_delim():
+                _expression(token, tokens, data)
+            elif token.is_keyword('def', 'ND', '|-'):
+                return data[init_pos:token.endpos()], token.endpos()
+
+    def transform(self, effects):
+        """
+        Return a new font that is slanted and/or extended.
+
+        Parameters
+        ----------
+        effects : dict
+            A dict with optional entries:
+
+            - 'slant' : float, default: 0
+                Tangent of the angle that the font is to be slanted to the
+                right. Negative values slant to the left.
+            - 'extend' : float, default: 1
+                Scaling factor for the font width. Values less than 1 condense
+                the glyphs.
+
+        Returns
+        -------
+        `Type1Font`
+        """
+        fontname = self.prop['FontName']
+        italicangle = self.prop['ItalicAngle']
+
+        array = [
+            float(x) for x in (self.prop['FontMatrix']
+                               .lstrip('[').rstrip(']').split())
+        ]
+        oldmatrix = np.eye(3, 3)
+        oldmatrix[0:3, 0] = array[::2]
+        oldmatrix[0:3, 1] = array[1::2]
+        modifier = np.eye(3, 3)
+
+        if 'slant' in effects:
+            slant = effects['slant']
+            fontname += f'_Slant_{int(1000 * slant)}'
+            italicangle = round(
+                float(italicangle) - np.arctan(slant) / np.pi * 180,
+                5
+            )
+            modifier[1, 0] = slant
+
+        if 'extend' in effects:
+            extend = effects['extend']
+            fontname += f'_Extend_{int(1000 * extend)}'
+            modifier[0, 0] = extend
+
+        newmatrix = np.dot(modifier, oldmatrix)
+        array[::2] = newmatrix[0:3, 0]
+        array[1::2] = newmatrix[0:3, 1]
+        fontmatrix = (
+            f"[{' '.join(_format_approx(x, 6) for x in array)}]"
+        )
+        replacements = (
+            [(x, f'/FontName/{fontname} def')
+             for x in self._pos['FontName']]
+            + [(x, f'/ItalicAngle {italicangle} def')
+               for x in self._pos['ItalicAngle']]
+            + [(x, f'/FontMatrix {fontmatrix} readonly def')
+               for x in self._pos['FontMatrix']]
+            + [(x, '') for x in self._pos.get('UniqueID', [])]
+        )
+
+        data = bytearray(self.parts[0])
+        data.extend(self.decrypted)
+        len0 = len(self.parts[0])
+        for (pos0, pos1), value in sorted(replacements, reverse=True):
+            data[pos0:pos1] = value.encode('ascii', 'replace')
+            if pos0 < len(self.parts[0]):
+                if pos1 >= len(self.parts[0]):
+                    raise RuntimeError(
+                        f"text to be replaced with {value} spans "
+                        "the eexec boundary"
+                    )
+                len0 += len(value) - pos1 + pos0
+
+        data = bytes(data)
+        return Type1Font((
+            data[:len0],
+            self._encrypt(data[len0:], 'eexec'),
+            self.parts[2]
+        ))
+
+
+_StandardEncoding = {
+    **{ord(letter): letter for letter in string.ascii_letters},
+    0: '.notdef',
+    32: 'space',
+    33: 'exclam',
+    34: 'quotedbl',
+    35: 'numbersign',
+    36: 'dollar',
+    37: 'percent',
+    38: 'ampersand',
+    39: 'quoteright',
+    40: 'parenleft',
+    41: 'parenright',
+    42: 'asterisk',
+    43: 'plus',
+    44: 'comma',
+    45: 'hyphen',
+    46: 'period',
+    47: 'slash',
+    48: 'zero',
+    49: 'one',
+    50: 'two',
+    51: 'three',
+    52: 'four',
+    53: 'five',
+    54: 'six',
+    55: 'seven',
+    56: 'eight',
+    57: 'nine',
+    58: 'colon',
+    59: 'semicolon',
+    60: 'less',
+    61: 'equal',
+    62: 'greater',
+    63: 'question',
+    64: 'at',
+    91: 'bracketleft',
+    92: 'backslash',
+    93: 'bracketright',
+    94: 'asciicircum',
+    95: 'underscore',
+    96: 'quoteleft',
+    123: 'braceleft',
+    124: 'bar',
+    125: 'braceright',
+    126: 'asciitilde',
+    161: 'exclamdown',
+    162: 'cent',
+    163: 'sterling',
+    164: 'fraction',
+    165: 'yen',
+    166: 'florin',
+    167: 'section',
+    168: 'currency',
+    169: 'quotesingle',
+    170: 'quotedblleft',
+    171: 'guillemotleft',
+    172: 'guilsinglleft',
+    173: 'guilsinglright',
+    174: 'fi',
+    175: 'fl',
+    177: 'endash',
+    178: 'dagger',
+    179: 'daggerdbl',
+    180: 'periodcentered',
+    182: 'paragraph',
+    183: 'bullet',
+    184: 'quotesinglbase',
+    185: 'quotedblbase',
+    186: 'quotedblright',
+    187: 'guillemotright',
+    188: 'ellipsis',
+    189: 'perthousand',
+    191: 'questiondown',
+    193: 'grave',
+    194: 'acute',
+    195: 'circumflex',
+    196: 'tilde',
+    197: 'macron',
+    198: 'breve',
+    199: 'dotaccent',
+    200: 'dieresis',
+    202: 'ring',
+    203: 'cedilla',
+    205: 'hungarumlaut',
+    206: 'ogonek',
+    207: 'caron',
+    208: 'emdash',
+    225: 'AE',
+    227: 'ordfeminine',
+    232: 'Lslash',
+    233: 'Oslash',
+    234: 'OE',
+    235: 'ordmasculine',
+    241: 'ae',
+    245: 'dotlessi',
+    248: 'lslash',
+    249: 'oslash',
+    250: 'oe',
+    251: 'germandbls',
+}

lib/python3.12/site-packages/matplotlib/backend_managers.pyi

@@ -0,0 +1,64 @@
+from matplotlib import backend_tools, widgets
+from matplotlib.backend_bases import FigureCanvasBase
+from matplotlib.figure import Figure
+
+from collections.abc import Callable, Iterable
+from typing import Any, TypeVar
+
+class ToolEvent:
+    name: str
+    sender: Any
+    tool: backend_tools.ToolBase
+    data: Any
+    def __init__(self, name, sender, tool, data: Any | None = ...) -> None: ...
+
+class ToolTriggerEvent(ToolEvent):
+    canvasevent: ToolEvent
+    def __init__(
+        self,
+        name,
+        sender,
+        tool,
+        canvasevent: ToolEvent | None = ...,
+        data: Any | None = ...,
+    ) -> None: ...
+
+class ToolManagerMessageEvent:
+    name: str
+    sender: Any
+    message: str
+    def __init__(self, name: str, sender: Any, message: str) -> None: ...
+
+class ToolManager:
+    keypresslock: widgets.LockDraw
+    messagelock: widgets.LockDraw
+    def __init__(self, figure: Figure | None = ...) -> None: ...
+    @property
+    def canvas(self) -> FigureCanvasBase | None: ...
+    @property
+    def figure(self) -> Figure | None: ...
+    @figure.setter
+    def figure(self, figure: Figure) -> None: ...
+    def set_figure(self, figure: Figure, update_tools: bool = ...) -> None: ...
+    def toolmanager_connect(self, s: str, func: Callable[[ToolEvent], Any]) -> int: ...
+    def toolmanager_disconnect(self, cid: int) -> None: ...
+    def message_event(self, message: str, sender: Any | None = ...) -> None: ...
+    @property
+    def active_toggle(self) -> dict[str | None, list[str] | str]: ...
+    def get_tool_keymap(self, name: str) -> list[str]: ...
+    def update_keymap(self, name: str, key: str | Iterable[str]) -> None: ...
+    def remove_tool(self, name: str) -> None: ...
+    _T = TypeVar("_T", bound=backend_tools.ToolBase)
+    def add_tool(self, name: str, tool: type[_T], *args, **kwargs) -> _T: ...
+    def trigger_tool(
+        self,
+        name: str | backend_tools.ToolBase,
+        sender: Any | None = ...,
+        canvasevent: ToolEvent | None = ...,
+        data: Any | None = ...,
+    ) -> None: ...
+    @property
+    def tools(self) -> dict[str, backend_tools.ToolBase]: ...
+    def get_tool(
+        self, name: str | backend_tools.ToolBase, warn: bool = ...
+    ) -> backend_tools.ToolBase | None: ...

lib/python3.12/site-packages/matplotlib/category.py

@@ -0,0 +1,235 @@
+"""
+Plotting of string "category" data: ``plot(['d', 'f', 'a'], [1, 2, 3])`` will
+plot three points with x-axis values of 'd', 'f', 'a'.
+
+See :doc:`/gallery/lines_bars_and_markers/categorical_variables` for an
+example.
+
+The module uses Matplotlib's `matplotlib.units` mechanism to convert from
+strings to integers and provides a tick locator, a tick formatter, and the
+`.UnitData` class that creates and stores the string-to-integer mapping.
+"""
+
+from collections import OrderedDict
+import dateutil.parser
+import itertools
+import logging
+
+import numpy as np
+
+from matplotlib import _api, cbook, ticker, units
+
+
+_log = logging.getLogger(__name__)
+
+
+class StrCategoryConverter(units.ConversionInterface):
+    @staticmethod
+    def convert(value, unit, axis):
+        """
+        Convert strings in *value* to floats using mapping information stored
+        in the *unit* object.
+
+        Parameters
+        ----------
+        value : str or iterable
+            Value or list of values to be converted.
+        unit : `.UnitData`
+            An object mapping strings to integers.
+        axis : `~matplotlib.axis.Axis`
+            The axis on which the converted value is plotted.
+
+            .. note:: *axis* is unused.
+
+        Returns
+        -------
+        float or `~numpy.ndarray` of float
+        """
+        if unit is None:
+            raise ValueError(
+                'Missing category information for StrCategoryConverter; '
+                'this might be caused by unintendedly mixing categorical and '
+                'numeric data')
+        StrCategoryConverter._validate_unit(unit)
+        # dtype = object preserves numerical pass throughs
+        values = np.atleast_1d(np.array(value, dtype=object))
+        # force an update so it also does type checking
+        unit.update(values)
+        s = np.vectorize(unit._mapping.__getitem__, otypes=[float])(values)
+        return s if not cbook.is_scalar_or_string(value) else s[0]
+
+    @staticmethod
+    def axisinfo(unit, axis):
+        """
+        Set the default axis ticks and labels.
+
+        Parameters
+        ----------
+        unit : `.UnitData`
+            object string unit information for value
+        axis : `~matplotlib.axis.Axis`
+            axis for which information is being set
+
+            .. note:: *axis* is not used
+
+        Returns
+        -------
+        `~matplotlib.units.AxisInfo`
+            Information to support default tick labeling
+
+        """
+        StrCategoryConverter._validate_unit(unit)
+        # locator and formatter take mapping dict because
+        # args need to be pass by reference for updates
+        majloc = StrCategoryLocator(unit._mapping)
+        majfmt = StrCategoryFormatter(unit._mapping)
+        return units.AxisInfo(majloc=majloc, majfmt=majfmt)
+
+    @staticmethod
+    def default_units(data, axis):
+        """
+        Set and update the `~matplotlib.axis.Axis` units.
+
+        Parameters
+        ----------
+        data : str or iterable of str
+        axis : `~matplotlib.axis.Axis`
+            axis on which the data is plotted
+
+        Returns
+        -------
+        `.UnitData`
+            object storing string to integer mapping
+        """
+        # the conversion call stack is default_units -> axis_info -> convert
+        if axis.units is None:
+            axis.set_units(UnitData(data))
+        else:
+            axis.units.update(data)
+        return axis.units
+
+    @staticmethod
+    def _validate_unit(unit):
+        if not hasattr(unit, '_mapping'):
+            raise ValueError(
+                f'Provided unit "{unit}" is not valid for a categorical '
+                'converter, as it does not have a _mapping attribute.')
+
+
+class StrCategoryLocator(ticker.Locator):
+    """Tick at every integer mapping of the string data."""
+    def __init__(self, units_mapping):
+        """
+        Parameters
+        ----------
+        units_mapping : dict
+            Mapping of category names (str) to indices (int).
+        """
+        self._units = units_mapping
+
+    def __call__(self):
+        # docstring inherited
+        return list(self._units.values())
+
+    def tick_values(self, vmin, vmax):
+        # docstring inherited
+        return self()
+
+
+class StrCategoryFormatter(ticker.Formatter):
+    """String representation of the data at every tick."""
+    def __init__(self, units_mapping):
+        """
+        Parameters
+        ----------
+        units_mapping : dict
+            Mapping of category names (str) to indices (int).
+        """
+        self._units = units_mapping
+
+    def __call__(self, x, pos=None):
+        # docstring inherited
+        return self.format_ticks([x])[0]
+
+    def format_ticks(self, values):
+        # docstring inherited
+        r_mapping = {v: self._text(k) for k, v in self._units.items()}
+        return [r_mapping.get(round(val), '') for val in values]
+
+    @staticmethod
+    def _text(value):
+        """Convert text values into utf-8 or ascii strings."""
+        if isinstance(value, bytes):
+            value = value.decode(encoding='utf-8')
+        elif not isinstance(value, str):
+            value = str(value)
+        return value
+
+
+class UnitData:
+    def __init__(self, data=None):
+        """
+        Create mapping between unique categorical values and integer ids.
+
+        Parameters
+        ----------
+        data : iterable
+            sequence of string values
+        """
+        self._mapping = OrderedDict()
+        self._counter = itertools.count()
+        if data is not None:
+            self.update(data)
+
+    @staticmethod
+    def _str_is_convertible(val):
+        """
+        Helper method to check whether a string can be parsed as float or date.
+        """
+        try:
+            float(val)
+        except ValueError:
+            try:
+                dateutil.parser.parse(val)
+            except (ValueError, TypeError):
+                # TypeError if dateutil >= 2.8.1 else ValueError
+                return False
+        return True
+
+    def update(self, data):
+        """
+        Map new values to integer identifiers.
+
+        Parameters
+        ----------
+        data : iterable of str or bytes
+
+        Raises
+        ------
+        TypeError
+            If elements in *data* are neither str nor bytes.
+        """
+        data = np.atleast_1d(np.array(data, dtype=object))
+        # check if convertible to number:
+        convertible = True
+        for val in OrderedDict.fromkeys(data):
+            # OrderedDict just iterates over unique values in data.
+            _api.check_isinstance((str, bytes), value=val)
+            if convertible:
+                # this will only be called so long as convertible is True.
+                convertible = self._str_is_convertible(val)
+            if val not in self._mapping:
+                self._mapping[val] = next(self._counter)
+        if data.size and convertible:
+            _log.info('Using categorical units to plot a list of strings '
+                      'that are all parsable as floats or dates. If these '
+                      'strings should be plotted as numbers, cast to the '
+                      'appropriate data type before plotting.')
+
+
+# Register the converter with Matplotlib's unit framework
+# Intentionally set to a single instance
+units.registry[str] = \
+    units.registry[np.str_] = \
+    units.registry[bytes] = \
+    units.registry[np.bytes_] = StrCategoryConverter()

lib/python3.12/site-packages/matplotlib/cbook.py

@@ -0,0 +1,2415 @@
+"""
+A collection of utility functions and classes.  Originally, many
+(but not all) were from the Python Cookbook -- hence the name cbook.
+"""
+
+import collections
+import collections.abc
+import contextlib
+import functools
+import gzip
+import itertools
+import math
+import operator
+import os
+from pathlib import Path
+import shlex
+import subprocess
+import sys
+import time
+import traceback
+import types
+import weakref
+
+import numpy as np
+
+try:
+    from numpy.exceptions import VisibleDeprecationWarning  # numpy >= 1.25
+except ImportError:
+    from numpy import VisibleDeprecationWarning
+
+import matplotlib
+from matplotlib import _api, _c_internal_utils
+
+
+class _ExceptionInfo:
+    """
+    A class to carry exception information around.
+
+    This is used to store and later raise exceptions. It's an alternative to
+    directly storing Exception instances that circumvents traceback-related
+    issues: caching tracebacks can keep user's objects in local namespaces
+    alive indefinitely, which can lead to very surprising memory issues for
+    users and result in incorrect tracebacks.
+    """
+
+    def __init__(self, cls, *args):
+        self._cls = cls
+        self._args = args
+
+    @classmethod
+    def from_exception(cls, exc):
+        return cls(type(exc), *exc.args)
+
+    def to_exception(self):
+        return self._cls(*self._args)
+
+
+def _get_running_interactive_framework():
+    """
+    Return the interactive framework whose event loop is currently running, if
+    any, or "headless" if no event loop can be started, or None.
+
+    Returns
+    -------
+    Optional[str]
+        One of the following values: "qt", "gtk3", "gtk4", "wx", "tk",
+        "macosx", "headless", ``None``.
+    """
+    # Use ``sys.modules.get(name)`` rather than ``name in sys.modules`` as
+    # entries can also have been explicitly set to None.
+    QtWidgets = (
+        sys.modules.get("PyQt6.QtWidgets")
+        or sys.modules.get("PySide6.QtWidgets")
+        or sys.modules.get("PyQt5.QtWidgets")
+        or sys.modules.get("PySide2.QtWidgets")
+    )
+    if QtWidgets and QtWidgets.QApplication.instance():
+        return "qt"
+    Gtk = sys.modules.get("gi.repository.Gtk")
+    if Gtk:
+        if Gtk.MAJOR_VERSION == 4:
+            from gi.repository import GLib
+            if GLib.main_depth():
+                return "gtk4"
+        if Gtk.MAJOR_VERSION == 3 and Gtk.main_level():
+            return "gtk3"
+    wx = sys.modules.get("wx")
+    if wx and wx.GetApp():
+        return "wx"
+    tkinter = sys.modules.get("tkinter")
+    if tkinter:
+        codes = {tkinter.mainloop.__code__, tkinter.Misc.mainloop.__code__}
+        for frame in sys._current_frames().values():
+            while frame:
+                if frame.f_code in codes:
+                    return "tk"
+                frame = frame.f_back
+        # Preemptively break reference cycle between locals and the frame.
+        del frame
+    macosx = sys.modules.get("matplotlib.backends._macosx")
+    if macosx and macosx.event_loop_is_running():
+        return "macosx"
+    if not _c_internal_utils.display_is_valid():
+        return "headless"
+    return None
+
+
+def _exception_printer(exc):
+    if _get_running_interactive_framework() in ["headless", None]:
+        raise exc
+    else:
+        traceback.print_exc()
+
+
+class _StrongRef:
+    """
+    Wrapper similar to a weakref, but keeping a strong reference to the object.
+    """
+
+    def __init__(self, obj):
+        self._obj = obj
+
+    def __call__(self):
+        return self._obj
+
+    def __eq__(self, other):
+        return isinstance(other, _StrongRef) and self._obj == other._obj
+
+    def __hash__(self):
+        return hash(self._obj)
+
+
+def _weak_or_strong_ref(func, callback):
+    """
+    Return a `WeakMethod` wrapping *func* if possible, else a `_StrongRef`.
+    """
+    try:
+        return weakref.WeakMethod(func, callback)
+    except TypeError:
+        return _StrongRef(func)
+
+
+class _UnhashDict:
+    """
+    A minimal dict-like class that also supports unhashable keys, storing them
+    in a list of key-value pairs.
+
+    This class only implements the interface needed for `CallbackRegistry`, and
+    tries to minimize the overhead for the hashable case.
+    """
+
+    def __init__(self, pairs):
+        self._dict = {}
+        self._pairs = []
+        for k, v in pairs:
+            self[k] = v
+
+    def __setitem__(self, key, value):
+        try:
+            self._dict[key] = value
+        except TypeError:
+            for i, (k, v) in enumerate(self._pairs):
+                if k == key:
+                    self._pairs[i] = (key, value)
+                    break
+            else:
+                self._pairs.append((key, value))
+
+    def __getitem__(self, key):
+        try:
+            return self._dict[key]
+        except TypeError:
+            pass
+        for k, v in self._pairs:
+            if k == key:
+                return v
+        raise KeyError(key)
+
+    def pop(self, key, *args):
+        try:
+            if key in self._dict:
+                return self._dict.pop(key)
+        except TypeError:
+            for i, (k, v) in enumerate(self._pairs):
+                if k == key:
+                    del self._pairs[i]
+                    return v
+        if args:
+            return args[0]
+        raise KeyError(key)
+
+    def __iter__(self):
+        yield from self._dict
+        for k, v in self._pairs:
+            yield k
+
+
+class CallbackRegistry:
+    """
+    Handle registering, processing, blocking, and disconnecting
+    for a set of signals and callbacks:
+
+        >>> def oneat(x):
+        ...     print('eat', x)
+        >>> def ondrink(x):
+        ...     print('drink', x)
+
+        >>> from matplotlib.cbook import CallbackRegistry
+        >>> callbacks = CallbackRegistry()
+
+        >>> id_eat = callbacks.connect('eat', oneat)
+        >>> id_drink = callbacks.connect('drink', ondrink)
+
+        >>> callbacks.process('drink', 123)
+        drink 123
+        >>> callbacks.process('eat', 456)
+        eat 456
+        >>> callbacks.process('be merry', 456)   # nothing will be called
+
+        >>> callbacks.disconnect(id_eat)
+        >>> callbacks.process('eat', 456)        # nothing will be called
+
+        >>> with callbacks.blocked(signal='drink'):
+        ...     callbacks.process('drink', 123)  # nothing will be called
+        >>> callbacks.process('drink', 123)
+        drink 123
+
+    In practice, one should always disconnect all callbacks when they are
+    no longer needed to avoid dangling references (and thus memory leaks).
+    However, real code in Matplotlib rarely does so, and due to its design,
+    it is rather difficult to place this kind of code.  To get around this,
+    and prevent this class of memory leaks, we instead store weak references
+    to bound methods only, so when the destination object needs to die, the
+    CallbackRegistry won't keep it alive.
+
+    Parameters
+    ----------
+    exception_handler : callable, optional
+       If not None, *exception_handler* must be a function that takes an
+       `Exception` as single parameter.  It gets called with any `Exception`
+       raised by the callbacks during `CallbackRegistry.process`, and may
+       either re-raise the exception or handle it in another manner.
+
+       The default handler prints the exception (with `traceback.print_exc`) if
+       an interactive event loop is running; it re-raises the exception if no
+       interactive event loop is running.
+
+    signals : list, optional
+        If not None, *signals* is a list of signals that this registry handles:
+        attempting to `process` or to `connect` to a signal not in the list
+        throws a `ValueError`.  The default, None, does not restrict the
+        handled signals.
+    """
+
+    # We maintain two mappings:
+    #   callbacks: signal -> {cid -> weakref-to-callback}
+    #   _func_cid_map: {(signal, weakref-to-callback) -> cid}
+
+    def __init__(self, exception_handler=_exception_printer, *, signals=None):
+        self._signals = None if signals is None else list(signals)  # Copy it.
+        self.exception_handler = exception_handler
+        self.callbacks = {}
+        self._cid_gen = itertools.count()
+        self._func_cid_map = _UnhashDict([])
+        # A hidden variable that marks cids that need to be pickled.
+        self._pickled_cids = set()
+
+    def __getstate__(self):
+        return {
+            **vars(self),
+            # In general, callbacks may not be pickled, so we just drop them,
+            # unless directed otherwise by self._pickled_cids.
+            "callbacks": {s: {cid: proxy() for cid, proxy in d.items()
+                              if cid in self._pickled_cids}
+                          for s, d in self.callbacks.items()},
+            # It is simpler to reconstruct this from callbacks in __setstate__.
+            "_func_cid_map": None,
+            "_cid_gen": next(self._cid_gen)
+        }
+
+    def __setstate__(self, state):
+        cid_count = state.pop('_cid_gen')
+        vars(self).update(state)
+        self.callbacks = {
+            s: {cid: _weak_or_strong_ref(func, functools.partial(self._remove_proxy, s))
+                for cid, func in d.items()}
+            for s, d in self.callbacks.items()}
+        self._func_cid_map = _UnhashDict(
+            ((s, proxy), cid)
+            for s, d in self.callbacks.items() for cid, proxy in d.items())
+        self._cid_gen = itertools.count(cid_count)
+
+    def connect(self, signal, func):
+        """Register *func* to be called when signal *signal* is generated."""
+        if self._signals is not None:
+            _api.check_in_list(self._signals, signal=signal)
+        proxy = _weak_or_strong_ref(func, functools.partial(self._remove_proxy, signal))
+        try:
+            return self._func_cid_map[signal, proxy]
+        except KeyError:
+            cid = self._func_cid_map[signal, proxy] = next(self._cid_gen)
+            self.callbacks.setdefault(signal, {})[cid] = proxy
+            return cid
+
+    def _connect_picklable(self, signal, func):
+        """
+        Like `.connect`, but the callback is kept when pickling/unpickling.
+
+        Currently internal-use only.
+        """
+        cid = self.connect(signal, func)
+        self._pickled_cids.add(cid)
+        return cid
+
+    # Keep a reference to sys.is_finalizing, as sys may have been cleared out
+    # at that point.
+    def _remove_proxy(self, signal, proxy, *, _is_finalizing=sys.is_finalizing):
+        if _is_finalizing():
+            # Weakrefs can't be properly torn down at that point anymore.
+            return
+        cid = self._func_cid_map.pop((signal, proxy), None)
+        if cid is not None:
+            del self.callbacks[signal][cid]
+            self._pickled_cids.discard(cid)
+        else:  # Not found
+            return
+        if len(self.callbacks[signal]) == 0:  # Clean up empty dicts
+            del self.callbacks[signal]
+
+    def disconnect(self, cid):
+        """
+        Disconnect the callback registered with callback id *cid*.
+
+        No error is raised if such a callback does not exist.
+        """
+        self._pickled_cids.discard(cid)
+        for signal, proxy in self._func_cid_map:
+            if self._func_cid_map[signal, proxy] == cid:
+                break
+        else:  # Not found
+            return
+        assert self.callbacks[signal][cid] == proxy
+        del self.callbacks[signal][cid]
+        self._func_cid_map.pop((signal, proxy))
+        if len(self.callbacks[signal]) == 0:  # Clean up empty dicts
+            del self.callbacks[signal]
+
+    def process(self, s, *args, **kwargs):
+        """
+        Process signal *s*.
+
+        All of the functions registered to receive callbacks on *s* will be
+        called with ``*args`` and ``**kwargs``.
+        """
+        if self._signals is not None:
+            _api.check_in_list(self._signals, signal=s)
+        for ref in list(self.callbacks.get(s, {}).values()):
+            func = ref()
+            if func is not None:
+                try:
+                    func(*args, **kwargs)
+                # this does not capture KeyboardInterrupt, SystemExit,
+                # and GeneratorExit
+                except Exception as exc:
+                    if self.exception_handler is not None:
+                        self.exception_handler(exc)
+                    else:
+                        raise
+
+    @contextlib.contextmanager
+    def blocked(self, *, signal=None):
+        """
+        Block callback signals from being processed.
+
+        A context manager to temporarily block/disable callback signals
+        from being processed by the registered listeners.
+
+        Parameters
+        ----------
+        signal : str, optional
+            The callback signal to block. The default is to block all signals.
+        """
+        orig = self.callbacks
+        try:
+            if signal is None:
+                # Empty out the callbacks
+                self.callbacks = {}
+            else:
+                # Only remove the specific signal
+                self.callbacks = {k: orig[k] for k in orig if k != signal}
+            yield
+        finally:
+            self.callbacks = orig
+
+
+class silent_list(list):
+    """
+    A list with a short ``repr()``.
+
+    This is meant to be used for a homogeneous list of artists, so that they
+    don't cause long, meaningless output.
+
+    Instead of ::
+
+        [<matplotlib.lines.Line2D object at 0x7f5749fed3c8>,
+         <matplotlib.lines.Line2D object at 0x7f5749fed4e0>,
+         <matplotlib.lines.Line2D object at 0x7f5758016550>]
+
+    one will get ::
+
+        <a list of 3 Line2D objects>
+
+    If ``self.type`` is None, the type name is obtained from the first item in
+    the list (if any).
+    """
+
+    def __init__(self, type, seq=None):
+        self.type = type
+        if seq is not None:
+            self.extend(seq)
+
+    def __repr__(self):
+        if self.type is not None or len(self) != 0:
+            tp = self.type if self.type is not None else type(self[0]).__name__
+            return f"<a list of {len(self)} {tp} objects>"
+        else:
+            return "<an empty list>"
+
+
+def _local_over_kwdict(
+        local_var, kwargs, *keys,
+        warning_cls=_api.MatplotlibDeprecationWarning):
+    out = local_var
+    for key in keys:
+        kwarg_val = kwargs.pop(key, None)
+        if kwarg_val is not None:
+            if out is None:
+                out = kwarg_val
+            else:
+                _api.warn_external(f'"{key}" keyword argument will be ignored',
+                                   warning_cls)
+    return out
+
+
+def strip_math(s):
+    """
+    Remove latex formatting from mathtext.
+
+    Only handles fully math and fully non-math strings.
+    """
+    if len(s) >= 2 and s[0] == s[-1] == "$":
+        s = s[1:-1]
+        for tex, plain in [
+                (r"\times", "x"),  # Specifically for Formatter support.
+                (r"\mathdefault", ""),
+                (r"\rm", ""),
+                (r"\cal", ""),
+                (r"\tt", ""),
+                (r"\it", ""),
+                ("\\", ""),
+                ("{", ""),
+                ("}", ""),
+        ]:
+            s = s.replace(tex, plain)
+    return s
+
+
+def _strip_comment(s):
+    """Strip everything from the first unquoted #."""
+    pos = 0
+    while True:
+        quote_pos = s.find('"', pos)
+        hash_pos = s.find('#', pos)
+        if quote_pos < 0:
+            without_comment = s if hash_pos < 0 else s[:hash_pos]
+            return without_comment.strip()
+        elif 0 <= hash_pos < quote_pos:
+            return s[:hash_pos].strip()
+        else:
+            closing_quote_pos = s.find('"', quote_pos + 1)
+            if closing_quote_pos < 0:
+                raise ValueError(
+                    f"Missing closing quote in: {s!r}. If you need a double-"
+                    'quote inside a string, use escaping: e.g. "the \" char"')
+            pos = closing_quote_pos + 1  # behind closing quote
+
+
+def is_writable_file_like(obj):
+    """Return whether *obj* looks like a file object with a *write* method."""
+    return callable(getattr(obj, 'write', None))
+
+
+def file_requires_unicode(x):
+    """
+    Return whether the given writable file-like object requires Unicode to be
+    written to it.
+    """
+    try:
+        x.write(b'')
+    except TypeError:
+        return True
+    else:
+        return False
+
+
+def to_filehandle(fname, flag='r', return_opened=False, encoding=None):
+    """
+    Convert a path to an open file handle or pass-through a file-like object.
+
+    Consider using `open_file_cm` instead, as it allows one to properly close
+    newly created file objects more easily.
+
+    Parameters
+    ----------
+    fname : str or path-like or file-like
+        If `str` or `os.PathLike`, the file is opened using the flags specified
+        by *flag* and *encoding*.  If a file-like object, it is passed through.
+    flag : str, default: 'r'
+        Passed as the *mode* argument to `open` when *fname* is `str` or
+        `os.PathLike`; ignored if *fname* is file-like.
+    return_opened : bool, default: False
+        If True, return both the file object and a boolean indicating whether
+        this was a new file (that the caller needs to close).  If False, return
+        only the new file.
+    encoding : str or None, default: None
+        Passed as the *mode* argument to `open` when *fname* is `str` or
+        `os.PathLike`; ignored if *fname* is file-like.
+
+    Returns
+    -------
+    fh : file-like
+    opened : bool
+        *opened* is only returned if *return_opened* is True.
+    """
+    if isinstance(fname, os.PathLike):
+        fname = os.fspath(fname)
+    if isinstance(fname, str):
+        if fname.endswith('.gz'):
+            fh = gzip.open(fname, flag)
+        elif fname.endswith('.bz2'):
+            # python may not be compiled with bz2 support,
+            # bury import until we need it
+            import bz2
+            fh = bz2.BZ2File(fname, flag)
+        else:
+            fh = open(fname, flag, encoding=encoding)
+        opened = True
+    elif hasattr(fname, 'seek'):
+        fh = fname
+        opened = False
+    else:
+        raise ValueError('fname must be a PathLike or file handle')
+    if return_opened:
+        return fh, opened
+    return fh
+
+
+def open_file_cm(path_or_file, mode="r", encoding=None):
+    r"""Pass through file objects and context-manage path-likes."""
+    fh, opened = to_filehandle(path_or_file, mode, True, encoding)
+    return fh if opened else contextlib.nullcontext(fh)
+
+
+def is_scalar_or_string(val):
+    """Return whether the given object is a scalar or string like."""
+    return isinstance(val, str) or not np.iterable(val)
+
+
+def get_sample_data(fname, asfileobj=True):
+    """
+    Return a sample data file.  *fname* is a path relative to the
+    :file:`mpl-data/sample_data` directory.  If *asfileobj* is `True`
+    return a file object, otherwise just a file path.
+
+    Sample data files are stored in the 'mpl-data/sample_data' directory within
+    the Matplotlib package.
+
+    If the filename ends in .gz, the file is implicitly ungzipped.  If the
+    filename ends with .npy or .npz, and *asfileobj* is `True`, the file is
+    loaded with `numpy.load`.
+    """
+    path = _get_data_path('sample_data', fname)
+    if asfileobj:
+        suffix = path.suffix.lower()
+        if suffix == '.gz':
+            return gzip.open(path)
+        elif suffix in ['.npy', '.npz']:
+            return np.load(path)
+        elif suffix in ['.csv', '.xrc', '.txt']:
+            return path.open('r')
+        else:
+            return path.open('rb')
+    else:
+        return str(path)
+
+
+def _get_data_path(*args):
+    """
+    Return the `pathlib.Path` to a resource file provided by Matplotlib.
+
+    ``*args`` specify a path relative to the base data path.
+    """
+    return Path(matplotlib.get_data_path(), *args)
+
+
+def flatten(seq, scalarp=is_scalar_or_string):
+    """
+    Return a generator of flattened nested containers.
+
+    For example:
+
+        >>> from matplotlib.cbook import flatten
+        >>> l = (('John', ['Hunter']), (1, 23), [[([42, (5, 23)], )]])
+        >>> print(list(flatten(l)))
+        ['John', 'Hunter', 1, 23, 42, 5, 23]
+
+    By: Composite of Holger Krekel and Luther Blissett
+    From: https://code.activestate.com/recipes/121294-simple-generator-for-flattening-nested-containers/
+    and Recipe 1.12 in cookbook
+    """  # noqa: E501
+    for item in seq:
+        if scalarp(item) or item is None:
+            yield item
+        else:
+            yield from flatten(item, scalarp)
+
+
+class _Stack:
+    """
+    Stack of elements with a movable cursor.
+
+    Mimics home/back/forward in a web browser.
+    """
+
+    def __init__(self):
+        self._pos = -1
+        self._elements = []
+
+    def clear(self):
+        """Empty the stack."""
+        self._pos = -1
+        self._elements = []
+
+    def __call__(self):
+        """Return the current element, or None."""
+        return self._elements[self._pos] if self._elements else None
+
+    def __len__(self):
+        return len(self._elements)
+
+    def __getitem__(self, ind):
+        return self._elements[ind]
+
+    def forward(self):
+        """Move the position forward and return the current element."""
+        self._pos = min(self._pos + 1, len(self._elements) - 1)
+        return self()
+
+    def back(self):
+        """Move the position back and return the current element."""
+        self._pos = max(self._pos - 1, 0)
+        return self()
+
+    def push(self, o):
+        """
+        Push *o* to the stack after the current position, and return *o*.
+
+        Discard all later elements.
+        """
+        self._elements[self._pos + 1:] = [o]
+        self._pos = len(self._elements) - 1
+        return o
+
+    def home(self):
+        """
+        Push the first element onto the top of the stack.
+
+        The first element is returned.
+        """
+        return self.push(self._elements[0]) if self._elements else None
+
+
+def safe_masked_invalid(x, copy=False):
+    x = np.array(x, subok=True, copy=copy)
+    if not x.dtype.isnative:
+        # If we have already made a copy, do the byteswap in place, else make a
+        # copy with the byte order swapped.
+        # Swap to native order.
+        x = x.byteswap(inplace=copy).view(x.dtype.newbyteorder('N'))
+    try:
+        xm = np.ma.masked_where(~(np.isfinite(x)), x, copy=False)
+    except TypeError:
+        return x
+    return xm
+
+
+def print_cycles(objects, outstream=sys.stdout, show_progress=False):
+    """
+    Print loops of cyclic references in the given *objects*.
+
+    It is often useful to pass in ``gc.garbage`` to find the cycles that are
+    preventing some objects from being garbage collected.
+
+    Parameters
+    ----------
+    objects
+        A list of objects to find cycles in.
+    outstream
+        The stream for output.
+    show_progress : bool
+        If True, print the number of objects reached as they are found.
+    """
+    import gc
+
+    def print_path(path):
+        for i, step in enumerate(path):
+            # next "wraps around"
+            next = path[(i + 1) % len(path)]
+
+            outstream.write("   %s -- " % type(step))
+            if isinstance(step, dict):
+                for key, val in step.items():
+                    if val is next:
+                        outstream.write(f"[{key!r}]")
+                        break
+                    if key is next:
+                        outstream.write(f"[key] = {val!r}")
+                        break
+            elif isinstance(step, list):
+                outstream.write("[%d]" % step.index(next))
+            elif isinstance(step, tuple):
+                outstream.write("( tuple )")
+            else:
+                outstream.write(repr(step))
+            outstream.write(" ->\n")
+        outstream.write("\n")
+
+    def recurse(obj, start, all, current_path):
+        if show_progress:
+            outstream.write("%d\r" % len(all))
+
+        all[id(obj)] = None
+
+        referents = gc.get_referents(obj)
+        for referent in referents:
+            # If we've found our way back to the start, this is
+            # a cycle, so print it out
+            if referent is start:
+                print_path(current_path)
+
+            # Don't go back through the original list of objects, or
+            # through temporary references to the object, since those
+            # are just an artifact of the cycle detector itself.
+            elif referent is objects or isinstance(referent, types.FrameType):
+                continue
+
+            # We haven't seen this object before, so recurse
+            elif id(referent) not in all:
+                recurse(referent, start, all, current_path + [obj])
+
+    for obj in objects:
+        outstream.write(f"Examining: {obj!r}\n")
+        recurse(obj, obj, {}, [])
+
+
+class Grouper:
+    """
+    A disjoint-set data structure.
+
+    Objects can be joined using :meth:`join`, tested for connectedness
+    using :meth:`joined`, and all disjoint sets can be retrieved by
+    using the object as an iterator.
+
+    The objects being joined must be hashable and weak-referenceable.
+
+    Examples
+    --------
+    >>> from matplotlib.cbook import Grouper
+    >>> class Foo:
+    ...     def __init__(self, s):
+    ...         self.s = s
+    ...     def __repr__(self):
+    ...         return self.s
+    ...
+    >>> a, b, c, d, e, f = [Foo(x) for x in 'abcdef']
+    >>> grp = Grouper()
+    >>> grp.join(a, b)
+    >>> grp.join(b, c)
+    >>> grp.join(d, e)
+    >>> list(grp)
+    [[a, b, c], [d, e]]
+    >>> grp.joined(a, b)
+    True
+    >>> grp.joined(a, c)
+    True
+    >>> grp.joined(a, d)
+    False
+    """
+
+    def __init__(self, init=()):
+        self._mapping = weakref.WeakKeyDictionary(
+            {x: weakref.WeakSet([x]) for x in init})
+        self._ordering = weakref.WeakKeyDictionary()
+        for x in init:
+            if x not in self._ordering:
+                self._ordering[x] = len(self._ordering)
+        self._next_order = len(self._ordering)  # Plain int to simplify pickling.
+
+    def __getstate__(self):
+        return {
+            **vars(self),
+            # Convert weak refs to strong ones.
+            "_mapping": {k: set(v) for k, v in self._mapping.items()},
+            "_ordering": {**self._ordering},
+        }
+
+    def __setstate__(self, state):
+        vars(self).update(state)
+        # Convert strong refs to weak ones.
+        self._mapping = weakref.WeakKeyDictionary(
+            {k: weakref.WeakSet(v) for k, v in self._mapping.items()})
+        self._ordering = weakref.WeakKeyDictionary(self._ordering)
+
+    def __contains__(self, item):
+        return item in self._mapping
+
+    def join(self, a, *args):
+        """
+        Join given arguments into the same set.  Accepts one or more arguments.
+        """
+        mapping = self._mapping
+        try:
+            set_a = mapping[a]
+        except KeyError:
+            set_a = mapping[a] = weakref.WeakSet([a])
+            self._ordering[a] = self._next_order
+            self._next_order += 1
+        for arg in args:
+            try:
+                set_b = mapping[arg]
+            except KeyError:
+                set_b = mapping[arg] = weakref.WeakSet([arg])
+                self._ordering[arg] = self._next_order
+                self._next_order += 1
+            if set_b is not set_a:
+                if len(set_b) > len(set_a):
+                    set_a, set_b = set_b, set_a
+                set_a.update(set_b)
+                for elem in set_b:
+                    mapping[elem] = set_a
+
+    def joined(self, a, b):
+        """Return whether *a* and *b* are members of the same set."""
+        return (self._mapping.get(a, object()) is self._mapping.get(b))
+
+    def remove(self, a):
+        """Remove *a* from the grouper, doing nothing if it is not there."""
+        self._mapping.pop(a, {a}).remove(a)
+        self._ordering.pop(a, None)
+
+    def __iter__(self):
+        """
+        Iterate over each of the disjoint sets as a list.
+
+        The iterator is invalid if interleaved with calls to join().
+        """
+        unique_groups = {id(group): group for group in self._mapping.values()}
+        for group in unique_groups.values():
+            yield sorted(group, key=self._ordering.__getitem__)
+
+    def get_siblings(self, a):
+        """Return all of the items joined with *a*, including itself."""
+        siblings = self._mapping.get(a, [a])
+        return sorted(siblings, key=self._ordering.get)
+
+
+class GrouperView:
+    """Immutable view over a `.Grouper`."""
+
+    def __init__(self, grouper): self._grouper = grouper
+    def __contains__(self, item): return item in self._grouper
+    def __iter__(self): return iter(self._grouper)
+
+    def joined(self, a, b):
+        """
+        Return whether *a* and *b* are members of the same set.
+        """
+        return self._grouper.joined(a, b)
+
+    def get_siblings(self, a):
+        """
+        Return all of the items joined with *a*, including itself.
+        """
+        return self._grouper.get_siblings(a)
+
+
+def simple_linear_interpolation(a, steps):
+    """
+    Resample an array with ``steps - 1`` points between original point pairs.
+
+    Along each column of *a*, ``(steps - 1)`` points are introduced between
+    each original values; the values are linearly interpolated.
+
+    Parameters
+    ----------
+    a : array, shape (n, ...)
+    steps : int
+
+    Returns
+    -------
+    array
+        shape ``((n - 1) * steps + 1, ...)``
+    """
+    fps = a.reshape((len(a), -1))
+    xp = np.arange(len(a)) * steps
+    x = np.arange((len(a) - 1) * steps + 1)
+    return (np.column_stack([np.interp(x, xp, fp) for fp in fps.T])
+            .reshape((len(x),) + a.shape[1:]))
+
+
+def delete_masked_points(*args):
+    """
+    Find all masked and/or non-finite points in a set of arguments,
+    and return the arguments with only the unmasked points remaining.
+
+    Arguments can be in any of 5 categories:
+
+    1) 1-D masked arrays
+    2) 1-D ndarrays
+    3) ndarrays with more than one dimension
+    4) other non-string iterables
+    5) anything else
+
+    The first argument must be in one of the first four categories;
+    any argument with a length differing from that of the first
+    argument (and hence anything in category 5) then will be
+    passed through unchanged.
+
+    Masks are obtained from all arguments of the correct length
+    in categories 1, 2, and 4; a point is bad if masked in a masked
+    array or if it is a nan or inf.  No attempt is made to
+    extract a mask from categories 2, 3, and 4 if `numpy.isfinite`
+    does not yield a Boolean array.
+
+    All input arguments that are not passed unchanged are returned
+    as ndarrays after removing the points or rows corresponding to
+    masks in any of the arguments.
+
+    A vastly simpler version of this function was originally
+    written as a helper for Axes.scatter().
+
+    """
+    if not len(args):
+        return ()
+    if is_scalar_or_string(args[0]):
+        raise ValueError("First argument must be a sequence")
+    nrecs = len(args[0])
+    margs = []
+    seqlist = [False] * len(args)
+    for i, x in enumerate(args):
+        if not isinstance(x, str) and np.iterable(x) and len(x) == nrecs:
+            seqlist[i] = True
+            if isinstance(x, np.ma.MaskedArray):
+                if x.ndim > 1:
+                    raise ValueError("Masked arrays must be 1-D")
+            else:
+                x = np.asarray(x)
+        margs.append(x)
+    masks = []  # List of masks that are True where good.
+    for i, x in enumerate(margs):
+        if seqlist[i]:
+            if x.ndim > 1:
+                continue  # Don't try to get nan locations unless 1-D.
+            if isinstance(x, np.ma.MaskedArray):
+                masks.append(~np.ma.getmaskarray(x))  # invert the mask
+                xd = x.data
+            else:
+                xd = x
+            try:
+                mask = np.isfinite(xd)
+                if isinstance(mask, np.ndarray):
+                    masks.append(mask)
+            except Exception:  # Fixme: put in tuple of possible exceptions?
+                pass
+    if len(masks):
+        mask = np.logical_and.reduce(masks)
+        igood = mask.nonzero()[0]
+        if len(igood) < nrecs:
+            for i, x in enumerate(margs):
+                if seqlist[i]:
+                    margs[i] = x[igood]
+    for i, x in enumerate(margs):
+        if seqlist[i] and isinstance(x, np.ma.MaskedArray):
+            margs[i] = x.filled()
+    return margs
+
+
+def _combine_masks(*args):
+    """
+    Find all masked and/or non-finite points in a set of arguments,
+    and return the arguments as masked arrays with a common mask.
+
+    Arguments can be in any of 5 categories:
+
+    1) 1-D masked arrays
+    2) 1-D ndarrays
+    3) ndarrays with more than one dimension
+    4) other non-string iterables
+    5) anything else
+
+    The first argument must be in one of the first four categories;
+    any argument with a length differing from that of the first
+    argument (and hence anything in category 5) then will be
+    passed through unchanged.
+
+    Masks are obtained from all arguments of the correct length
+    in categories 1, 2, and 4; a point is bad if masked in a masked
+    array or if it is a nan or inf.  No attempt is made to
+    extract a mask from categories 2 and 4 if `numpy.isfinite`
+    does not yield a Boolean array.  Category 3 is included to
+    support RGB or RGBA ndarrays, which are assumed to have only
+    valid values and which are passed through unchanged.
+
+    All input arguments that are not passed unchanged are returned
+    as masked arrays if any masked points are found, otherwise as
+    ndarrays.
+
+    """
+    if not len(args):
+        return ()
+    if is_scalar_or_string(args[0]):
+        raise ValueError("First argument must be a sequence")
+    nrecs = len(args[0])
+    margs = []  # Output args; some may be modified.
+    seqlist = [False] * len(args)  # Flags: True if output will be masked.
+    masks = []    # List of masks.
+    for i, x in enumerate(args):
+        if is_scalar_or_string(x) or len(x) != nrecs:
+            margs.append(x)  # Leave it unmodified.
+        else:
+            if isinstance(x, np.ma.MaskedArray) and x.ndim > 1:
+                raise ValueError("Masked arrays must be 1-D")
+            try:
+                x = np.asanyarray(x)
+            except (VisibleDeprecationWarning, ValueError):
+                # NumPy 1.19 raises a warning about ragged arrays, but we want
+                # to accept basically anything here.
+                x = np.asanyarray(x, dtype=object)
+            if x.ndim == 1:
+                x = safe_masked_invalid(x)
+                seqlist[i] = True
+                if np.ma.is_masked(x):
+                    masks.append(np.ma.getmaskarray(x))
+            margs.append(x)  # Possibly modified.
+    if len(masks):
+        mask = np.logical_or.reduce(masks)
+        for i, x in enumerate(margs):
+            if seqlist[i]:
+                margs[i] = np.ma.array(x, mask=mask)
+    return margs
+
+
+def _broadcast_with_masks(*args, compress=False):
+    """
+    Broadcast inputs, combining all masked arrays.
+
+    Parameters
+    ----------
+    *args : array-like
+        The inputs to broadcast.
+    compress : bool, default: False
+        Whether to compress the masked arrays. If False, the masked values
+        are replaced by NaNs.
+
+    Returns
+    -------
+    list of array-like
+        The broadcasted and masked inputs.
+    """
+    # extract the masks, if any
+    masks = [k.mask for k in args if isinstance(k, np.ma.MaskedArray)]
+    # broadcast to match the shape
+    bcast = np.broadcast_arrays(*args, *masks)
+    inputs = bcast[:len(args)]
+    masks = bcast[len(args):]
+    if masks:
+        # combine the masks into one
+        mask = np.logical_or.reduce(masks)
+        # put mask on and compress
+        if compress:
+            inputs = [np.ma.array(k, mask=mask).compressed()
+                      for k in inputs]
+        else:
+            inputs = [np.ma.array(k, mask=mask, dtype=float).filled(np.nan).ravel()
+                      for k in inputs]
+    else:
+        inputs = [np.ravel(k) for k in inputs]
+    return inputs
+
+
+def boxplot_stats(X, whis=1.5, bootstrap=None, labels=None, autorange=False):
+    r"""
+    Return a list of dictionaries of statistics used to draw a series of box
+    and whisker plots using `~.Axes.bxp`.
+
+    Parameters
+    ----------
+    X : array-like
+        Data that will be represented in the boxplots. Should have 2 or
+        fewer dimensions.
+
+    whis : float or (float, float), default: 1.5
+        The position of the whiskers.
+
+        If a float, the lower whisker is at the lowest datum above
+        ``Q1 - whis*(Q3-Q1)``, and the upper whisker at the highest datum below
+        ``Q3 + whis*(Q3-Q1)``, where Q1 and Q3 are the first and third
+        quartiles.  The default value of ``whis = 1.5`` corresponds to Tukey's
+        original definition of boxplots.
+
+        If a pair of floats, they indicate the percentiles at which to draw the
+        whiskers (e.g., (5, 95)).  In particular, setting this to (0, 100)
+        results in whiskers covering the whole range of the data.
+
+        In the edge case where ``Q1 == Q3``, *whis* is automatically set to
+        (0, 100) (cover the whole range of the data) if *autorange* is True.
+
+        Beyond the whiskers, data are considered outliers and are plotted as
+        individual points.
+
+    bootstrap : int, optional
+        Number of times the confidence intervals around the median
+        should be bootstrapped (percentile method).
+
+    labels : list of str, optional
+        Labels for each dataset. Length must be compatible with
+        dimensions of *X*.
+
+    autorange : bool, optional (False)
+        When `True` and the data are distributed such that the 25th and 75th
+        percentiles are equal, ``whis`` is set to (0, 100) such that the
+        whisker ends are at the minimum and maximum of the data.
+
+    Returns
+    -------
+    list of dict
+        A list of dictionaries containing the results for each column
+        of data. Keys of each dictionary are the following:
+
+        ========   ===================================
+        Key        Value Description
+        ========   ===================================
+        label      tick label for the boxplot
+        mean       arithmetic mean value
+        med        50th percentile
+        q1         first quartile (25th percentile)
+        q3         third quartile (75th percentile)
+        iqr        interquartile range
+        cilo       lower notch around the median
+        cihi       upper notch around the median
+        whislo     end of the lower whisker
+        whishi     end of the upper whisker
+        fliers     outliers
+        ========   ===================================
+
+    Notes
+    -----
+    Non-bootstrapping approach to confidence interval uses Gaussian-based
+    asymptotic approximation:
+
+    .. math::
+
+        \mathrm{med} \pm 1.57 \times \frac{\mathrm{iqr}}{\sqrt{N}}
+
+    General approach from:
+    McGill, R., Tukey, J.W., and Larsen, W.A. (1978) "Variations of
+    Boxplots", The American Statistician, 32:12-16.
+    """
+
+    def _bootstrap_median(data, N=5000):
+        # determine 95% confidence intervals of the median
+        M = len(data)
+        percentiles = [2.5, 97.5]
+
+        bs_index = np.random.randint(M, size=(N, M))
+        bsData = data[bs_index]
+        estimate = np.median(bsData, axis=1, overwrite_input=True)
+
+        CI = np.percentile(estimate, percentiles)
+        return CI
+
+    def _compute_conf_interval(data, med, iqr, bootstrap):
+        if bootstrap is not None:
+            # Do a bootstrap estimate of notch locations.
+            # get conf. intervals around median
+            CI = _bootstrap_median(data, N=bootstrap)
+            notch_min = CI[0]
+            notch_max = CI[1]
+        else:
+
+            N = len(data)
+            notch_min = med - 1.57 * iqr / np.sqrt(N)
+            notch_max = med + 1.57 * iqr / np.sqrt(N)
+
+        return notch_min, notch_max
+
+    # output is a list of dicts
+    bxpstats = []
+
+    # convert X to a list of lists
+    X = _reshape_2D(X, "X")
+
+    ncols = len(X)
+    if labels is None:
+        labels = itertools.repeat(None)
+    elif len(labels) != ncols:
+        raise ValueError("Dimensions of labels and X must be compatible")
+
+    input_whis = whis
+    for ii, (x, label) in enumerate(zip(X, labels)):
+
+        # empty dict
+        stats = {}
+        if label is not None:
+            stats['label'] = label
+
+        # restore whis to the input values in case it got changed in the loop
+        whis = input_whis
+
+        # note tricksiness, append up here and then mutate below
+        bxpstats.append(stats)
+
+        # if empty, bail
+        if len(x) == 0:
+            stats['fliers'] = np.array([])
+            stats['mean'] = np.nan
+            stats['med'] = np.nan
+            stats['q1'] = np.nan
+            stats['q3'] = np.nan
+            stats['iqr'] = np.nan
+            stats['cilo'] = np.nan
+            stats['cihi'] = np.nan
+            stats['whislo'] = np.nan
+            stats['whishi'] = np.nan
+            continue
+
+        # up-convert to an array, just to be safe
+        x = np.ma.asarray(x)
+        x = x.data[~x.mask].ravel()
+
+        # arithmetic mean
+        stats['mean'] = np.mean(x)
+
+        # medians and quartiles
+        q1, med, q3 = np.percentile(x, [25, 50, 75])
+
+        # interquartile range
+        stats['iqr'] = q3 - q1
+        if stats['iqr'] == 0 and autorange:
+            whis = (0, 100)
+
+        # conf. interval around median
+        stats['cilo'], stats['cihi'] = _compute_conf_interval(
+            x, med, stats['iqr'], bootstrap
+        )
+
+        # lowest/highest non-outliers
+        if np.iterable(whis) and not isinstance(whis, str):
+            loval, hival = np.percentile(x, whis)
+        elif np.isreal(whis):
+            loval = q1 - whis * stats['iqr']
+            hival = q3 + whis * stats['iqr']
+        else:
+            raise ValueError('whis must be a float or list of percentiles')
+
+        # get high extreme
+        wiskhi = x[x <= hival]
+        if len(wiskhi) == 0 or np.max(wiskhi) < q3:
+            stats['whishi'] = q3
+        else:
+            stats['whishi'] = np.max(wiskhi)
+
+        # get low extreme
+        wisklo = x[x >= loval]
+        if len(wisklo) == 0 or np.min(wisklo) > q1:
+            stats['whislo'] = q1
+        else:
+            stats['whislo'] = np.min(wisklo)
+
+        # compute a single array of outliers
+        stats['fliers'] = np.concatenate([
+            x[x < stats['whislo']],
+            x[x > stats['whishi']],
+        ])
+
+        # add in the remaining stats
+        stats['q1'], stats['med'], stats['q3'] = q1, med, q3
+
+    return bxpstats
+
+
+#: Maps short codes for line style to their full name used by backends.
+ls_mapper = {'-': 'solid', '--': 'dashed', '-.': 'dashdot', ':': 'dotted'}
+#: Maps full names for line styles used by backends to their short codes.
+ls_mapper_r = {v: k for k, v in ls_mapper.items()}
+
+
+def contiguous_regions(mask):
+    """
+    Return a list of (ind0, ind1) such that ``mask[ind0:ind1].all()`` is
+    True and we cover all such regions.
+    """
+    mask = np.asarray(mask, dtype=bool)
+
+    if not mask.size:
+        return []
+
+    # Find the indices of region changes, and correct offset
+    idx, = np.nonzero(mask[:-1] != mask[1:])
+    idx += 1
+
+    # List operations are faster for moderately sized arrays
+    idx = idx.tolist()
+
+    # Add first and/or last index if needed
+    if mask[0]:
+        idx = [0] + idx
+    if mask[-1]:
+        idx.append(len(mask))
+
+    return list(zip(idx[::2], idx[1::2]))
+
+
+def is_math_text(s):
+    """
+    Return whether the string *s* contains math expressions.
+
+    This is done by checking whether *s* contains an even number of
+    non-escaped dollar signs.
+    """
+    s = str(s)
+    dollar_count = s.count(r'$') - s.count(r'\$')
+    even_dollars = (dollar_count > 0 and dollar_count % 2 == 0)
+    return even_dollars
+
+
+def _to_unmasked_float_array(x):
+    """
+    Convert a sequence to a float array; if input was a masked array, masked
+    values are converted to nans.
+    """
+    if hasattr(x, 'mask'):
+        return np.ma.asarray(x, float).filled(np.nan)
+    else:
+        return np.asarray(x, float)
+
+
+def _check_1d(x):
+    """Convert scalars to 1D arrays; pass-through arrays as is."""
+    # Unpack in case of e.g. Pandas or xarray object
+    x = _unpack_to_numpy(x)
+    # plot requires `shape` and `ndim`.  If passed an
+    # object that doesn't provide them, then force to numpy array.
+    # Note this will strip unit information.
+    if (not hasattr(x, 'shape') or
+            not hasattr(x, 'ndim') or
+            len(x.shape) < 1):
+        return np.atleast_1d(x)
+    else:
+        return x
+
+
+def _reshape_2D(X, name):
+    """
+    Use Fortran ordering to convert ndarrays and lists of iterables to lists of
+    1D arrays.
+
+    Lists of iterables are converted by applying `numpy.asanyarray` to each of
+    their elements.  1D ndarrays are returned in a singleton list containing
+    them.  2D ndarrays are converted to the list of their *columns*.
+
+    *name* is used to generate the error message for invalid inputs.
+    """
+
+    # Unpack in case of e.g. Pandas or xarray object
+    X = _unpack_to_numpy(X)
+
+    # Iterate over columns for ndarrays.
+    if isinstance(X, np.ndarray):
+        X = X.transpose()
+
+        if len(X) == 0:
+            return [[]]
+        elif X.ndim == 1 and np.ndim(X[0]) == 0:
+            # 1D array of scalars: directly return it.
+            return [X]
+        elif X.ndim in [1, 2]:
+            # 2D array, or 1D array of iterables: flatten them first.
+            return [np.reshape(x, -1) for x in X]
+        else:
+            raise ValueError(f'{name} must have 2 or fewer dimensions')
+
+    # Iterate over list of iterables.
+    if len(X) == 0:
+        return [[]]
+
+    result = []
+    is_1d = True
+    for xi in X:
+        # check if this is iterable, except for strings which we
+        # treat as singletons.
+        if not isinstance(xi, str):
+            try:
+                iter(xi)
+            except TypeError:
+                pass
+            else:
+                is_1d = False
+        xi = np.asanyarray(xi)
+        nd = np.ndim(xi)
+        if nd > 1:
+            raise ValueError(f'{name} must have 2 or fewer dimensions')
+        result.append(xi.reshape(-1))
+
+    if is_1d:
+        # 1D array of scalars: directly return it.
+        return [np.reshape(result, -1)]
+    else:
+        # 2D array, or 1D array of iterables: use flattened version.
+        return result
+
+
+def violin_stats(X, method, points=100, quantiles=None):
+    """
+    Return a list of dictionaries of data which can be used to draw a series
+    of violin plots.
+
+    See the ``Returns`` section below to view the required keys of the
+    dictionary.
+
+    Users can skip this function and pass a user-defined set of dictionaries
+    with the same keys to `~.axes.Axes.violinplot` instead of using Matplotlib
+    to do the calculations. See the *Returns* section below for the keys
+    that must be present in the dictionaries.
+
+    Parameters
+    ----------
+    X : array-like
+        Sample data that will be used to produce the gaussian kernel density
+        estimates. Must have 2 or fewer dimensions.
+
+    method : callable
+        The method used to calculate the kernel density estimate for each
+        column of data. When called via ``method(v, coords)``, it should
+        return a vector of the values of the KDE evaluated at the values
+        specified in coords.
+
+    points : int, default: 100
+        Defines the number of points to evaluate each of the gaussian kernel
+        density estimates at.
+
+    quantiles : array-like, default: None
+        Defines (if not None) a list of floats in interval [0, 1] for each
+        column of data, which represents the quantiles that will be rendered
+        for that column of data. Must have 2 or fewer dimensions. 1D array will
+        be treated as a singleton list containing them.
+
+    Returns
+    -------
+    list of dict
+        A list of dictionaries containing the results for each column of data.
+        The dictionaries contain at least the following:
+
+        - coords: A list of scalars containing the coordinates this particular
+          kernel density estimate was evaluated at.
+        - vals: A list of scalars containing the values of the kernel density
+          estimate at each of the coordinates given in *coords*.
+        - mean: The mean value for this column of data.
+        - median: The median value for this column of data.
+        - min: The minimum value for this column of data.
+        - max: The maximum value for this column of data.
+        - quantiles: The quantile values for this column of data.
+    """
+
+    # List of dictionaries describing each of the violins.
+    vpstats = []
+
+    # Want X to be a list of data sequences
+    X = _reshape_2D(X, "X")
+
+    # Want quantiles to be as the same shape as data sequences
+    if quantiles is not None and len(quantiles) != 0:
+        quantiles = _reshape_2D(quantiles, "quantiles")
+    # Else, mock quantiles if it's none or empty
+    else:
+        quantiles = [[]] * len(X)
+
+    # quantiles should have the same size as dataset
+    if len(X) != len(quantiles):
+        raise ValueError("List of violinplot statistics and quantiles values"
+                         " must have the same length")
+
+    # Zip x and quantiles
+    for (x, q) in zip(X, quantiles):
+        # Dictionary of results for this distribution
+        stats = {}
+
+        # Calculate basic stats for the distribution
+        min_val = np.min(x)
+        max_val = np.max(x)
+        quantile_val = np.percentile(x, 100 * q)
+
+        # Evaluate the kernel density estimate
+        coords = np.linspace(min_val, max_val, points)
+        stats['vals'] = method(x, coords)
+        stats['coords'] = coords
+
+        # Store additional statistics for this distribution
+        stats['mean'] = np.mean(x)
+        stats['median'] = np.median(x)
+        stats['min'] = min_val
+        stats['max'] = max_val
+        stats['quantiles'] = np.atleast_1d(quantile_val)
+
+        # Append to output
+        vpstats.append(stats)
+
+    return vpstats
+
+
+def pts_to_prestep(x, *args):
+    """
+    Convert continuous line to pre-steps.
+
+    Given a set of ``N`` points, convert to ``2N - 1`` points, which when
+    connected linearly give a step function which changes values at the
+    beginning of the intervals.
+
+    Parameters
+    ----------
+    x : array
+        The x location of the steps. May be empty.
+
+    y1, ..., yp : array
+        y arrays to be turned into steps; all must be the same length as ``x``.
+
+    Returns
+    -------
+    array
+        The x and y values converted to steps in the same order as the input;
+        can be unpacked as ``x_out, y1_out, ..., yp_out``.  If the input is
+        length ``N``, each of these arrays will be length ``2N + 1``. For
+        ``N=0``, the length will be 0.
+
+    Examples
+    --------
+    >>> x_s, y1_s, y2_s = pts_to_prestep(x, y1, y2)
+    """
+    steps = np.zeros((1 + len(args), max(2 * len(x) - 1, 0)))
+    # In all `pts_to_*step` functions, only assign once using *x* and *args*,
+    # as converting to an array may be expensive.
+    steps[0, 0::2] = x
+    steps[0, 1::2] = steps[0, 0:-2:2]
+    steps[1:, 0::2] = args
+    steps[1:, 1::2] = steps[1:, 2::2]
+    return steps
+
+
+def pts_to_poststep(x, *args):
+    """
+    Convert continuous line to post-steps.
+
+    Given a set of ``N`` points convert to ``2N + 1`` points, which when
+    connected linearly give a step function which changes values at the end of
+    the intervals.
+
+    Parameters
+    ----------
+    x : array
+        The x location of the steps. May be empty.
+
+    y1, ..., yp : array
+        y arrays to be turned into steps; all must be the same length as ``x``.
+
+    Returns
+    -------
+    array
+        The x and y values converted to steps in the same order as the input;
+        can be unpacked as ``x_out, y1_out, ..., yp_out``.  If the input is
+        length ``N``, each of these arrays will be length ``2N + 1``. For
+        ``N=0``, the length will be 0.
+
+    Examples
+    --------
+    >>> x_s, y1_s, y2_s = pts_to_poststep(x, y1, y2)
+    """
+    steps = np.zeros((1 + len(args), max(2 * len(x) - 1, 0)))
+    steps[0, 0::2] = x
+    steps[0, 1::2] = steps[0, 2::2]
+    steps[1:, 0::2] = args
+    steps[1:, 1::2] = steps[1:, 0:-2:2]
+    return steps
+
+
+def pts_to_midstep(x, *args):
+    """
+    Convert continuous line to mid-steps.
+
+    Given a set of ``N`` points convert to ``2N`` points which when connected
+    linearly give a step function which changes values at the middle of the
+    intervals.
+
+    Parameters
+    ----------
+    x : array
+        The x location of the steps. May be empty.
+
+    y1, ..., yp : array
+        y arrays to be turned into steps; all must be the same length as
+        ``x``.
+
+    Returns
+    -------
+    array
+        The x and y values converted to steps in the same order as the input;
+        can be unpacked as ``x_out, y1_out, ..., yp_out``.  If the input is
+        length ``N``, each of these arrays will be length ``2N``.
+
+    Examples
+    --------
+    >>> x_s, y1_s, y2_s = pts_to_midstep(x, y1, y2)
+    """
+    steps = np.zeros((1 + len(args), 2 * len(x)))
+    x = np.asanyarray(x)
+    steps[0, 1:-1:2] = steps[0, 2::2] = (x[:-1] + x[1:]) / 2
+    steps[0, :1] = x[:1]  # Also works for zero-sized input.
+    steps[0, -1:] = x[-1:]
+    steps[1:, 0::2] = args
+    steps[1:, 1::2] = steps[1:, 0::2]
+    return steps
+
+
+STEP_LOOKUP_MAP = {'default': lambda x, y: (x, y),
+                   'steps': pts_to_prestep,
+                   'steps-pre': pts_to_prestep,
+                   'steps-post': pts_to_poststep,
+                   'steps-mid': pts_to_midstep}
+
+
+def index_of(y):
+    """
+    A helper function to create reasonable x values for the given *y*.
+
+    This is used for plotting (x, y) if x values are not explicitly given.
+
+    First try ``y.index`` (assuming *y* is a `pandas.Series`), if that
+    fails, use ``range(len(y))``.
+
+    This will be extended in the future to deal with more types of
+    labeled data.
+
+    Parameters
+    ----------
+    y : float or array-like
+
+    Returns
+    -------
+    x, y : ndarray
+       The x and y values to plot.
+    """
+    try:
+        return y.index.to_numpy(), y.to_numpy()
+    except AttributeError:
+        pass
+    try:
+        y = _check_1d(y)
+    except (VisibleDeprecationWarning, ValueError):
+        # NumPy 1.19 will warn on ragged input, and we can't actually use it.
+        pass
+    else:
+        return np.arange(y.shape[0], dtype=float), y
+    raise ValueError('Input could not be cast to an at-least-1D NumPy array')
+
+
+def safe_first_element(obj):
+    """
+    Return the first element in *obj*.
+
+    This is a type-independent way of obtaining the first element,
+    supporting both index access and the iterator protocol.
+    """
+    if isinstance(obj, collections.abc.Iterator):
+        # needed to accept `array.flat` as input.
+        # np.flatiter reports as an instance of collections.Iterator but can still be
+        # indexed via []. This has the side effect of re-setting the iterator, but
+        # that is acceptable.
+        try:
+            return obj[0]
+        except TypeError:
+            pass
+        raise RuntimeError("matplotlib does not support generators as input")
+    return next(iter(obj))
+
+
+def _safe_first_finite(obj):
+    """
+    Return the first finite element in *obj* if one is available and skip_nonfinite is
+    True. Otherwise, return the first element.
+
+    This is a method for internal use.
+
+    This is a type-independent way of obtaining the first finite element, supporting
+    both index access and the iterator protocol.
+    """
+    def safe_isfinite(val):
+        if val is None:
+            return False
+        try:
+            return math.isfinite(val)
+        except (TypeError, ValueError):
+            # if the outer object is 2d, then val is a 1d array, and
+            # - math.isfinite(numpy.zeros(3)) raises TypeError
+            # - math.isfinite(torch.zeros(3)) raises ValueError
+            pass
+        try:
+            return np.isfinite(val) if np.isscalar(val) else True
+        except TypeError:
+            # This is something that NumPy cannot make heads or tails of,
+            # assume "finite"
+            return True
+
+    if isinstance(obj, np.flatiter):
+        # TODO do the finite filtering on this
+        return obj[0]
+    elif isinstance(obj, collections.abc.Iterator):
+        raise RuntimeError("matplotlib does not support generators as input")
+    else:
+        for val in obj:
+            if safe_isfinite(val):
+                return val
+        return safe_first_element(obj)
+
+
+def sanitize_sequence(data):
+    """
+    Convert dictview objects to list. Other inputs are returned unchanged.
+    """
+    return (list(data) if isinstance(data, collections.abc.MappingView)
+            else data)
+
+
+def normalize_kwargs(kw, alias_mapping=None):
+    """
+    Helper function to normalize kwarg inputs.
+
+    Parameters
+    ----------
+    kw : dict or None
+        A dict of keyword arguments.  None is explicitly supported and treated
+        as an empty dict, to support functions with an optional parameter of
+        the form ``props=None``.
+
+    alias_mapping : dict or Artist subclass or Artist instance, optional
+        A mapping between a canonical name to a list of aliases, in order of
+        precedence from lowest to highest.
+
+        If the canonical value is not in the list it is assumed to have the
+        highest priority.
+
+        If an Artist subclass or instance is passed, use its properties alias
+        mapping.
+
+    Raises
+    ------
+    TypeError
+        To match what Python raises if invalid arguments/keyword arguments are
+        passed to a callable.
+    """
+    from matplotlib.artist import Artist
+
+    if kw is None:
+        return {}
+
+    # deal with default value of alias_mapping
+    if alias_mapping is None:
+        alias_mapping = {}
+    elif (isinstance(alias_mapping, type) and issubclass(alias_mapping, Artist)
+          or isinstance(alias_mapping, Artist)):
+        alias_mapping = getattr(alias_mapping, "_alias_map", {})
+
+    to_canonical = {alias: canonical
+                    for canonical, alias_list in alias_mapping.items()
+                    for alias in alias_list}
+    canonical_to_seen = {}
+    ret = {}  # output dictionary
+
+    for k, v in kw.items():
+        canonical = to_canonical.get(k, k)
+        if canonical in canonical_to_seen:
+            raise TypeError(f"Got both {canonical_to_seen[canonical]!r} and "
+                            f"{k!r}, which are aliases of one another")
+        canonical_to_seen[canonical] = k
+        ret[canonical] = v
+
+    return ret
+
+
+@contextlib.contextmanager
+def _lock_path(path):
+    """
+    Context manager for locking a path.
+
+    Usage::
+
+        with _lock_path(path):
+            ...
+
+    Another thread or process that attempts to lock the same path will wait
+    until this context manager is exited.
+
+    The lock is implemented by creating a temporary file in the parent
+    directory, so that directory must exist and be writable.
+    """
+    path = Path(path)
+    lock_path = path.with_name(path.name + ".matplotlib-lock")
+    retries = 50
+    sleeptime = 0.1
+    for _ in range(retries):
+        try:
+            with lock_path.open("xb"):
+                break
+        except FileExistsError:
+            time.sleep(sleeptime)
+    else:
+        raise TimeoutError("""\
+Lock error: Matplotlib failed to acquire the following lock file:
+    {}
+This maybe due to another process holding this lock file.  If you are sure no
+other Matplotlib process is running, remove this file and try again.""".format(
+            lock_path))
+    try:
+        yield
+    finally:
+        lock_path.unlink()
+
+
+def _topmost_artist(
+        artists,
+        _cached_max=functools.partial(max, key=operator.attrgetter("zorder"))):
+    """
+    Get the topmost artist of a list.
+
+    In case of a tie, return the *last* of the tied artists, as it will be
+    drawn on top of the others. `max` returns the first maximum in case of
+    ties, so we need to iterate over the list in reverse order.
+    """
+    return _cached_max(reversed(artists))
+
+
+def _str_equal(obj, s):
+    """
+    Return whether *obj* is a string equal to string *s*.
+
+    This helper solely exists to handle the case where *obj* is a numpy array,
+    because in such cases, a naive ``obj == s`` would yield an array, which
+    cannot be used in a boolean context.
+    """
+    return isinstance(obj, str) and obj == s
+
+
+def _str_lower_equal(obj, s):
+    """
+    Return whether *obj* is a string equal, when lowercased, to string *s*.
+
+    This helper solely exists to handle the case where *obj* is a numpy array,
+    because in such cases, a naive ``obj == s`` would yield an array, which
+    cannot be used in a boolean context.
+    """
+    return isinstance(obj, str) and obj.lower() == s
+
+
+def _array_perimeter(arr):
+    """
+    Get the elements on the perimeter of *arr*.
+
+    Parameters
+    ----------
+    arr : ndarray, shape (M, N)
+        The input array.
+
+    Returns
+    -------
+    ndarray, shape (2*(M - 1) + 2*(N - 1),)
+        The elements on the perimeter of the array::
+
+           [arr[0, 0], ..., arr[0, -1], ..., arr[-1, -1], ..., arr[-1, 0], ...]
+
+    Examples
+    --------
+    >>> i, j = np.ogrid[:3, :4]
+    >>> a = i*10 + j
+    >>> a
+    array([[ 0,  1,  2,  3],
+           [10, 11, 12, 13],
+           [20, 21, 22, 23]])
+    >>> _array_perimeter(a)
+    array([ 0,  1,  2,  3, 13, 23, 22, 21, 20, 10])
+    """
+    # note we use Python's half-open ranges to avoid repeating
+    # the corners
+    forward = np.s_[0:-1]      # [0 ... -1)
+    backward = np.s_[-1:0:-1]  # [-1 ... 0)
+    return np.concatenate((
+        arr[0, forward],
+        arr[forward, -1],
+        arr[-1, backward],
+        arr[backward, 0],
+    ))
+
+
+def _unfold(arr, axis, size, step):
+    """
+    Append an extra dimension containing sliding windows along *axis*.
+
+    All windows are of size *size* and begin with every *step* elements.
+
+    Parameters
+    ----------
+    arr : ndarray, shape (N_1, ..., N_k)
+        The input array
+    axis : int
+        Axis along which the windows are extracted
+    size : int
+        Size of the windows
+    step : int
+        Stride between first elements of subsequent windows.
+
+    Returns
+    -------
+    ndarray, shape (N_1, ..., 1 + (N_axis-size)/step, ..., N_k, size)
+
+    Examples
+    --------
+    >>> i, j = np.ogrid[:3, :7]
+    >>> a = i*10 + j
+    >>> a
+    array([[ 0,  1,  2,  3,  4,  5,  6],
+           [10, 11, 12, 13, 14, 15, 16],
+           [20, 21, 22, 23, 24, 25, 26]])
+    >>> _unfold(a, axis=1, size=3, step=2)
+    array([[[ 0,  1,  2],
+            [ 2,  3,  4],
+            [ 4,  5,  6]],
+           [[10, 11, 12],
+            [12, 13, 14],
+            [14, 15, 16]],
+           [[20, 21, 22],
+            [22, 23, 24],
+            [24, 25, 26]]])
+    """
+    new_shape = [*arr.shape, size]
+    new_strides = [*arr.strides, arr.strides[axis]]
+    new_shape[axis] = (new_shape[axis] - size) // step + 1
+    new_strides[axis] = new_strides[axis] * step
+    return np.lib.stride_tricks.as_strided(arr,
+                                           shape=new_shape,
+                                           strides=new_strides,
+                                           writeable=False)
+
+
+def _array_patch_perimeters(x, rstride, cstride):
+    """
+    Extract perimeters of patches from *arr*.
+
+    Extracted patches are of size (*rstride* + 1) x (*cstride* + 1) and
+    share perimeters with their neighbors. The ordering of the vertices matches
+    that returned by ``_array_perimeter``.
+
+    Parameters
+    ----------
+    x : ndarray, shape (N, M)
+        Input array
+    rstride : int
+        Vertical (row) stride between corresponding elements of each patch
+    cstride : int
+        Horizontal (column) stride between corresponding elements of each patch
+
+    Returns
+    -------
+    ndarray, shape (N/rstride * M/cstride, 2 * (rstride + cstride))
+    """
+    assert rstride > 0 and cstride > 0
+    assert (x.shape[0] - 1) % rstride == 0
+    assert (x.shape[1] - 1) % cstride == 0
+    # We build up each perimeter from four half-open intervals. Here is an
+    # illustrated explanation for rstride == cstride == 3
+    #
+    #       T T T R
+    #       L     R
+    #       L     R
+    #       L B B B
+    #
+    # where T means that this element will be in the top array, R for right,
+    # B for bottom and L for left. Each of the arrays below has a shape of:
+    #
+    #    (number of perimeters that can be extracted vertically,
+    #     number of perimeters that can be extracted horizontally,
+    #     cstride for top and bottom and rstride for left and right)
+    #
+    # Note that _unfold doesn't incur any memory copies, so the only costly
+    # operation here is the np.concatenate.
+    top = _unfold(x[:-1:rstride, :-1], 1, cstride, cstride)
+    bottom = _unfold(x[rstride::rstride, 1:], 1, cstride, cstride)[..., ::-1]
+    right = _unfold(x[:-1, cstride::cstride], 0, rstride, rstride)
+    left = _unfold(x[1:, :-1:cstride], 0, rstride, rstride)[..., ::-1]
+    return (np.concatenate((top, right, bottom, left), axis=2)
+              .reshape(-1, 2 * (rstride + cstride)))
+
+
+@contextlib.contextmanager
+def _setattr_cm(obj, **kwargs):
+    """
+    Temporarily set some attributes; restore original state at context exit.
+    """
+    sentinel = object()
+    origs = {}
+    for attr in kwargs:
+        orig = getattr(obj, attr, sentinel)
+        if attr in obj.__dict__ or orig is sentinel:
+            # if we are pulling from the instance dict or the object
+            # does not have this attribute we can trust the above
+            origs[attr] = orig
+        else:
+            # if the attribute is not in the instance dict it must be
+            # from the class level
+            cls_orig = getattr(type(obj), attr)
+            # if we are dealing with a property (but not a general descriptor)
+            # we want to set the original value back.
+            if isinstance(cls_orig, property):
+                origs[attr] = orig
+            # otherwise this is _something_ we are going to shadow at
+            # the instance dict level from higher up in the MRO.  We
+            # are going to assume we can delattr(obj, attr) to clean
+            # up after ourselves.  It is possible that this code will
+            # fail if used with a non-property custom descriptor which
+            # implements __set__ (and __delete__ does not act like a
+            # stack).  However, this is an internal tool and we do not
+            # currently have any custom descriptors.
+            else:
+                origs[attr] = sentinel
+
+    try:
+        for attr, val in kwargs.items():
+            setattr(obj, attr, val)
+        yield
+    finally:
+        for attr, orig in origs.items():
+            if orig is sentinel:
+                delattr(obj, attr)
+            else:
+                setattr(obj, attr, orig)
+
+
+class _OrderedSet(collections.abc.MutableSet):
+    def __init__(self):
+        self._od = collections.OrderedDict()
+
+    def __contains__(self, key):
+        return key in self._od
+
+    def __iter__(self):
+        return iter(self._od)
+
+    def __len__(self):
+        return len(self._od)
+
+    def add(self, key):
+        self._od.pop(key, None)
+        self._od[key] = None
+
+    def discard(self, key):
+        self._od.pop(key, None)
+
+
+# Agg's buffers are unmultiplied RGBA8888, which neither PyQt<=5.1 nor cairo
+# support; however, both do support premultiplied ARGB32.
+
+
+def _premultiplied_argb32_to_unmultiplied_rgba8888(buf):
+    """
+    Convert a premultiplied ARGB32 buffer to an unmultiplied RGBA8888 buffer.
+    """
+    rgba = np.take(  # .take() ensures C-contiguity of the result.
+        buf,
+        [2, 1, 0, 3] if sys.byteorder == "little" else [1, 2, 3, 0], axis=2)
+    rgb = rgba[..., :-1]
+    alpha = rgba[..., -1]
+    # Un-premultiply alpha.  The formula is the same as in cairo-png.c.
+    mask = alpha != 0
+    for channel in np.rollaxis(rgb, -1):
+        channel[mask] = (
+            (channel[mask].astype(int) * 255 + alpha[mask] // 2)
+            // alpha[mask])
+    return rgba
+
+
+def _unmultiplied_rgba8888_to_premultiplied_argb32(rgba8888):
+    """
+    Convert an unmultiplied RGBA8888 buffer to a premultiplied ARGB32 buffer.
+    """
+    if sys.byteorder == "little":
+        argb32 = np.take(rgba8888, [2, 1, 0, 3], axis=2)
+        rgb24 = argb32[..., :-1]
+        alpha8 = argb32[..., -1:]
+    else:
+        argb32 = np.take(rgba8888, [3, 0, 1, 2], axis=2)
+        alpha8 = argb32[..., :1]
+        rgb24 = argb32[..., 1:]
+    # Only bother premultiplying when the alpha channel is not fully opaque,
+    # as the cost is not negligible.  The unsafe cast is needed to do the
+    # multiplication in-place in an integer buffer.
+    if alpha8.min() != 0xff:
+        np.multiply(rgb24, alpha8 / 0xff, out=rgb24, casting="unsafe")
+    return argb32
+
+
+def _get_nonzero_slices(buf):
+    """
+    Return the bounds of the nonzero region of a 2D array as a pair of slices.
+
+    ``buf[_get_nonzero_slices(buf)]`` is the smallest sub-rectangle in *buf*
+    that encloses all non-zero entries in *buf*.  If *buf* is fully zero, then
+    ``(slice(0, 0), slice(0, 0))`` is returned.
+    """
+    x_nz, = buf.any(axis=0).nonzero()
+    y_nz, = buf.any(axis=1).nonzero()
+    if len(x_nz) and len(y_nz):
+        l, r = x_nz[[0, -1]]
+        b, t = y_nz[[0, -1]]
+        return slice(b, t + 1), slice(l, r + 1)
+    else:
+        return slice(0, 0), slice(0, 0)
+
+
+def _pformat_subprocess(command):
+    """Pretty-format a subprocess command for printing/logging purposes."""
+    return (command if isinstance(command, str)
+            else " ".join(shlex.quote(os.fspath(arg)) for arg in command))
+
+
+def _check_and_log_subprocess(command, logger, **kwargs):
+    """
+    Run *command*, returning its stdout output if it succeeds.
+
+    If it fails (exits with nonzero return code), raise an exception whose text
+    includes the failed command and captured stdout and stderr output.
+
+    Regardless of the return code, the command is logged at DEBUG level on
+    *logger*.  In case of success, the output is likewise logged.
+    """
+    logger.debug('%s', _pformat_subprocess(command))
+    proc = subprocess.run(command, capture_output=True, **kwargs)
+    if proc.returncode:
+        stdout = proc.stdout
+        if isinstance(stdout, bytes):
+            stdout = stdout.decode()
+        stderr = proc.stderr
+        if isinstance(stderr, bytes):
+            stderr = stderr.decode()
+        raise RuntimeError(
+            f"The command\n"
+            f"    {_pformat_subprocess(command)}\n"
+            f"failed and generated the following output:\n"
+            f"{stdout}\n"
+            f"and the following error:\n"
+            f"{stderr}")
+    if proc.stdout:
+        logger.debug("stdout:\n%s", proc.stdout)
+    if proc.stderr:
+        logger.debug("stderr:\n%s", proc.stderr)
+    return proc.stdout
+
+
+def _setup_new_guiapp():
+    """
+    Perform OS-dependent setup when Matplotlib creates a new GUI application.
+    """
+    # Windows: If not explicit app user model id has been set yet (so we're not
+    # already embedded), then set it to "matplotlib", so that taskbar icons are
+    # correct.
+    try:
+        _c_internal_utils.Win32_GetCurrentProcessExplicitAppUserModelID()
+    except OSError:
+        _c_internal_utils.Win32_SetCurrentProcessExplicitAppUserModelID(
+            "matplotlib")
+
+
+def _format_approx(number, precision):
+    """
+    Format the number with at most the number of decimals given as precision.
+    Remove trailing zeros and possibly the decimal point.
+    """
+    return f'{number:.{precision}f}'.rstrip('0').rstrip('.') or '0'
+
+
+def _g_sig_digits(value, delta):
+    """
+    Return the number of significant digits to %g-format *value*, assuming that
+    it is known with an error of *delta*.
+    """
+    if delta == 0:
+        if value == 0:
+            # if both value and delta are 0, np.spacing below returns 5e-324
+            # which results in rather silly results
+            return 3
+        # delta = 0 may occur when trying to format values over a tiny range;
+        # in that case, replace it by the distance to the closest float.
+        delta = abs(np.spacing(value))
+    # If e.g. value = 45.67 and delta = 0.02, then we want to round to 2 digits
+    # after the decimal point (floor(log10(0.02)) = -2); 45.67 contributes 2
+    # digits before the decimal point (floor(log10(45.67)) + 1 = 2): the total
+    # is 4 significant digits.  A value of 0 contributes 1 "digit" before the
+    # decimal point.
+    # For inf or nan, the precision doesn't matter.
+    return max(
+        0,
+        (math.floor(math.log10(abs(value))) + 1 if value else 1)
+        - math.floor(math.log10(delta))) if math.isfinite(value) else 0
+
+
+def _unikey_or_keysym_to_mplkey(unikey, keysym):
+    """
+    Convert a Unicode key or X keysym to a Matplotlib key name.
+
+    The Unicode key is checked first; this avoids having to list most printable
+    keysyms such as ``EuroSign``.
+    """
+    # For non-printable characters, gtk3 passes "\0" whereas tk passes an "".
+    if unikey and unikey.isprintable():
+        return unikey
+    key = keysym.lower()
+    if key.startswith("kp_"):  # keypad_x (including kp_enter).
+        key = key[3:]
+    if key.startswith("page_"):  # page_{up,down}
+        key = key.replace("page_", "page")
+    if key.endswith(("_l", "_r")):  # alt_l, ctrl_l, shift_l.
+        key = key[:-2]
+    if sys.platform == "darwin" and key == "meta":
+        # meta should be reported as command on mac
+        key = "cmd"
+    key = {
+        "return": "enter",
+        "prior": "pageup",  # Used by tk.
+        "next": "pagedown",  # Used by tk.
+    }.get(key, key)
+    return key
+
+
+@functools.cache
+def _make_class_factory(mixin_class, fmt, attr_name=None):
+    """
+    Return a function that creates picklable classes inheriting from a mixin.
+
+    After ::
+
+        factory = _make_class_factory(FooMixin, fmt, attr_name)
+        FooAxes = factory(Axes)
+
+    ``Foo`` is a class that inherits from ``FooMixin`` and ``Axes`` and **is
+    picklable** (picklability is what differentiates this from a plain call to
+    `type`).  Its ``__name__`` is set to ``fmt.format(Axes.__name__)`` and the
+    base class is stored in the ``attr_name`` attribute, if not None.
+
+    Moreover, the return value of ``factory`` is memoized: calls with the same
+    ``Axes`` class always return the same subclass.
+    """
+
+    @functools.cache
+    def class_factory(axes_class):
+        # if we have already wrapped this class, declare victory!
+        if issubclass(axes_class, mixin_class):
+            return axes_class
+
+        # The parameter is named "axes_class" for backcompat but is really just
+        # a base class; no axes semantics are used.
+        base_class = axes_class
+
+        class subcls(mixin_class, base_class):
+            # Better approximation than __module__ = "matplotlib.cbook".
+            __module__ = mixin_class.__module__
+
+            def __reduce__(self):
+                return (_picklable_class_constructor,
+                        (mixin_class, fmt, attr_name, base_class),
+                        self.__getstate__())
+
+        subcls.__name__ = subcls.__qualname__ = fmt.format(base_class.__name__)
+        if attr_name is not None:
+            setattr(subcls, attr_name, base_class)
+        return subcls
+
+    class_factory.__module__ = mixin_class.__module__
+    return class_factory
+
+
+def _picklable_class_constructor(mixin_class, fmt, attr_name, base_class):
+    """Internal helper for _make_class_factory."""
+    factory = _make_class_factory(mixin_class, fmt, attr_name)
+    cls = factory(base_class)
+    return cls.__new__(cls)
+
+
+def _is_torch_array(x):
+    """Check if 'x' is a PyTorch Tensor."""
+    try:
+        # we're intentionally not attempting to import torch. If somebody
+        # has created a torch array, torch should already be in sys.modules
+        return isinstance(x, sys.modules['torch'].Tensor)
+    except Exception:  # TypeError, KeyError, AttributeError, maybe others?
+        # we're attempting to access attributes on imported modules which
+        # may have arbitrary user code, so we deliberately catch all exceptions
+        return False
+
+
+def _is_jax_array(x):
+    """Check if 'x' is a JAX Array."""
+    try:
+        # we're intentionally not attempting to import jax. If somebody
+        # has created a jax array, jax should already be in sys.modules
+        return isinstance(x, sys.modules['jax'].Array)
+    except Exception:  # TypeError, KeyError, AttributeError, maybe others?
+        # we're attempting to access attributes on imported modules which
+        # may have arbitrary user code, so we deliberately catch all exceptions
+        return False
+
+
+def _is_tensorflow_array(x):
+    """Check if 'x' is a TensorFlow Tensor or Variable."""
+    try:
+        # we're intentionally not attempting to import TensorFlow. If somebody
+        # has created a TensorFlow array, TensorFlow should already be in sys.modules
+        # we use `is_tensor` to not depend on the class structure of TensorFlow
+        # arrays, as `tf.Variables` are not instances of `tf.Tensor`
+        # (they both convert the same way)
+        return isinstance(x, sys.modules['tensorflow'].is_tensor(x))
+    except Exception:  # TypeError, KeyError, AttributeError, maybe others?
+        # we're attempting to access attributes on imported modules which
+        # may have arbitrary user code, so we deliberately catch all exceptions
+        return False
+
+
+def _unpack_to_numpy(x):
+    """Internal helper to extract data from e.g. pandas and xarray objects."""
+    if isinstance(x, np.ndarray):
+        # If numpy, return directly
+        return x
+    if hasattr(x, 'to_numpy'):
+        # Assume that any to_numpy() method actually returns a numpy array
+        return x.to_numpy()
+    if hasattr(x, 'values'):
+        xtmp = x.values
+        # For example a dict has a 'values' attribute, but it is not a property
+        # so in this case we do not want to return a function
+        if isinstance(xtmp, np.ndarray):
+            return xtmp
+    if _is_torch_array(x) or _is_jax_array(x) or _is_tensorflow_array(x):
+        # using np.asarray() instead of explicitly __array__(), as the latter is
+        # only _one_ of many methods, and it's the last resort, see also
+        # https://numpy.org/devdocs/user/basics.interoperability.html#using-arbitrary-objects-in-numpy
+        # therefore, let arrays do better if they can
+        xtmp = np.asarray(x)
+
+        # In case np.asarray method does not return a numpy array in future
+        if isinstance(xtmp, np.ndarray):
+            return xtmp
+    return x
+
+
+def _auto_format_str(fmt, value):
+    """
+    Apply *value* to the format string *fmt*.
+
+    This works both with unnamed %-style formatting and
+    unnamed {}-style formatting. %-style formatting has priority.
+    If *fmt* is %-style formattable that will be used. Otherwise,
+    {}-formatting is applied. Strings without formatting placeholders
+    are passed through as is.
+
+    Examples
+    --------
+    >>> _auto_format_str('%.2f m', 0.2)
+    '0.20 m'
+    >>> _auto_format_str('{} m', 0.2)
+    '0.2 m'
+    >>> _auto_format_str('const', 0.2)
+    'const'
+    >>> _auto_format_str('%d or {}', 0.2)
+    '0 or {}'
+    """
+    try:
+        return fmt % (value,)
+    except (TypeError, ValueError):
+        return fmt.format(value)
+
+
+def _is_pandas_dataframe(x):
+    """Check if 'x' is a Pandas DataFrame."""
+    try:
+        # we're intentionally not attempting to import Pandas. If somebody
+        # has created a Pandas DataFrame, Pandas should already be in sys.modules
+        return isinstance(x, sys.modules['pandas'].DataFrame)
+    except Exception:  # TypeError, KeyError, AttributeError, maybe others?
+        # we're attempting to access attributes on imported modules which
+        # may have arbitrary user code, so we deliberately catch all exceptions
+        return False

lib/python3.12/site-packages/matplotlib/dviread.py

@@ -0,0 +1,1139 @@
+"""
+A module for reading dvi files output by TeX. Several limitations make
+this not (currently) useful as a general-purpose dvi preprocessor, but
+it is currently used by the pdf backend for processing usetex text.
+
+Interface::
+
+  with Dvi(filename, 72) as dvi:
+      # iterate over pages:
+      for page in dvi:
+          w, h, d = page.width, page.height, page.descent
+          for x, y, font, glyph, width in page.text:
+              fontname = font.texname
+              pointsize = font.size
+              ...
+          for x, y, height, width in page.boxes:
+              ...
+"""
+
+from collections import namedtuple
+import enum
+from functools import lru_cache, partial, wraps
+import logging
+import os
+from pathlib import Path
+import re
+import struct
+import subprocess
+import sys
+
+import numpy as np
+
+from matplotlib import _api, cbook
+
+_log = logging.getLogger(__name__)
+
+# Many dvi related files are looked for by external processes, require
+# additional parsing, and are used many times per rendering, which is why they
+# are cached using lru_cache().
+
+# Dvi is a bytecode format documented in
+# https://ctan.org/pkg/dvitype
+# https://texdoc.org/serve/dvitype.pdf/0
+#
+# The file consists of a preamble, some number of pages, a postamble,
+# and a finale. Different opcodes are allowed in different contexts,
+# so the Dvi object has a parser state:
+#
+#   pre:       expecting the preamble
+#   outer:     between pages (followed by a page or the postamble,
+#              also e.g. font definitions are allowed)
+#   page:      processing a page
+#   post_post: state after the postamble (our current implementation
+#              just stops reading)
+#   finale:    the finale (unimplemented in our current implementation)
+
+_dvistate = enum.Enum('DviState', 'pre outer inpage post_post finale')
+
+# The marks on a page consist of text and boxes. A page also has dimensions.
+Page = namedtuple('Page', 'text boxes height width descent')
+Box = namedtuple('Box', 'x y height width')
+
+
+# Also a namedtuple, for backcompat.
+class Text(namedtuple('Text', 'x y font glyph width')):
+    """
+    A glyph in the dvi file.
+
+    The *x* and *y* attributes directly position the glyph.  The *font*,
+    *glyph*, and *width* attributes are kept public for back-compatibility,
+    but users wanting to draw the glyph themselves are encouraged to instead
+    load the font specified by `font_path` at `font_size`, warp it with the
+    effects specified by `font_effects`, and load the glyph specified by
+    `glyph_name_or_index`.
+    """
+
+    def _get_pdftexmap_entry(self):
+        return PsfontsMap(find_tex_file("pdftex.map"))[self.font.texname]
+
+    @property
+    def font_path(self):
+        """The `~pathlib.Path` to the font for this glyph."""
+        psfont = self._get_pdftexmap_entry()
+        if psfont.filename is None:
+            raise ValueError("No usable font file found for {} ({}); "
+                             "the font may lack a Type-1 version"
+                             .format(psfont.psname.decode("ascii"),
+                                     psfont.texname.decode("ascii")))
+        return Path(psfont.filename)
+
+    @property
+    def font_size(self):
+        """The font size."""
+        return self.font.size
+
+    @property
+    def font_effects(self):
+        """
+        The "font effects" dict for this glyph.
+
+        This dict contains the values for this glyph of SlantFont and
+        ExtendFont (if any), read off :file:`pdftex.map`.
+        """
+        return self._get_pdftexmap_entry().effects
+
+    @property
+    def glyph_name_or_index(self):
+        """
+        Either the glyph name or the native charmap glyph index.
+
+        If :file:`pdftex.map` specifies an encoding for this glyph's font, that
+        is a mapping of glyph indices to Adobe glyph names; use it to convert
+        dvi indices to glyph names.  Callers can then convert glyph names to
+        glyph indices (with FT_Get_Name_Index/get_name_index), and load the
+        glyph using FT_Load_Glyph/load_glyph.
+
+        If :file:`pdftex.map` specifies no encoding, the indices directly map
+        to the font's "native" charmap; glyphs should directly load using
+        FT_Load_Char/load_char after selecting the native charmap.
+        """
+        entry = self._get_pdftexmap_entry()
+        return (_parse_enc(entry.encoding)[self.glyph]
+                if entry.encoding is not None else self.glyph)
+
+
+# Opcode argument parsing
+#
+# Each of the following functions takes a Dvi object and delta, which is the
+# difference between the opcode and the minimum opcode with the same meaning.
+# Dvi opcodes often encode the number of argument bytes in this delta.
+_arg_mapping = dict(
+    # raw: Return delta as is.
+    raw=lambda dvi, delta: delta,
+    # u1: Read 1 byte as an unsigned number.
+    u1=lambda dvi, delta: dvi._read_arg(1, signed=False),
+    # u4: Read 4 bytes as an unsigned number.
+    u4=lambda dvi, delta: dvi._read_arg(4, signed=False),
+    # s4: Read 4 bytes as a signed number.
+    s4=lambda dvi, delta: dvi._read_arg(4, signed=True),
+    # slen: Read delta bytes as a signed number, or None if delta is None.
+    slen=lambda dvi, delta: dvi._read_arg(delta, signed=True) if delta else None,
+    # slen1: Read (delta + 1) bytes as a signed number.
+    slen1=lambda dvi, delta: dvi._read_arg(delta + 1, signed=True),
+    # ulen1: Read (delta + 1) bytes as an unsigned number.
+    ulen1=lambda dvi, delta: dvi._read_arg(delta + 1, signed=False),
+    # olen1: Read (delta + 1) bytes as an unsigned number if less than 4 bytes,
+    # as a signed number if 4 bytes.
+    olen1=lambda dvi, delta: dvi._read_arg(delta + 1, signed=(delta == 3)),
+)
+
+
+def _dispatch(table, min, max=None, state=None, args=('raw',)):
+    """
+    Decorator for dispatch by opcode. Sets the values in *table*
+    from *min* to *max* to this method, adds a check that the Dvi state
+    matches *state* if not None, reads arguments from the file according
+    to *args*.
+
+    Parameters
+    ----------
+    table : dict[int, callable]
+        The dispatch table to be filled in.
+
+    min, max : int
+        Range of opcodes that calls the registered function; *max* defaults to
+        *min*.
+
+    state : _dvistate, optional
+        State of the Dvi object in which these opcodes are allowed.
+
+    args : list[str], default: ['raw']
+        Sequence of argument specifications:
+
+        - 'raw': opcode minus minimum
+        - 'u1': read one unsigned byte
+        - 'u4': read four bytes, treat as an unsigned number
+        - 's4': read four bytes, treat as a signed number
+        - 'slen': read (opcode - minimum) bytes, treat as signed
+        - 'slen1': read (opcode - minimum + 1) bytes, treat as signed
+        - 'ulen1': read (opcode - minimum + 1) bytes, treat as unsigned
+        - 'olen1': read (opcode - minimum + 1) bytes, treat as unsigned
+          if under four bytes, signed if four bytes
+    """
+    def decorate(method):
+        get_args = [_arg_mapping[x] for x in args]
+
+        @wraps(method)
+        def wrapper(self, byte):
+            if state is not None and self.state != state:
+                raise ValueError("state precondition failed")
+            return method(self, *[f(self, byte-min) for f in get_args])
+        if max is None:
+            table[min] = wrapper
+        else:
+            for i in range(min, max+1):
+                assert table[i] is None
+                table[i] = wrapper
+        return wrapper
+    return decorate
+
+
+class Dvi:
+    """
+    A reader for a dvi ("device-independent") file, as produced by TeX.
+
+    The current implementation can only iterate through pages in order,
+    and does not even attempt to verify the postamble.
+
+    This class can be used as a context manager to close the underlying
+    file upon exit. Pages can be read via iteration. Here is an overly
+    simple way to extract text without trying to detect whitespace::
+
+        >>> with matplotlib.dviread.Dvi('input.dvi', 72) as dvi:
+        ...     for page in dvi:
+        ...         print(''.join(chr(t.glyph) for t in page.text))
+    """
+    # dispatch table
+    _dtable = [None] * 256
+    _dispatch = partial(_dispatch, _dtable)
+
+    def __init__(self, filename, dpi):
+        """
+        Read the data from the file named *filename* and convert
+        TeX's internal units to units of *dpi* per inch.
+        *dpi* only sets the units and does not limit the resolution.
+        Use None to return TeX's internal units.
+        """
+        _log.debug('Dvi: %s', filename)
+        self.file = open(filename, 'rb')
+        self.dpi = dpi
+        self.fonts = {}
+        self.state = _dvistate.pre
+        self._missing_font = None
+
+    def __enter__(self):
+        """Context manager enter method, does nothing."""
+        return self
+
+    def __exit__(self, etype, evalue, etrace):
+        """
+        Context manager exit method, closes the underlying file if it is open.
+        """
+        self.close()
+
+    def __iter__(self):
+        """
+        Iterate through the pages of the file.
+
+        Yields
+        ------
+        Page
+            Details of all the text and box objects on the page.
+            The Page tuple contains lists of Text and Box tuples and
+            the page dimensions, and the Text and Box tuples contain
+            coordinates transformed into a standard Cartesian
+            coordinate system at the dpi value given when initializing.
+            The coordinates are floating point numbers, but otherwise
+            precision is not lost and coordinate values are not clipped to
+            integers.
+        """
+        while self._read():
+            yield self._output()
+
+    def close(self):
+        """Close the underlying file if it is open."""
+        if not self.file.closed:
+            self.file.close()
+
+    def _output(self):
+        """
+        Output the text and boxes belonging to the most recent page.
+        page = dvi._output()
+        """
+        minx = miny = np.inf
+        maxx = maxy = -np.inf
+        maxy_pure = -np.inf
+        for elt in self.text + self.boxes:
+            if isinstance(elt, Box):
+                x, y, h, w = elt
+                e = 0  # zero depth
+            else:  # glyph
+                x, y, font, g, w = elt
+                h, e = font._height_depth_of(g)
+            minx = min(minx, x)
+            miny = min(miny, y - h)
+            maxx = max(maxx, x + w)
+            maxy = max(maxy, y + e)
+            maxy_pure = max(maxy_pure, y)
+        if self._baseline_v is not None:
+            maxy_pure = self._baseline_v  # This should normally be the case.
+            self._baseline_v = None
+
+        if not self.text and not self.boxes:  # Avoid infs/nans from inf+/-inf.
+            return Page(text=[], boxes=[], width=0, height=0, descent=0)
+
+        if self.dpi is None:
+            # special case for ease of debugging: output raw dvi coordinates
+            return Page(text=self.text, boxes=self.boxes,
+                        width=maxx-minx, height=maxy_pure-miny,
+                        descent=maxy-maxy_pure)
+
+        # convert from TeX's "scaled points" to dpi units
+        d = self.dpi / (72.27 * 2**16)
+        descent = (maxy - maxy_pure) * d
+
+        text = [Text((x-minx)*d, (maxy-y)*d - descent, f, g, w*d)
+                for (x, y, f, g, w) in self.text]
+        boxes = [Box((x-minx)*d, (maxy-y)*d - descent, h*d, w*d)
+                 for (x, y, h, w) in self.boxes]
+
+        return Page(text=text, boxes=boxes, width=(maxx-minx)*d,
+                    height=(maxy_pure-miny)*d, descent=descent)
+
+    def _read(self):
+        """
+        Read one page from the file. Return True if successful,
+        False if there were no more pages.
+        """
+        # Pages appear to start with the sequence
+        #   bop (begin of page)
+        #   xxx comment
+        #   <push, ..., pop>  # if using chemformula
+        #   down
+        #   push
+        #     down
+        #     <push, push, xxx, right, xxx, pop, pop>  # if using xcolor
+        #     down
+        #     push
+        #       down (possibly multiple)
+        #       push  <=  here, v is the baseline position.
+        #         etc.
+        # (dviasm is useful to explore this structure.)
+        # Thus, we use the vertical position at the first time the stack depth
+        # reaches 3, while at least three "downs" have been executed (excluding
+        # those popped out (corresponding to the chemformula preamble)), as the
+        # baseline (the "down" count is necessary to handle xcolor).
+        down_stack = [0]
+        self._baseline_v = None
+        while True:
+            byte = self.file.read(1)[0]
+            self._dtable[byte](self, byte)
+            if self._missing_font:
+                raise self._missing_font.to_exception()
+            name = self._dtable[byte].__name__
+            if name == "_push":
+                down_stack.append(down_stack[-1])
+            elif name == "_pop":
+                down_stack.pop()
+            elif name == "_down":
+                down_stack[-1] += 1
+            if (self._baseline_v is None
+                    and len(getattr(self, "stack", [])) == 3
+                    and down_stack[-1] >= 4):
+                self._baseline_v = self.v
+            if byte == 140:                         # end of page
+                return True
+            if self.state is _dvistate.post_post:   # end of file
+                self.close()
+                return False
+
+    def _read_arg(self, nbytes, signed=False):
+        """
+        Read and return a big-endian integer *nbytes* long.
+        Signedness is determined by the *signed* keyword.
+        """
+        return int.from_bytes(self.file.read(nbytes), "big", signed=signed)
+
+    @_dispatch(min=0, max=127, state=_dvistate.inpage)
+    def _set_char_immediate(self, char):
+        self._put_char_real(char)
+        if isinstance(self.fonts[self.f], cbook._ExceptionInfo):
+            return
+        self.h += self.fonts[self.f]._width_of(char)
+
+    @_dispatch(min=128, max=131, state=_dvistate.inpage, args=('olen1',))
+    def _set_char(self, char):
+        self._put_char_real(char)
+        if isinstance(self.fonts[self.f], cbook._ExceptionInfo):
+            return
+        self.h += self.fonts[self.f]._width_of(char)
+
+    @_dispatch(132, state=_dvistate.inpage, args=('s4', 's4'))
+    def _set_rule(self, a, b):
+        self._put_rule_real(a, b)
+        self.h += b
+
+    @_dispatch(min=133, max=136, state=_dvistate.inpage, args=('olen1',))
+    def _put_char(self, char):
+        self._put_char_real(char)
+
+    def _put_char_real(self, char):
+        font = self.fonts[self.f]
+        if isinstance(font, cbook._ExceptionInfo):
+            self._missing_font = font
+        elif font._vf is None:
+            self.text.append(Text(self.h, self.v, font, char,
+                                  font._width_of(char)))
+        else:
+            scale = font._scale
+            for x, y, f, g, w in font._vf[char].text:
+                newf = DviFont(scale=_mul2012(scale, f._scale),
+                               tfm=f._tfm, texname=f.texname, vf=f._vf)
+                self.text.append(Text(self.h + _mul2012(x, scale),
+                                      self.v + _mul2012(y, scale),
+                                      newf, g, newf._width_of(g)))
+            self.boxes.extend([Box(self.h + _mul2012(x, scale),
+                                   self.v + _mul2012(y, scale),
+                                   _mul2012(a, scale), _mul2012(b, scale))
+                               for x, y, a, b in font._vf[char].boxes])
+
+    @_dispatch(137, state=_dvistate.inpage, args=('s4', 's4'))
+    def _put_rule(self, a, b):
+        self._put_rule_real(a, b)
+
+    def _put_rule_real(self, a, b):
+        if a > 0 and b > 0:
+            self.boxes.append(Box(self.h, self.v, a, b))
+
+    @_dispatch(138)
+    def _nop(self, _):
+        pass
+
+    @_dispatch(139, state=_dvistate.outer, args=('s4',)*11)
+    def _bop(self, c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, p):
+        self.state = _dvistate.inpage
+        self.h = self.v = self.w = self.x = self.y = self.z = 0
+        self.stack = []
+        self.text = []          # list of Text objects
+        self.boxes = []         # list of Box objects
+
+    @_dispatch(140, state=_dvistate.inpage)
+    def _eop(self, _):
+        self.state = _dvistate.outer
+        del self.h, self.v, self.w, self.x, self.y, self.z, self.stack
+
+    @_dispatch(141, state=_dvistate.inpage)
+    def _push(self, _):
+        self.stack.append((self.h, self.v, self.w, self.x, self.y, self.z))
+
+    @_dispatch(142, state=_dvistate.inpage)
+    def _pop(self, _):
+        self.h, self.v, self.w, self.x, self.y, self.z = self.stack.pop()
+
+    @_dispatch(min=143, max=146, state=_dvistate.inpage, args=('slen1',))
+    def _right(self, b):
+        self.h += b
+
+    @_dispatch(min=147, max=151, state=_dvistate.inpage, args=('slen',))
+    def _right_w(self, new_w):
+        if new_w is not None:
+            self.w = new_w
+        self.h += self.w
+
+    @_dispatch(min=152, max=156, state=_dvistate.inpage, args=('slen',))
+    def _right_x(self, new_x):
+        if new_x is not None:
+            self.x = new_x
+        self.h += self.x
+
+    @_dispatch(min=157, max=160, state=_dvistate.inpage, args=('slen1',))
+    def _down(self, a):
+        self.v += a
+
+    @_dispatch(min=161, max=165, state=_dvistate.inpage, args=('slen',))
+    def _down_y(self, new_y):
+        if new_y is not None:
+            self.y = new_y
+        self.v += self.y
+
+    @_dispatch(min=166, max=170, state=_dvistate.inpage, args=('slen',))
+    def _down_z(self, new_z):
+        if new_z is not None:
+            self.z = new_z
+        self.v += self.z
+
+    @_dispatch(min=171, max=234, state=_dvistate.inpage)
+    def _fnt_num_immediate(self, k):
+        self.f = k
+
+    @_dispatch(min=235, max=238, state=_dvistate.inpage, args=('olen1',))
+    def _fnt_num(self, new_f):
+        self.f = new_f
+
+    @_dispatch(min=239, max=242, args=('ulen1',))
+    def _xxx(self, datalen):
+        special = self.file.read(datalen)
+        _log.debug(
+            'Dvi._xxx: encountered special: %s',
+            ''.join([chr(ch) if 32 <= ch < 127 else '<%02x>' % ch
+                     for ch in special]))
+
+    @_dispatch(min=243, max=246, args=('olen1', 'u4', 'u4', 'u4', 'u1', 'u1'))
+    def _fnt_def(self, k, c, s, d, a, l):
+        self._fnt_def_real(k, c, s, d, a, l)
+
+    def _fnt_def_real(self, k, c, s, d, a, l):
+        n = self.file.read(a + l)
+        fontname = n[-l:].decode('ascii')
+        try:
+            tfm = _tfmfile(fontname)
+        except FileNotFoundError as exc:
+            # Explicitly allow defining missing fonts for Vf support; we only
+            # register an error when trying to load a glyph from a missing font
+            # and throw that error in Dvi._read.  For Vf, _finalize_packet
+            # checks whether a missing glyph has been used, and in that case
+            # skips the glyph definition.
+            self.fonts[k] = cbook._ExceptionInfo.from_exception(exc)
+            return
+        if c != 0 and tfm.checksum != 0 and c != tfm.checksum:
+            raise ValueError(f'tfm checksum mismatch: {n}')
+        try:
+            vf = _vffile(fontname)
+        except FileNotFoundError:
+            vf = None
+        self.fonts[k] = DviFont(scale=s, tfm=tfm, texname=n, vf=vf)
+
+    @_dispatch(247, state=_dvistate.pre, args=('u1', 'u4', 'u4', 'u4', 'u1'))
+    def _pre(self, i, num, den, mag, k):
+        self.file.read(k)  # comment in the dvi file
+        if i != 2:
+            raise ValueError(f"Unknown dvi format {i}")
+        if num != 25400000 or den != 7227 * 2**16:
+            raise ValueError("Nonstandard units in dvi file")
+            # meaning: TeX always uses those exact values, so it
+            # should be enough for us to support those
+            # (There are 72.27 pt to an inch so 7227 pt =
+            # 7227 * 2**16 sp to 100 in. The numerator is multiplied
+            # by 10^5 to get units of 10**-7 meters.)
+        if mag != 1000:
+            raise ValueError("Nonstandard magnification in dvi file")
+            # meaning: LaTeX seems to frown on setting \mag, so
+            # I think we can assume this is constant
+        self.state = _dvistate.outer
+
+    @_dispatch(248, state=_dvistate.outer)
+    def _post(self, _):
+        self.state = _dvistate.post_post
+        # TODO: actually read the postamble and finale?
+        # currently post_post just triggers closing the file
+
+    @_dispatch(249)
+    def _post_post(self, _):
+        raise NotImplementedError
+
+    @_dispatch(min=250, max=255)
+    def _malformed(self, offset):
+        raise ValueError(f"unknown command: byte {250 + offset}")
+
+
+class DviFont:
+    """
+    Encapsulation of a font that a DVI file can refer to.
+
+    This class holds a font's texname and size, supports comparison,
+    and knows the widths of glyphs in the same units as the AFM file.
+    There are also internal attributes (for use by dviread.py) that
+    are *not* used for comparison.
+
+    The size is in Adobe points (converted from TeX points).
+
+    Parameters
+    ----------
+    scale : float
+        Factor by which the font is scaled from its natural size.
+    tfm : Tfm
+        TeX font metrics for this font
+    texname : bytes
+       Name of the font as used internally by TeX and friends, as an ASCII
+       bytestring.  This is usually very different from any external font
+       names; `PsfontsMap` can be used to find the external name of the font.
+    vf : Vf
+       A TeX "virtual font" file, or None if this font is not virtual.
+
+    Attributes
+    ----------
+    texname : bytes
+    size : float
+       Size of the font in Adobe points, converted from the slightly
+       smaller TeX points.
+    widths : list
+       Widths of glyphs in glyph-space units, typically 1/1000ths of
+       the point size.
+
+    """
+    __slots__ = ('texname', 'size', 'widths', '_scale', '_vf', '_tfm')
+
+    def __init__(self, scale, tfm, texname, vf):
+        _api.check_isinstance(bytes, texname=texname)
+        self._scale = scale
+        self._tfm = tfm
+        self.texname = texname
+        self._vf = vf
+        self.size = scale * (72.0 / (72.27 * 2**16))
+        try:
+            nchars = max(tfm.width) + 1
+        except ValueError:
+            nchars = 0
+        self.widths = [(1000*tfm.width.get(char, 0)) >> 20
+                       for char in range(nchars)]
+
+    def __eq__(self, other):
+        return (type(self) is type(other)
+                and self.texname == other.texname and self.size == other.size)
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __repr__(self):
+        return f"<{type(self).__name__}: {self.texname}>"
+
+    def _width_of(self, char):
+        """Width of char in dvi units."""
+        width = self._tfm.width.get(char, None)
+        if width is not None:
+            return _mul2012(width, self._scale)
+        _log.debug('No width for char %d in font %s.', char, self.texname)
+        return 0
+
+    def _height_depth_of(self, char):
+        """Height and depth of char in dvi units."""
+        result = []
+        for metric, name in ((self._tfm.height, "height"),
+                             (self._tfm.depth, "depth")):
+            value = metric.get(char, None)
+            if value is None:
+                _log.debug('No %s for char %d in font %s',
+                           name, char, self.texname)
+                result.append(0)
+            else:
+                result.append(_mul2012(value, self._scale))
+        # cmsyXX (symbols font) glyph 0 ("minus") has a nonzero descent
+        # so that TeX aligns equations properly
+        # (https://tex.stackexchange.com/q/526103/)
+        # but we actually care about the rasterization depth to align
+        # the dvipng-generated images.
+        if re.match(br'^cmsy\d+$', self.texname) and char == 0:
+            result[-1] = 0
+        return result
+
+
+class Vf(Dvi):
+    r"""
+    A virtual font (\*.vf file) containing subroutines for dvi files.
+
+    Parameters
+    ----------
+    filename : str or path-like
+
+    Notes
+    -----
+    The virtual font format is a derivative of dvi:
+    http://mirrors.ctan.org/info/knuth/virtual-fonts
+    This class reuses some of the machinery of `Dvi`
+    but replaces the `_read` loop and dispatch mechanism.
+
+    Examples
+    --------
+    ::
+
+        vf = Vf(filename)
+        glyph = vf[code]
+        glyph.text, glyph.boxes, glyph.width
+    """
+
+    def __init__(self, filename):
+        super().__init__(filename, 0)
+        try:
+            self._first_font = None
+            self._chars = {}
+            self._read()
+        finally:
+            self.close()
+
+    def __getitem__(self, code):
+        return self._chars[code]
+
+    def _read(self):
+        """
+        Read one page from the file. Return True if successful,
+        False if there were no more pages.
+        """
+        packet_char = packet_ends = None
+        packet_len = packet_width = None
+        while True:
+            byte = self.file.read(1)[0]
+            # If we are in a packet, execute the dvi instructions
+            if self.state is _dvistate.inpage:
+                byte_at = self.file.tell()-1
+                if byte_at == packet_ends:
+                    self._finalize_packet(packet_char, packet_width)
+                    packet_len = packet_char = packet_width = None
+                    # fall through to out-of-packet code
+                elif byte_at > packet_ends:
+                    raise ValueError("Packet length mismatch in vf file")
+                else:
+                    if byte in (139, 140) or byte >= 243:
+                        raise ValueError(f"Inappropriate opcode {byte} in vf file")
+                    Dvi._dtable[byte](self, byte)
+                    continue
+
+            # We are outside a packet
+            if byte < 242:          # a short packet (length given by byte)
+                packet_len = byte
+                packet_char = self._read_arg(1)
+                packet_width = self._read_arg(3)
+                packet_ends = self._init_packet(byte)
+                self.state = _dvistate.inpage
+            elif byte == 242:       # a long packet
+                packet_len = self._read_arg(4)
+                packet_char = self._read_arg(4)
+                packet_width = self._read_arg(4)
+                self._init_packet(packet_len)
+            elif 243 <= byte <= 246:
+                k = self._read_arg(byte - 242, byte == 246)
+                c = self._read_arg(4)
+                s = self._read_arg(4)
+                d = self._read_arg(4)
+                a = self._read_arg(1)
+                l = self._read_arg(1)
+                self._fnt_def_real(k, c, s, d, a, l)
+                if self._first_font is None:
+                    self._first_font = k
+            elif byte == 247:       # preamble
+                i = self._read_arg(1)
+                k = self._read_arg(1)
+                x = self.file.read(k)
+                cs = self._read_arg(4)
+                ds = self._read_arg(4)
+                self._pre(i, x, cs, ds)
+            elif byte == 248:       # postamble (just some number of 248s)
+                break
+            else:
+                raise ValueError(f"Unknown vf opcode {byte}")
+
+    def _init_packet(self, pl):
+        if self.state != _dvistate.outer:
+            raise ValueError("Misplaced packet in vf file")
+        self.h = self.v = self.w = self.x = self.y = self.z = 0
+        self.stack = []
+        self.text = []
+        self.boxes = []
+        self.f = self._first_font
+        self._missing_font = None
+        return self.file.tell() + pl
+
+    def _finalize_packet(self, packet_char, packet_width):
+        if not self._missing_font:  # Otherwise we don't have full glyph definition.
+            self._chars[packet_char] = Page(
+                text=self.text, boxes=self.boxes, width=packet_width,
+                height=None, descent=None)
+        self.state = _dvistate.outer
+
+    def _pre(self, i, x, cs, ds):
+        if self.state is not _dvistate.pre:
+            raise ValueError("pre command in middle of vf file")
+        if i != 202:
+            raise ValueError(f"Unknown vf format {i}")
+        if len(x):
+            _log.debug('vf file comment: %s', x)
+        self.state = _dvistate.outer
+        # cs = checksum, ds = design size
+
+
+def _mul2012(num1, num2):
+    """Multiply two numbers in 20.12 fixed point format."""
+    # Separated into a function because >> has surprising precedence
+    return (num1*num2) >> 20
+
+
+class Tfm:
+    """
+    A TeX Font Metric file.
+
+    This implementation covers only the bare minimum needed by the Dvi class.
+
+    Parameters
+    ----------
+    filename : str or path-like
+
+    Attributes
+    ----------
+    checksum : int
+       Used for verifying against the dvi file.
+    design_size : int
+       Design size of the font (unknown units)
+    width, height, depth : dict
+       Dimensions of each character, need to be scaled by the factor
+       specified in the dvi file. These are dicts because indexing may
+       not start from 0.
+    """
+    __slots__ = ('checksum', 'design_size', 'width', 'height', 'depth')
+
+    def __init__(self, filename):
+        _log.debug('opening tfm file %s', filename)
+        with open(filename, 'rb') as file:
+            header1 = file.read(24)
+            lh, bc, ec, nw, nh, nd = struct.unpack('!6H', header1[2:14])
+            _log.debug('lh=%d, bc=%d, ec=%d, nw=%d, nh=%d, nd=%d',
+                       lh, bc, ec, nw, nh, nd)
+            header2 = file.read(4*lh)
+            self.checksum, self.design_size = struct.unpack('!2I', header2[:8])
+            # there is also encoding information etc.
+            char_info = file.read(4*(ec-bc+1))
+            widths = struct.unpack(f'!{nw}i', file.read(4*nw))
+            heights = struct.unpack(f'!{nh}i', file.read(4*nh))
+            depths = struct.unpack(f'!{nd}i', file.read(4*nd))
+        self.width = {}
+        self.height = {}
+        self.depth = {}
+        for idx, char in enumerate(range(bc, ec+1)):
+            byte0 = char_info[4*idx]
+            byte1 = char_info[4*idx+1]
+            self.width[char] = widths[byte0]
+            self.height[char] = heights[byte1 >> 4]
+            self.depth[char] = depths[byte1 & 0xf]
+
+
+PsFont = namedtuple('PsFont', 'texname psname effects encoding filename')
+
+
+class PsfontsMap:
+    """
+    A psfonts.map formatted file, mapping TeX fonts to PS fonts.
+
+    Parameters
+    ----------
+    filename : str or path-like
+
+    Notes
+    -----
+    For historical reasons, TeX knows many Type-1 fonts by different
+    names than the outside world. (For one thing, the names have to
+    fit in eight characters.) Also, TeX's native fonts are not Type-1
+    but Metafont, which is nontrivial to convert to PostScript except
+    as a bitmap. While high-quality conversions to Type-1 format exist
+    and are shipped with modern TeX distributions, we need to know
+    which Type-1 fonts are the counterparts of which native fonts. For
+    these reasons a mapping is needed from internal font names to font
+    file names.
+
+    A texmf tree typically includes mapping files called e.g.
+    :file:`psfonts.map`, :file:`pdftex.map`, or :file:`dvipdfm.map`.
+    The file :file:`psfonts.map` is used by :program:`dvips`,
+    :file:`pdftex.map` by :program:`pdfTeX`, and :file:`dvipdfm.map`
+    by :program:`dvipdfm`. :file:`psfonts.map` might avoid embedding
+    the 35 PostScript fonts (i.e., have no filename for them, as in
+    the Times-Bold example above), while the pdf-related files perhaps
+    only avoid the "Base 14" pdf fonts. But the user may have
+    configured these files differently.
+
+    Examples
+    --------
+    >>> map = PsfontsMap(find_tex_file('pdftex.map'))
+    >>> entry = map[b'ptmbo8r']
+    >>> entry.texname
+    b'ptmbo8r'
+    >>> entry.psname
+    b'Times-Bold'
+    >>> entry.encoding
+    '/usr/local/texlive/2008/texmf-dist/fonts/enc/dvips/base/8r.enc'
+    >>> entry.effects
+    {'slant': 0.16700000000000001}
+    >>> entry.filename
+    """
+    __slots__ = ('_filename', '_unparsed', '_parsed')
+
+    # Create a filename -> PsfontsMap cache, so that calling
+    # `PsfontsMap(filename)` with the same filename a second time immediately
+    # returns the same object.
+    @lru_cache
+    def __new__(cls, filename):
+        self = object.__new__(cls)
+        self._filename = os.fsdecode(filename)
+        # Some TeX distributions have enormous pdftex.map files which would
+        # take hundreds of milliseconds to parse, but it is easy enough to just
+        # store the unparsed lines (keyed by the first word, which is the
+        # texname) and parse them on-demand.
+        with open(filename, 'rb') as file:
+            self._unparsed = {}
+            for line in file:
+                tfmname = line.split(b' ', 1)[0]
+                self._unparsed.setdefault(tfmname, []).append(line)
+        self._parsed = {}
+        return self
+
+    def __getitem__(self, texname):
+        assert isinstance(texname, bytes)
+        if texname in self._unparsed:
+            for line in self._unparsed.pop(texname):
+                if self._parse_and_cache_line(line):
+                    break
+        try:
+            return self._parsed[texname]
+        except KeyError:
+            raise LookupError(
+                f"An associated PostScript font (required by Matplotlib) "
+                f"could not be found for TeX font {texname.decode('ascii')!r} "
+                f"in {self._filename!r}; this problem can often be solved by "
+                f"installing a suitable PostScript font package in your TeX "
+                f"package manager") from None
+
+    def _parse_and_cache_line(self, line):
+        """
+        Parse a line in the font mapping file.
+
+        The format is (partially) documented at
+        http://mirrors.ctan.org/systems/doc/pdftex/manual/pdftex-a.pdf
+        https://tug.org/texinfohtml/dvips.html#psfonts_002emap
+        Each line can have the following fields:
+
+        - tfmname (first, only required field),
+        - psname (defaults to tfmname, must come immediately after tfmname if
+          present),
+        - fontflags (integer, must come immediately after psname if present,
+          ignored by us),
+        - special (SlantFont and ExtendFont, only field that is double-quoted),
+        - fontfile, encodingfile (optional, prefixed by <, <<, or <[; << always
+          precedes a font, <[ always precedes an encoding, < can precede either
+          but then an encoding file must have extension .enc; < and << also
+          request different font subsetting behaviors but we ignore that; < can
+          be separated from the filename by whitespace).
+
+        special, fontfile, and encodingfile can appear in any order.
+        """
+        # If the map file specifies multiple encodings for a font, we
+        # follow pdfTeX in choosing the last one specified. Such
+        # entries are probably mistakes but they have occurred.
+        # https://tex.stackexchange.com/q/10826/
+
+        if not line or line.startswith((b" ", b"%", b"*", b";", b"#")):
+            return
+        tfmname = basename = special = encodingfile = fontfile = None
+        is_subsetted = is_t1 = is_truetype = False
+        matches = re.finditer(br'"([^"]*)(?:"|$)|(\S+)', line)
+        for match in matches:
+            quoted, unquoted = match.groups()
+            if unquoted:
+                if unquoted.startswith(b"<<"):  # font
+                    fontfile = unquoted[2:]
+                elif unquoted.startswith(b"<["):  # encoding
+                    encodingfile = unquoted[2:]
+                elif unquoted.startswith(b"<"):  # font or encoding
+                    word = (
+                        # <foo => foo
+                        unquoted[1:]
+                        # < by itself => read the next word
+                        or next(filter(None, next(matches).groups())))
+                    if word.endswith(b".enc"):
+                        encodingfile = word
+                    else:
+                        fontfile = word
+                        is_subsetted = True
+                elif tfmname is None:
+                    tfmname = unquoted
+                elif basename is None:
+                    basename = unquoted
+            elif quoted:
+                special = quoted
+        effects = {}
+        if special:
+            words = reversed(special.split())
+            for word in words:
+                if word == b"SlantFont":
+                    effects["slant"] = float(next(words))
+                elif word == b"ExtendFont":
+                    effects["extend"] = float(next(words))
+
+        # Verify some properties of the line that would cause it to be ignored
+        # otherwise.
+        if fontfile is not None:
+            if fontfile.endswith((b".ttf", b".ttc")):
+                is_truetype = True
+            elif not fontfile.endswith(b".otf"):
+                is_t1 = True
+        elif basename is not None:
+            is_t1 = True
+        if is_truetype and is_subsetted and encodingfile is None:
+            return
+        if not is_t1 and ("slant" in effects or "extend" in effects):
+            return
+        if abs(effects.get("slant", 0)) > 1:
+            return
+        if abs(effects.get("extend", 0)) > 2:
+            return
+
+        if basename is None:
+            basename = tfmname
+        if encodingfile is not None:
+            encodingfile = find_tex_file(encodingfile)
+        if fontfile is not None:
+            fontfile = find_tex_file(fontfile)
+        self._parsed[tfmname] = PsFont(
+            texname=tfmname, psname=basename, effects=effects,
+            encoding=encodingfile, filename=fontfile)
+        return True
+
+
+def _parse_enc(path):
+    r"""
+    Parse a \*.enc file referenced from a psfonts.map style file.
+
+    The format supported by this function is a tiny subset of PostScript.
+
+    Parameters
+    ----------
+    path : `os.PathLike`
+
+    Returns
+    -------
+    list
+        The nth entry of the list is the PostScript glyph name of the nth
+        glyph.
+    """
+    no_comments = re.sub("%.*", "", Path(path).read_text(encoding="ascii"))
+    array = re.search(r"(?s)\[(.*)\]", no_comments).group(1)
+    lines = [line for line in array.split() if line]
+    if all(line.startswith("/") for line in lines):
+        return [line[1:] for line in lines]
+    else:
+        raise ValueError(f"Failed to parse {path} as Postscript encoding")
+
+
+class _LuatexKpsewhich:
+    @lru_cache  # A singleton.
+    def __new__(cls):
+        self = object.__new__(cls)
+        self._proc = self._new_proc()
+        return self
+
+    def _new_proc(self):
+        return subprocess.Popen(
+            ["luatex", "--luaonly",
+             str(cbook._get_data_path("kpsewhich.lua"))],
+            stdin=subprocess.PIPE, stdout=subprocess.PIPE)
+
+    def search(self, filename):
+        if self._proc.poll() is not None:  # Dead, restart it.
+            self._proc = self._new_proc()
+        self._proc.stdin.write(os.fsencode(filename) + b"\n")
+        self._proc.stdin.flush()
+        out = self._proc.stdout.readline().rstrip()
+        return None if out == b"nil" else os.fsdecode(out)
+
+
+@lru_cache
+def find_tex_file(filename):
+    """
+    Find a file in the texmf tree using kpathsea_.
+
+    The kpathsea library, provided by most existing TeX distributions, both
+    on Unix-like systems and on Windows (MikTeX), is invoked via a long-lived
+    luatex process if luatex is installed, or via kpsewhich otherwise.
+
+    .. _kpathsea: https://www.tug.org/kpathsea/
+
+    Parameters
+    ----------
+    filename : str or path-like
+
+    Raises
+    ------
+    FileNotFoundError
+        If the file is not found.
+    """
+
+    # we expect these to always be ascii encoded, but use utf-8
+    # out of caution
+    if isinstance(filename, bytes):
+        filename = filename.decode('utf-8', errors='replace')
+
+    try:
+        lk = _LuatexKpsewhich()
+    except FileNotFoundError:
+        lk = None  # Fallback to directly calling kpsewhich, as below.
+
+    if lk:
+        path = lk.search(filename)
+    else:
+        if sys.platform == 'win32':
+            # On Windows only, kpathsea can use utf-8 for cmd args and output.
+            # The `command_line_encoding` environment variable is set to force
+            # it to always use utf-8 encoding.  See Matplotlib issue #11848.
+            kwargs = {'env': {**os.environ, 'command_line_encoding': 'utf-8'},
+                      'encoding': 'utf-8'}
+        else:  # On POSIX, run through the equivalent of os.fsdecode().
+            kwargs = {'encoding': sys.getfilesystemencoding(),
+                      'errors': 'surrogateescape'}
+
+        try:
+            path = (cbook._check_and_log_subprocess(['kpsewhich', filename],
+                                                    _log, **kwargs)
+                    .rstrip('\n'))
+        except (FileNotFoundError, RuntimeError):
+            path = None
+
+    if path:
+        return path
+    else:
+        raise FileNotFoundError(
+            f"Matplotlib's TeX implementation searched for a file named "
+            f"{filename!r} in your texmf tree, but could not find it")
+
+
+@lru_cache
+def _fontfile(cls, suffix, texname):
+    return cls(find_tex_file(texname + suffix))
+
+
+_tfmfile = partial(_fontfile, Tfm, ".tfm")
+_vffile = partial(_fontfile, Vf, ".vf")
+
+
+if __name__ == '__main__':
+    from argparse import ArgumentParser
+    import itertools
+
+    parser = ArgumentParser()
+    parser.add_argument("filename")
+    parser.add_argument("dpi", nargs="?", type=float, default=None)
+    args = parser.parse_args()
+    with Dvi(args.filename, args.dpi) as dvi:
+        fontmap = PsfontsMap(find_tex_file('pdftex.map'))
+        for page in dvi:
+            print(f"=== new page === "
+                  f"(w: {page.width}, h: {page.height}, d: {page.descent})")
+            for font, group in itertools.groupby(
+                    page.text, lambda text: text.font):
+                print(f"font: {font.texname.decode('latin-1')!r}\t"
+                      f"scale: {font._scale / 2 ** 20}")
+                print("x", "y", "glyph", "chr", "w", "(glyphs)", sep="\t")
+                for text in group:
+                    print(text.x, text.y, text.glyph,
+                          chr(text.glyph) if chr(text.glyph).isprintable()
+                          else ".",
+                          text.width, sep="\t")
+            if page.boxes:
+                print("x", "y", "h", "w", "", "(boxes)", sep="\t")
+                for box in page.boxes:
+                    print(box.x, box.y, box.height, box.width, sep="\t")

lib/python3.12/site-packages/matplotlib/dviread.pyi

@@ -0,0 +1,89 @@
+from pathlib import Path
+import io
+import os
+from enum import Enum
+from collections.abc import Generator
+
+from typing import NamedTuple
+from typing_extensions import Self  # < Py 3.11
+
+class _dvistate(Enum):
+    pre = ...
+    outer = ...
+    inpage = ...
+    post_post = ...
+    finale = ...
+
+class Page(NamedTuple):
+    text: list[Text]
+    boxes: list[Box]
+    height: int
+    width: int
+    descent: int
+
+class Box(NamedTuple):
+    x: int
+    y: int
+    height: int
+    width: int
+
+class Text(NamedTuple):
+    x: int
+    y: int
+    font: DviFont
+    glyph: int
+    width: int
+    @property
+    def font_path(self) -> Path: ...
+    @property
+    def font_size(self) -> float: ...
+    @property
+    def font_effects(self) -> dict[str, float]: ...
+    @property
+    def glyph_name_or_index(self) -> int | str: ...
+
+class Dvi:
+    file: io.BufferedReader
+    dpi: float | None
+    fonts: dict[int, DviFont]
+    state: _dvistate
+    def __init__(self, filename: str | os.PathLike, dpi: float | None) -> None: ...
+    def __enter__(self) -> Self: ...
+    def __exit__(self, etype, evalue, etrace) -> None: ...
+    def __iter__(self) -> Generator[Page, None, None]: ...
+    def close(self) -> None: ...
+
+class DviFont:
+    texname: bytes
+    size: float
+    widths: list[int]
+    def __init__(
+        self, scale: float, tfm: Tfm, texname: bytes, vf: Vf | None
+    ) -> None: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __ne__(self, other: object) -> bool: ...
+
+class Vf(Dvi):
+    def __init__(self, filename: str | os.PathLike) -> None: ...
+    def __getitem__(self, code: int) -> Page: ...
+
+class Tfm:
+    checksum: int
+    design_size: int
+    width: dict[int, int]
+    height: dict[int, int]
+    depth: dict[int, int]
+    def __init__(self, filename: str | os.PathLike) -> None: ...
+
+class PsFont(NamedTuple):
+    texname: bytes
+    psname: bytes
+    effects: dict[str, float]
+    encoding: None | bytes
+    filename: str
+
+class PsfontsMap:
+    def __new__(cls, filename: str | os.PathLike) -> Self: ...
+    def __getitem__(self, texname: bytes) -> PsFont: ...
+
+def find_tex_file(filename: str | os.PathLike) -> str: ...

lib/python3.12/site-packages/matplotlib/inset.py

@@ -0,0 +1,269 @@
+"""
+The inset module defines the InsetIndicator class, which draws the rectangle and
+connectors required for `.Axes.indicate_inset` and `.Axes.indicate_inset_zoom`.
+"""
+
+from . import _api, artist, transforms
+from matplotlib.patches import ConnectionPatch, PathPatch, Rectangle
+from matplotlib.path import Path
+
+
+_shared_properties = ('alpha', 'edgecolor', 'linestyle', 'linewidth')
+
+
+class InsetIndicator(artist.Artist):
+    """
+    An artist to highlight an area of interest.
+
+    An inset indicator is a rectangle on the plot at the position indicated by
+    *bounds* that optionally has lines that connect the rectangle to an inset
+    Axes (`.Axes.inset_axes`).
+
+    .. versionadded:: 3.10
+    """
+    zorder = 4.99
+
+    def __init__(self, bounds=None, inset_ax=None, zorder=None, **kwargs):
+        """
+        Parameters
+        ----------
+        bounds : [x0, y0, width, height], optional
+            Lower-left corner of rectangle to be marked, and its width
+            and height.  If not set, the bounds will be calculated from the
+            data limits of inset_ax, which must be supplied.
+
+        inset_ax : `~.axes.Axes`, optional
+            An optional inset Axes to draw connecting lines to.  Two lines are
+            drawn connecting the indicator box to the inset Axes on corners
+            chosen so as to not overlap with the indicator box.
+
+        zorder : float, default: 4.99
+            Drawing order of the rectangle and connector lines.  The default,
+            4.99, is just below the default level of inset Axes.
+
+        **kwargs
+            Other keyword arguments are passed on to the `.Rectangle` patch.
+        """
+        if bounds is None and inset_ax is None:
+            raise ValueError("At least one of bounds or inset_ax must be supplied")
+
+        self._inset_ax = inset_ax
+
+        if bounds is None:
+            # Work out bounds from inset_ax
+            self._auto_update_bounds = True
+            bounds = self._bounds_from_inset_ax()
+        else:
+            self._auto_update_bounds = False
+
+        x, y, width, height = bounds
+
+        self._rectangle = Rectangle((x, y), width, height, clip_on=False, **kwargs)
+
+        # Connector positions cannot be calculated till the artist has been added
+        # to an axes, so just make an empty list for now.
+        self._connectors = []
+
+        super().__init__()
+        self.set_zorder(zorder)
+
+        # Initial style properties for the artist should match the rectangle.
+        for prop in _shared_properties:
+            setattr(self, f'_{prop}', artist.getp(self._rectangle, prop))
+
+    def _shared_setter(self, prop, val):
+        """
+        Helper function to set the same style property on the artist and its children.
+        """
+        setattr(self, f'_{prop}', val)
+
+        artist.setp([self._rectangle, *self._connectors], prop, val)
+
+    def set_alpha(self, alpha):
+        # docstring inherited
+        self._shared_setter('alpha', alpha)
+
+    def set_edgecolor(self, color):
+        """
+        Set the edge color of the rectangle and the connectors.
+
+        Parameters
+        ----------
+        color : :mpltype:`color` or None
+        """
+        self._shared_setter('edgecolor', color)
+
+    def set_color(self, c):
+        """
+        Set the edgecolor of the rectangle and the connectors, and the
+        facecolor for the rectangle.
+
+        Parameters
+        ----------
+        c : :mpltype:`color`
+        """
+        self._shared_setter('edgecolor', c)
+        self._shared_setter('facecolor', c)
+
+    def set_linewidth(self, w):
+        """
+        Set the linewidth in points of the rectangle and the connectors.
+
+        Parameters
+        ----------
+        w : float or None
+        """
+        self._shared_setter('linewidth', w)
+
+    def set_linestyle(self, ls):
+        """
+        Set the linestyle of the rectangle and the connectors.
+
+        ==========================================  =================
+        linestyle                                   description
+        ==========================================  =================
+        ``'-'`` or ``'solid'``                      solid line
+        ``'--'`` or ``'dashed'``                    dashed line
+        ``'-.'`` or ``'dashdot'``                   dash-dotted line
+        ``':'`` or ``'dotted'``                     dotted line
+        ``'none'``, ``'None'``, ``' '``, or ``''``  draw nothing
+        ==========================================  =================
+
+        Alternatively a dash tuple of the following form can be provided::
+
+            (offset, onoffseq)
+
+        where ``onoffseq`` is an even length tuple of on and off ink in points.
+
+        Parameters
+        ----------
+        ls : {'-', '--', '-.', ':', '', (offset, on-off-seq), ...}
+            The line style.
+        """
+        self._shared_setter('linestyle', ls)
+
+    def _bounds_from_inset_ax(self):
+        xlim = self._inset_ax.get_xlim()
+        ylim = self._inset_ax.get_ylim()
+        return (xlim[0], ylim[0], xlim[1] - xlim[0], ylim[1] - ylim[0])
+
+    def _update_connectors(self):
+        (x, y) = self._rectangle.get_xy()
+        width = self._rectangle.get_width()
+        height = self._rectangle.get_height()
+
+        existing_connectors = self._connectors or [None] * 4
+
+        # connect the inset_axes to the rectangle
+        for xy_inset_ax, existing in zip([(0, 0), (0, 1), (1, 0), (1, 1)],
+                                         existing_connectors):
+            # inset_ax positions are in axes coordinates
+            # The 0, 1 values define the four edges if the inset_ax
+            # lower_left, upper_left, lower_right upper_right.
+            ex, ey = xy_inset_ax
+            if self.axes.xaxis.get_inverted():
+                ex = 1 - ex
+            if self.axes.yaxis.get_inverted():
+                ey = 1 - ey
+            xy_data = x + ex * width, y + ey * height
+            if existing is None:
+                # Create new connection patch with styles inherited from the
+                # parent artist.
+                p = ConnectionPatch(
+                    xyA=xy_inset_ax, coordsA=self._inset_ax.transAxes,
+                    xyB=xy_data, coordsB=self.axes.transData,
+                    arrowstyle="-",
+                    edgecolor=self._edgecolor, alpha=self.get_alpha(),
+                    linestyle=self._linestyle, linewidth=self._linewidth)
+                self._connectors.append(p)
+            else:
+                # Only update positioning of existing connection patch.  We
+                # do not want to override any style settings made by the user.
+                existing.xy1 = xy_inset_ax
+                existing.xy2 = xy_data
+                existing.coords1 = self._inset_ax.transAxes
+                existing.coords2 = self.axes.transData
+
+        if existing is None:
+            # decide which two of the lines to keep visible....
+            pos = self._inset_ax.get_position()
+            bboxins = pos.transformed(self.get_figure(root=False).transSubfigure)
+            rectbbox = transforms.Bbox.from_bounds(x, y, width, height).transformed(
+                self._rectangle.get_transform())
+            x0 = rectbbox.x0 < bboxins.x0
+            x1 = rectbbox.x1 < bboxins.x1
+            y0 = rectbbox.y0 < bboxins.y0
+            y1 = rectbbox.y1 < bboxins.y1
+            self._connectors[0].set_visible(x0 ^ y0)
+            self._connectors[1].set_visible(x0 == y1)
+            self._connectors[2].set_visible(x1 == y0)
+            self._connectors[3].set_visible(x1 ^ y1)
+
+    @property
+    def rectangle(self):
+        """`.Rectangle`: the indicator frame."""
+        return self._rectangle
+
+    @property
+    def connectors(self):
+        """
+        4-tuple of `.patches.ConnectionPatch` or None
+            The four connector lines connecting to (lower_left, upper_left,
+            lower_right upper_right) corners of *inset_ax*. Two lines are
+            set with visibility to *False*,  but the user can set the
+            visibility to True if the automatic choice is not deemed correct.
+        """
+        if self._inset_ax is None:
+            return
+
+        if self._auto_update_bounds:
+            self._rectangle.set_bounds(self._bounds_from_inset_ax())
+        self._update_connectors()
+        return tuple(self._connectors)
+
+    def draw(self, renderer):
+        # docstring inherited
+        conn_same_style = []
+
+        # Figure out which connectors have the same style as the box, so should
+        # be drawn as a single path.
+        for conn in self.connectors or []:
+            if conn.get_visible():
+                drawn = False
+                for s in _shared_properties:
+                    if artist.getp(self._rectangle, s) != artist.getp(conn, s):
+                        # Draw this connector by itself
+                        conn.draw(renderer)
+                        drawn = True
+                        break
+
+                if not drawn:
+                    # Connector has same style as box.
+                    conn_same_style.append(conn)
+
+        if conn_same_style:
+            # Since at least one connector has the same style as the rectangle, draw
+            # them as a compound path.
+            artists = [self._rectangle] + conn_same_style
+            paths = [a.get_transform().transform_path(a.get_path()) for a in artists]
+            path = Path.make_compound_path(*paths)
+
+            # Create a temporary patch to draw the path.
+            p = PathPatch(path)
+            p.update_from(self._rectangle)
+            p.set_transform(transforms.IdentityTransform())
+            p.draw(renderer)
+
+            return
+
+        # Just draw the rectangle
+        self._rectangle.draw(renderer)
+
+    @_api.deprecated(
+        '3.10',
+        message=('Since Matplotlib 3.10 indicate_inset_[zoom] returns a single '
+                 'InsetIndicator artist with a rectangle property and a connectors '
+                 'property.  From 3.12 it will no longer be possible to unpack the '
+                 'return value into two elements.'))
+    def __getitem__(self, key):
+        return [self._rectangle, self.connectors][key]

lib/python3.12/site-packages/matplotlib/layout_engine.py

@@ -0,0 +1,309 @@
+"""
+Classes to layout elements in a `.Figure`.
+
+Figures have a ``layout_engine`` property that holds a subclass of
+`~.LayoutEngine` defined here (or *None* for no layout).  At draw time
+``figure.get_layout_engine().execute()`` is called, the goal of which is
+usually to rearrange Axes on the figure to produce a pleasing layout. This is
+like a ``draw`` callback but with two differences.  First, when printing we
+disable the layout engine for the final draw. Second, it is useful to know the
+layout engine while the figure is being created.  In particular, colorbars are
+made differently with different layout engines (for historical reasons).
+
+Matplotlib has two built-in layout engines:
+
+- `.TightLayoutEngine` was the first layout engine added to Matplotlib.
+  See also :ref:`tight_layout_guide`.
+- `.ConstrainedLayoutEngine` is more modern and generally gives better results.
+  See also :ref:`constrainedlayout_guide`.
+
+Third parties can create their own layout engine by subclassing `.LayoutEngine`.
+"""
+
+from contextlib import nullcontext
+
+import matplotlib as mpl
+
+from matplotlib._constrained_layout import do_constrained_layout
+from matplotlib._tight_layout import (get_subplotspec_list,
+                                      get_tight_layout_figure)
+
+
+class LayoutEngine:
+    """
+    Base class for Matplotlib layout engines.
+
+    A layout engine can be passed to a figure at instantiation or at any time
+    with `~.figure.Figure.set_layout_engine`.  Once attached to a figure, the
+    layout engine ``execute`` function is called at draw time by
+    `~.figure.Figure.draw`, providing a special draw-time hook.
+
+    .. note::
+
+       However, note that layout engines affect the creation of colorbars, so
+       `~.figure.Figure.set_layout_engine` should be called before any
+       colorbars are created.
+
+    Currently, there are two properties of `LayoutEngine` classes that are
+    consulted while manipulating the figure:
+
+    - ``engine.colorbar_gridspec`` tells `.Figure.colorbar` whether to make the
+       axes using the gridspec method (see `.colorbar.make_axes_gridspec`) or
+       not (see `.colorbar.make_axes`);
+    - ``engine.adjust_compatible`` stops `.Figure.subplots_adjust` from being
+        run if it is not compatible with the layout engine.
+
+    To implement a custom `LayoutEngine`:
+
+    1. override ``_adjust_compatible`` and ``_colorbar_gridspec``
+    2. override `LayoutEngine.set` to update *self._params*
+    3. override `LayoutEngine.execute` with your implementation
+
+    """
+    # override these in subclass
+    _adjust_compatible = None
+    _colorbar_gridspec = None
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._params = {}
+
+    def set(self, **kwargs):
+        """
+        Set the parameters for the layout engine.
+        """
+        raise NotImplementedError
+
+    @property
+    def colorbar_gridspec(self):
+        """
+        Return a boolean if the layout engine creates colorbars using a
+        gridspec.
+        """
+        if self._colorbar_gridspec is None:
+            raise NotImplementedError
+        return self._colorbar_gridspec
+
+    @property
+    def adjust_compatible(self):
+        """
+        Return a boolean if the layout engine is compatible with
+        `~.Figure.subplots_adjust`.
+        """
+        if self._adjust_compatible is None:
+            raise NotImplementedError
+        return self._adjust_compatible
+
+    def get(self):
+        """
+        Return copy of the parameters for the layout engine.
+        """
+        return dict(self._params)
+
+    def execute(self, fig):
+        """
+        Execute the layout on the figure given by *fig*.
+        """
+        # subclasses must implement this.
+        raise NotImplementedError
+
+
+class PlaceHolderLayoutEngine(LayoutEngine):
+    """
+    This layout engine does not adjust the figure layout at all.
+
+    The purpose of this `.LayoutEngine` is to act as a placeholder when the user removes
+    a layout engine to ensure an incompatible `.LayoutEngine` cannot be set later.
+
+    Parameters
+    ----------
+    adjust_compatible, colorbar_gridspec : bool
+        Allow the PlaceHolderLayoutEngine to mirror the behavior of whatever
+        layout engine it is replacing.
+
+    """
+    def __init__(self, adjust_compatible, colorbar_gridspec, **kwargs):
+        self._adjust_compatible = adjust_compatible
+        self._colorbar_gridspec = colorbar_gridspec
+        super().__init__(**kwargs)
+
+    def execute(self, fig):
+        """
+        Do nothing.
+        """
+        return
+
+
+class TightLayoutEngine(LayoutEngine):
+    """
+    Implements the ``tight_layout`` geometry management.  See
+    :ref:`tight_layout_guide` for details.
+    """
+    _adjust_compatible = True
+    _colorbar_gridspec = True
+
+    def __init__(self, *, pad=1.08, h_pad=None, w_pad=None,
+                 rect=(0, 0, 1, 1), **kwargs):
+        """
+        Initialize tight_layout engine.
+
+        Parameters
+        ----------
+        pad : float, default: 1.08
+            Padding between the figure edge and the edges of subplots, as a
+            fraction of the font size.
+        h_pad, w_pad : float
+            Padding (height/width) between edges of adjacent subplots.
+            Defaults to *pad*.
+        rect : tuple (left, bottom, right, top), default: (0, 0, 1, 1).
+            rectangle in normalized figure coordinates that the subplots
+            (including labels) will fit into.
+        """
+        super().__init__(**kwargs)
+        for td in ['pad', 'h_pad', 'w_pad', 'rect']:
+            # initialize these in case None is passed in above:
+            self._params[td] = None
+        self.set(pad=pad, h_pad=h_pad, w_pad=w_pad, rect=rect)
+
+    def execute(self, fig):
+        """
+        Execute tight_layout.
+
+        This decides the subplot parameters given the padding that
+        will allow the Axes labels to not be covered by other labels
+        and Axes.
+
+        Parameters
+        ----------
+        fig : `.Figure` to perform layout on.
+
+        See Also
+        --------
+        .figure.Figure.tight_layout
+        .pyplot.tight_layout
+        """
+        info = self._params
+        renderer = fig._get_renderer()
+        with getattr(renderer, "_draw_disabled", nullcontext)():
+            kwargs = get_tight_layout_figure(
+                fig, fig.axes, get_subplotspec_list(fig.axes), renderer,
+                pad=info['pad'], h_pad=info['h_pad'], w_pad=info['w_pad'],
+                rect=info['rect'])
+        if kwargs:
+            fig.subplots_adjust(**kwargs)
+
+    def set(self, *, pad=None, w_pad=None, h_pad=None, rect=None):
+        """
+        Set the pads for tight_layout.
+
+        Parameters
+        ----------
+        pad : float
+            Padding between the figure edge and the edges of subplots, as a
+            fraction of the font size.
+        w_pad, h_pad : float
+            Padding (width/height) between edges of adjacent subplots.
+            Defaults to *pad*.
+        rect : tuple (left, bottom, right, top)
+            rectangle in normalized figure coordinates that the subplots
+            (including labels) will fit into.
+        """
+        for td in self.set.__kwdefaults__:
+            if locals()[td] is not None:
+                self._params[td] = locals()[td]
+
+
+class ConstrainedLayoutEngine(LayoutEngine):
+    """
+    Implements the ``constrained_layout`` geometry management.  See
+    :ref:`constrainedlayout_guide` for details.
+    """
+
+    _adjust_compatible = False
+    _colorbar_gridspec = False
+
+    def __init__(self, *, h_pad=None, w_pad=None,
+                 hspace=None, wspace=None, rect=(0, 0, 1, 1),
+                 compress=False, **kwargs):
+        """
+        Initialize ``constrained_layout`` settings.
+
+        Parameters
+        ----------
+        h_pad, w_pad : float
+            Padding around the Axes elements in inches.
+            Default to :rc:`figure.constrained_layout.h_pad` and
+            :rc:`figure.constrained_layout.w_pad`.
+        hspace, wspace : float
+            Fraction of the figure to dedicate to space between the
+            axes.  These are evenly spread between the gaps between the Axes.
+            A value of 0.2 for a three-column layout would have a space
+            of 0.1 of the figure width between each column.
+            If h/wspace < h/w_pad, then the pads are used instead.
+            Default to :rc:`figure.constrained_layout.hspace` and
+            :rc:`figure.constrained_layout.wspace`.
+        rect : tuple of 4 floats
+            Rectangle in figure coordinates to perform constrained layout in
+            (left, bottom, width, height), each from 0-1.
+        compress : bool
+            Whether to shift Axes so that white space in between them is
+            removed. This is useful for simple grids of fixed-aspect Axes (e.g.
+            a grid of images).  See :ref:`compressed_layout`.
+        """
+        super().__init__(**kwargs)
+        # set the defaults:
+        self.set(w_pad=mpl.rcParams['figure.constrained_layout.w_pad'],
+                 h_pad=mpl.rcParams['figure.constrained_layout.h_pad'],
+                 wspace=mpl.rcParams['figure.constrained_layout.wspace'],
+                 hspace=mpl.rcParams['figure.constrained_layout.hspace'],
+                 rect=(0, 0, 1, 1))
+        # set anything that was passed in (None will be ignored):
+        self.set(w_pad=w_pad, h_pad=h_pad, wspace=wspace, hspace=hspace,
+                 rect=rect)
+        self._compress = compress
+
+    def execute(self, fig):
+        """
+        Perform constrained_layout and move and resize Axes accordingly.
+
+        Parameters
+        ----------
+        fig : `.Figure` to perform layout on.
+        """
+        width, height = fig.get_size_inches()
+        # pads are relative to the current state of the figure...
+        w_pad = self._params['w_pad'] / width
+        h_pad = self._params['h_pad'] / height
+
+        return do_constrained_layout(fig, w_pad=w_pad, h_pad=h_pad,
+                                     wspace=self._params['wspace'],
+                                     hspace=self._params['hspace'],
+                                     rect=self._params['rect'],
+                                     compress=self._compress)
+
+    def set(self, *, h_pad=None, w_pad=None,
+            hspace=None, wspace=None, rect=None):
+        """
+        Set the pads for constrained_layout.
+
+        Parameters
+        ----------
+        h_pad, w_pad : float
+            Padding around the Axes elements in inches.
+            Default to :rc:`figure.constrained_layout.h_pad` and
+            :rc:`figure.constrained_layout.w_pad`.
+        hspace, wspace : float
+            Fraction of the figure to dedicate to space between the
+            axes.  These are evenly spread between the gaps between the Axes.
+            A value of 0.2 for a three-column layout would have a space
+            of 0.1 of the figure width between each column.
+            If h/wspace < h/w_pad, then the pads are used instead.
+            Default to :rc:`figure.constrained_layout.hspace` and
+            :rc:`figure.constrained_layout.wspace`.
+        rect : tuple of 4 floats
+            Rectangle in figure coordinates to perform constrained layout in
+            (left, bottom, width, height), each from 0-1.
+        """
+        for td in self.set.__kwdefaults__:
+            if locals()[td] is not None:
+                self._params[td] = locals()[td]

lib/python3.12/site-packages/matplotlib/legend.py

@@ -0,0 +1,1370 @@
+"""
+The legend module defines the Legend class, which is responsible for
+drawing legends associated with Axes and/or figures.
+
+.. important::
+
+    It is unlikely that you would ever create a Legend instance manually.
+    Most users would normally create a legend via the `~.Axes.legend`
+    function. For more details on legends there is also a :ref:`legend guide
+    <legend_guide>`.
+
+The `Legend` class is a container of legend handles and legend texts.
+
+The legend handler map specifies how to create legend handles from artists
+(lines, patches, etc.) in the Axes or figures. Default legend handlers are
+defined in the :mod:`~matplotlib.legend_handler` module. While not all artist
+types are covered by the default legend handlers, custom legend handlers can be
+defined to support arbitrary objects.
+
+See the :ref`<legend_guide>` for more
+information.
+"""
+
+import itertools
+import logging
+import numbers
+import time
+
+import numpy as np
+
+import matplotlib as mpl
+from matplotlib import _api, _docstring, cbook, colors, offsetbox
+from matplotlib.artist import Artist, allow_rasterization
+from matplotlib.cbook import silent_list
+from matplotlib.font_manager import FontProperties
+from matplotlib.lines import Line2D
+from matplotlib.patches import (Patch, Rectangle, Shadow, FancyBboxPatch,
+                                StepPatch)
+from matplotlib.collections import (
+    Collection, CircleCollection, LineCollection, PathCollection,
+    PolyCollection, RegularPolyCollection)
+from matplotlib.text import Text
+from matplotlib.transforms import Bbox, BboxBase, TransformedBbox
+from matplotlib.transforms import BboxTransformTo, BboxTransformFrom
+from matplotlib.offsetbox import (
+    AnchoredOffsetbox, DraggableOffsetBox,
+    HPacker, VPacker,
+    DrawingArea, TextArea,
+)
+from matplotlib.container import ErrorbarContainer, BarContainer, StemContainer
+from . import legend_handler
+
+
+class DraggableLegend(DraggableOffsetBox):
+    def __init__(self, legend, use_blit=False, update="loc"):
+        """
+        Wrapper around a `.Legend` to support mouse dragging.
+
+        Parameters
+        ----------
+        legend : `.Legend`
+            The `.Legend` instance to wrap.
+        use_blit : bool, optional
+            Use blitting for faster image composition. For details see
+            :ref:`func-animation`.
+        update : {'loc', 'bbox'}, optional
+            If "loc", update the *loc* parameter of the legend upon finalizing.
+            If "bbox", update the *bbox_to_anchor* parameter.
+        """
+        self.legend = legend
+
+        _api.check_in_list(["loc", "bbox"], update=update)
+        self._update = update
+
+        super().__init__(legend, legend._legend_box, use_blit=use_blit)
+
+    def finalize_offset(self):
+        if self._update == "loc":
+            self._update_loc(self.get_loc_in_canvas())
+        elif self._update == "bbox":
+            self._update_bbox_to_anchor(self.get_loc_in_canvas())
+
+    def _update_loc(self, loc_in_canvas):
+        bbox = self.legend.get_bbox_to_anchor()
+        # if bbox has zero width or height, the transformation is
+        # ill-defined. Fall back to the default bbox_to_anchor.
+        if bbox.width == 0 or bbox.height == 0:
+            self.legend.set_bbox_to_anchor(None)
+            bbox = self.legend.get_bbox_to_anchor()
+        _bbox_transform = BboxTransformFrom(bbox)
+        self.legend._loc = tuple(_bbox_transform.transform(loc_in_canvas))
+
+    def _update_bbox_to_anchor(self, loc_in_canvas):
+        loc_in_bbox = self.legend.axes.transAxes.transform(loc_in_canvas)
+        self.legend.set_bbox_to_anchor(loc_in_bbox)
+
+
+_legend_kw_doc_base = """
+bbox_to_anchor : `.BboxBase`, 2-tuple, or 4-tuple of floats
+    Box that is used to position the legend in conjunction with *loc*.
+    Defaults to ``axes.bbox`` (if called as a method to `.Axes.legend`) or
+    ``figure.bbox`` (if ``figure.legend``).  This argument allows arbitrary
+    placement of the legend.
+
+    Bbox coordinates are interpreted in the coordinate system given by
+    *bbox_transform*, with the default transform
+    Axes or Figure coordinates, depending on which ``legend`` is called.
+
+    If a 4-tuple or `.BboxBase` is given, then it specifies the bbox
+    ``(x, y, width, height)`` that the legend is placed in.
+    To put the legend in the best location in the bottom right
+    quadrant of the Axes (or figure)::
+
+        loc='best', bbox_to_anchor=(0.5, 0., 0.5, 0.5)
+
+    A 2-tuple ``(x, y)`` places the corner of the legend specified by *loc* at
+    x, y.  For example, to put the legend's upper right-hand corner in the
+    center of the Axes (or figure) the following keywords can be used::
+
+        loc='upper right', bbox_to_anchor=(0.5, 0.5)
+
+ncols : int, default: 1
+    The number of columns that the legend has.
+
+    For backward compatibility, the spelling *ncol* is also supported
+    but it is discouraged. If both are given, *ncols* takes precedence.
+
+prop : None or `~matplotlib.font_manager.FontProperties` or dict
+    The font properties of the legend. If None (default), the current
+    :data:`matplotlib.rcParams` will be used.
+
+fontsize : int or {'xx-small', 'x-small', 'small', 'medium', 'large', \
+'x-large', 'xx-large'}
+    The font size of the legend. If the value is numeric the size will be the
+    absolute font size in points. String values are relative to the current
+    default font size. This argument is only used if *prop* is not specified.
+
+labelcolor : str or list, default: :rc:`legend.labelcolor`
+    The color of the text in the legend. Either a valid color string
+    (for example, 'red'), or a list of color strings. The labelcolor can
+    also be made to match the color of the line or marker using 'linecolor',
+    'markerfacecolor' (or 'mfc'), or 'markeredgecolor' (or 'mec').
+
+    Labelcolor can be set globally using :rc:`legend.labelcolor`. If None,
+    use :rc:`text.color`.
+
+numpoints : int, default: :rc:`legend.numpoints`
+    The number of marker points in the legend when creating a legend
+    entry for a `.Line2D` (line).
+
+scatterpoints : int, default: :rc:`legend.scatterpoints`
+    The number of marker points in the legend when creating
+    a legend entry for a `.PathCollection` (scatter plot).
+
+scatteryoffsets : iterable of floats, default: ``[0.375, 0.5, 0.3125]``
+    The vertical offset (relative to the font size) for the markers
+    created for a scatter plot legend entry. 0.0 is at the base the
+    legend text, and 1.0 is at the top. To draw all markers at the
+    same height, set to ``[0.5]``.
+
+markerscale : float, default: :rc:`legend.markerscale`
+    The relative size of legend markers compared to the originally drawn ones.
+
+markerfirst : bool, default: True
+    If *True*, legend marker is placed to the left of the legend label.
+    If *False*, legend marker is placed to the right of the legend label.
+
+reverse : bool, default: False
+    If *True*, the legend labels are displayed in reverse order from the input.
+    If *False*, the legend labels are displayed in the same order as the input.
+
+    .. versionadded:: 3.7
+
+frameon : bool, default: :rc:`legend.frameon`
+    Whether the legend should be drawn on a patch (frame).
+
+fancybox : bool, default: :rc:`legend.fancybox`
+    Whether round edges should be enabled around the `.FancyBboxPatch` which
+    makes up the legend's background.
+
+shadow : None, bool or dict, default: :rc:`legend.shadow`
+    Whether to draw a shadow behind the legend.
+    The shadow can be configured using `.Patch` keywords.
+    Customization via :rc:`legend.shadow` is currently not supported.
+
+framealpha : float, default: :rc:`legend.framealpha`
+    The alpha transparency of the legend's background.
+    If *shadow* is activated and *framealpha* is ``None``, the default value is
+    ignored.
+
+facecolor : "inherit" or color, default: :rc:`legend.facecolor`
+    The legend's background color.
+    If ``"inherit"``, use :rc:`axes.facecolor`.
+
+edgecolor : "inherit" or color, default: :rc:`legend.edgecolor`
+    The legend's background patch edge color.
+    If ``"inherit"``, use :rc:`axes.edgecolor`.
+
+mode : {"expand", None}
+    If *mode* is set to ``"expand"`` the legend will be horizontally
+    expanded to fill the Axes area (or *bbox_to_anchor* if defines
+    the legend's size).
+
+bbox_transform : None or `~matplotlib.transforms.Transform`
+    The transform for the bounding box (*bbox_to_anchor*). For a value
+    of ``None`` (default) the Axes'
+    :data:`~matplotlib.axes.Axes.transAxes` transform will be used.
+
+title : str or None
+    The legend's title. Default is no title (``None``).
+
+title_fontproperties : None or `~matplotlib.font_manager.FontProperties` or dict
+    The font properties of the legend's title. If None (default), the
+    *title_fontsize* argument will be used if present; if *title_fontsize* is
+    also None, the current :rc:`legend.title_fontsize` will be used.
+
+title_fontsize : int or {'xx-small', 'x-small', 'small', 'medium', 'large', \
+'x-large', 'xx-large'}, default: :rc:`legend.title_fontsize`
+    The font size of the legend's title.
+    Note: This cannot be combined with *title_fontproperties*. If you want
+    to set the fontsize alongside other font properties, use the *size*
+    parameter in *title_fontproperties*.
+
+alignment : {'center', 'left', 'right'}, default: 'center'
+    The alignment of the legend title and the box of entries. The entries
+    are aligned as a single block, so that markers always lined up.
+
+borderpad : float, default: :rc:`legend.borderpad`
+    The fractional whitespace inside the legend border, in font-size units.
+
+labelspacing : float, default: :rc:`legend.labelspacing`
+    The vertical space between the legend entries, in font-size units.
+
+handlelength : float, default: :rc:`legend.handlelength`
+    The length of the legend handles, in font-size units.
+
+handleheight : float, default: :rc:`legend.handleheight`
+    The height of the legend handles, in font-size units.
+
+handletextpad : float, default: :rc:`legend.handletextpad`
+    The pad between the legend handle and text, in font-size units.
+
+borderaxespad : float, default: :rc:`legend.borderaxespad`
+    The pad between the Axes and legend border, in font-size units.
+
+columnspacing : float, default: :rc:`legend.columnspacing`
+    The spacing between columns, in font-size units.
+
+handler_map : dict or None
+    The custom dictionary mapping instances or types to a legend
+    handler. This *handler_map* updates the default handler map
+    found at `matplotlib.legend.Legend.get_legend_handler_map`.
+
+draggable : bool, default: False
+    Whether the legend can be dragged with the mouse.
+"""
+
+_loc_doc_base = """
+loc : str or pair of floats, default: {default}
+    The location of the legend.
+
+    The strings ``'upper left'``, ``'upper right'``, ``'lower left'``,
+    ``'lower right'`` place the legend at the corresponding corner of the
+    {parent}.
+
+    The strings ``'upper center'``, ``'lower center'``, ``'center left'``,
+    ``'center right'`` place the legend at the center of the corresponding edge
+    of the {parent}.
+
+    The string ``'center'`` places the legend at the center of the {parent}.
+{best}
+    The location can also be a 2-tuple giving the coordinates of the lower-left
+    corner of the legend in {parent} coordinates (in which case *bbox_to_anchor*
+    will be ignored).
+
+    For back-compatibility, ``'center right'`` (but no other location) can also
+    be spelled ``'right'``, and each "string" location can also be given as a
+    numeric value:
+
+    ==================   =============
+    Location String      Location Code
+    ==================   =============
+    'best' (Axes only)   0
+    'upper right'        1
+    'upper left'         2
+    'lower left'         3
+    'lower right'        4
+    'right'              5
+    'center left'        6
+    'center right'       7
+    'lower center'       8
+    'upper center'       9
+    'center'             10
+    ==================   =============
+    {outside}"""
+
+_loc_doc_best = """
+    The string ``'best'`` places the legend at the location, among the nine
+    locations defined so far, with the minimum overlap with other drawn
+    artists.  This option can be quite slow for plots with large amounts of
+    data; your plotting speed may benefit from providing a specific location.
+"""
+
+_legend_kw_axes_st = (
+    _loc_doc_base.format(parent='axes', default=':rc:`legend.loc`',
+                         best=_loc_doc_best, outside='') +
+    _legend_kw_doc_base)
+_docstring.interpd.register(_legend_kw_axes=_legend_kw_axes_st)
+
+_outside_doc = """
+    If a figure is using the constrained layout manager, the string codes
+    of the *loc* keyword argument can get better layout behaviour using the
+    prefix 'outside'. There is ambiguity at the corners, so 'outside
+    upper right' will make space for the legend above the rest of the
+    axes in the layout, and 'outside right upper' will make space on the
+    right side of the layout.  In addition to the values of *loc*
+    listed above, we have 'outside right upper', 'outside right lower',
+    'outside left upper', and 'outside left lower'.  See
+    :ref:`legend_guide` for more details.
+"""
+
+_legend_kw_figure_st = (
+    _loc_doc_base.format(parent='figure', default="'upper right'",
+                         best='', outside=_outside_doc) +
+    _legend_kw_doc_base)
+_docstring.interpd.register(_legend_kw_figure=_legend_kw_figure_st)
+
+_legend_kw_both_st = (
+    _loc_doc_base.format(parent='axes/figure',
+                         default=":rc:`legend.loc` for Axes, 'upper right' for Figure",
+                         best=_loc_doc_best, outside=_outside_doc) +
+    _legend_kw_doc_base)
+_docstring.interpd.register(_legend_kw_doc=_legend_kw_both_st)
+
+_legend_kw_set_loc_st = (
+    _loc_doc_base.format(parent='axes/figure',
+                         default=":rc:`legend.loc` for Axes, 'upper right' for Figure",
+                         best=_loc_doc_best, outside=_outside_doc))
+_docstring.interpd.register(_legend_kw_set_loc_doc=_legend_kw_set_loc_st)
+
+
+class Legend(Artist):
+    """
+    Place a legend on the figure/axes.
+    """
+
+    # 'best' is only implemented for Axes legends
+    codes = {'best': 0, **AnchoredOffsetbox.codes}
+    zorder = 5
+
+    def __str__(self):
+        return "Legend"
+
+    @_docstring.interpd
+    def __init__(
+        self, parent, handles, labels,
+        *,
+        loc=None,
+        numpoints=None,      # number of points in the legend line
+        markerscale=None,    # relative size of legend markers vs. original
+        markerfirst=True,    # left/right ordering of legend marker and label
+        reverse=False,       # reverse ordering of legend marker and label
+        scatterpoints=None,  # number of scatter points
+        scatteryoffsets=None,
+        prop=None,           # properties for the legend texts
+        fontsize=None,       # keyword to set font size directly
+        labelcolor=None,     # keyword to set the text color
+
+        # spacing & pad defined as a fraction of the font-size
+        borderpad=None,      # whitespace inside the legend border
+        labelspacing=None,   # vertical space between the legend entries
+        handlelength=None,   # length of the legend handles
+        handleheight=None,   # height of the legend handles
+        handletextpad=None,  # pad between the legend handle and text
+        borderaxespad=None,  # pad between the Axes and legend border
+        columnspacing=None,  # spacing between columns
+
+        ncols=1,     # number of columns
+        mode=None,  # horizontal distribution of columns: None or "expand"
+
+        fancybox=None,  # True: fancy box, False: rounded box, None: rcParam
+        shadow=None,
+        title=None,           # legend title
+        title_fontsize=None,  # legend title font size
+        framealpha=None,      # set frame alpha
+        edgecolor=None,       # frame patch edgecolor
+        facecolor=None,       # frame patch facecolor
+
+        bbox_to_anchor=None,  # bbox to which the legend will be anchored
+        bbox_transform=None,  # transform for the bbox
+        frameon=None,         # draw frame
+        handler_map=None,
+        title_fontproperties=None,  # properties for the legend title
+        alignment="center",       # control the alignment within the legend box
+        ncol=1,  # synonym for ncols (backward compatibility)
+        draggable=False  # whether the legend can be dragged with the mouse
+    ):
+        """
+        Parameters
+        ----------
+        parent : `~matplotlib.axes.Axes` or `.Figure`
+            The artist that contains the legend.
+
+        handles : list of (`.Artist` or tuple of `.Artist`)
+            A list of Artists (lines, patches) to be added to the legend.
+
+        labels : list of str
+            A list of labels to show next to the artists. The length of handles
+            and labels should be the same. If they are not, they are truncated
+            to the length of the shorter list.
+
+        Other Parameters
+        ----------------
+        %(_legend_kw_doc)s
+
+        Attributes
+        ----------
+        legend_handles
+            List of `.Artist` objects added as legend entries.
+
+            .. versionadded:: 3.7
+        """
+        # local import only to avoid circularity
+        from matplotlib.axes import Axes
+        from matplotlib.figure import FigureBase
+
+        super().__init__()
+
+        if prop is None:
+            self.prop = FontProperties(size=mpl._val_or_rc(fontsize, "legend.fontsize"))
+        else:
+            self.prop = FontProperties._from_any(prop)
+            if isinstance(prop, dict) and "size" not in prop:
+                self.prop.set_size(mpl.rcParams["legend.fontsize"])
+
+        self._fontsize = self.prop.get_size_in_points()
+
+        self.texts = []
+        self.legend_handles = []
+        self._legend_title_box = None
+
+        #: A dictionary with the extra handler mappings for this Legend
+        #: instance.
+        self._custom_handler_map = handler_map
+
+        self.numpoints = mpl._val_or_rc(numpoints, 'legend.numpoints')
+        self.markerscale = mpl._val_or_rc(markerscale, 'legend.markerscale')
+        self.scatterpoints = mpl._val_or_rc(scatterpoints, 'legend.scatterpoints')
+        self.borderpad = mpl._val_or_rc(borderpad, 'legend.borderpad')
+        self.labelspacing = mpl._val_or_rc(labelspacing, 'legend.labelspacing')
+        self.handlelength = mpl._val_or_rc(handlelength, 'legend.handlelength')
+        self.handleheight = mpl._val_or_rc(handleheight, 'legend.handleheight')
+        self.handletextpad = mpl._val_or_rc(handletextpad, 'legend.handletextpad')
+        self.borderaxespad = mpl._val_or_rc(borderaxespad, 'legend.borderaxespad')
+        self.columnspacing = mpl._val_or_rc(columnspacing, 'legend.columnspacing')
+        self.shadow = mpl._val_or_rc(shadow, 'legend.shadow')
+
+        if reverse:
+            labels = [*reversed(labels)]
+            handles = [*reversed(handles)]
+
+        if len(handles) < 2:
+            ncols = 1
+        self._ncols = ncols if ncols != 1 else ncol
+
+        if self.numpoints <= 0:
+            raise ValueError("numpoints must be > 0; it was %d" % numpoints)
+
+        # introduce y-offset for handles of the scatter plot
+        if scatteryoffsets is None:
+            self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.])
+        else:
+            self._scatteryoffsets = np.asarray(scatteryoffsets)
+        reps = self.scatterpoints // len(self._scatteryoffsets) + 1
+        self._scatteryoffsets = np.tile(self._scatteryoffsets,
+                                        reps)[:self.scatterpoints]
+
+        # _legend_box is a VPacker instance that contains all
+        # legend items and will be initialized from _init_legend_box()
+        # method.
+        self._legend_box = None
+
+        if isinstance(parent, Axes):
+            self.isaxes = True
+            self.axes = parent
+            self.set_figure(parent.get_figure(root=False))
+        elif isinstance(parent, FigureBase):
+            self.isaxes = False
+            self.set_figure(parent)
+        else:
+            raise TypeError(
+                "Legend needs either Axes or FigureBase as parent"
+            )
+        self.parent = parent
+
+        self._mode = mode
+        self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform)
+
+        # Figure out if self.shadow is valid
+        # If shadow was None, rcParams loads False
+        # So it shouldn't be None here
+
+        self._shadow_props = {'ox': 2, 'oy': -2}  # default location offsets
+        if isinstance(self.shadow, dict):
+            self._shadow_props.update(self.shadow)
+            self.shadow = True
+        elif self.shadow in (0, 1, True, False):
+            self.shadow = bool(self.shadow)
+        else:
+            raise ValueError(
+                'Legend shadow must be a dict or bool, not '
+                f'{self.shadow!r} of type {type(self.shadow)}.'
+            )
+
+        # We use FancyBboxPatch to draw a legend frame. The location
+        # and size of the box will be updated during the drawing time.
+
+        facecolor = mpl._val_or_rc(facecolor, "legend.facecolor")
+        if facecolor == 'inherit':
+            facecolor = mpl.rcParams["axes.facecolor"]
+
+        edgecolor = mpl._val_or_rc(edgecolor, "legend.edgecolor")
+        if edgecolor == 'inherit':
+            edgecolor = mpl.rcParams["axes.edgecolor"]
+
+        fancybox = mpl._val_or_rc(fancybox, "legend.fancybox")
+
+        self.legendPatch = FancyBboxPatch(
+            xy=(0, 0), width=1, height=1,
+            facecolor=facecolor, edgecolor=edgecolor,
+            # If shadow is used, default to alpha=1 (#8943).
+            alpha=(framealpha if framealpha is not None
+                   else 1 if shadow
+                   else mpl.rcParams["legend.framealpha"]),
+            # The width and height of the legendPatch will be set (in draw())
+            # to the length that includes the padding. Thus we set pad=0 here.
+            boxstyle=("round,pad=0,rounding_size=0.2" if fancybox
+                      else "square,pad=0"),
+            mutation_scale=self._fontsize,
+            snap=True,
+            visible=mpl._val_or_rc(frameon, "legend.frameon")
+        )
+        self._set_artist_props(self.legendPatch)
+
+        _api.check_in_list(["center", "left", "right"], alignment=alignment)
+        self._alignment = alignment
+
+        # init with null renderer
+        self._init_legend_box(handles, labels, markerfirst)
+
+        # Set legend location
+        self.set_loc(loc)
+
+        # figure out title font properties:
+        if title_fontsize is not None and title_fontproperties is not None:
+            raise ValueError(
+                "title_fontsize and title_fontproperties can't be specified "
+                "at the same time. Only use one of them. ")
+        title_prop_fp = FontProperties._from_any(title_fontproperties)
+        if isinstance(title_fontproperties, dict):
+            if "size" not in title_fontproperties:
+                title_fontsize = mpl.rcParams["legend.title_fontsize"]
+                title_prop_fp.set_size(title_fontsize)
+        elif title_fontsize is not None:
+            title_prop_fp.set_size(title_fontsize)
+        elif not isinstance(title_fontproperties, FontProperties):
+            title_fontsize = mpl.rcParams["legend.title_fontsize"]
+            title_prop_fp.set_size(title_fontsize)
+
+        self.set_title(title, prop=title_prop_fp)
+
+        self._draggable = None
+        self.set_draggable(state=draggable)
+
+        # set the text color
+
+        color_getters = {  # getter function depends on line or patch
+            'linecolor':       ['get_color',           'get_facecolor'],
+            'markerfacecolor': ['get_markerfacecolor', 'get_facecolor'],
+            'mfc':             ['get_markerfacecolor', 'get_facecolor'],
+            'markeredgecolor': ['get_markeredgecolor', 'get_edgecolor'],
+            'mec':             ['get_markeredgecolor', 'get_edgecolor'],
+        }
+        labelcolor = mpl._val_or_rc(labelcolor, 'legend.labelcolor')
+        if labelcolor is None:
+            labelcolor = mpl.rcParams['text.color']
+        if isinstance(labelcolor, str) and labelcolor in color_getters:
+            getter_names = color_getters[labelcolor]
+            for handle, text in zip(self.legend_handles, self.texts):
+                try:
+                    if handle.get_array() is not None:
+                        continue
+                except AttributeError:
+                    pass
+                for getter_name in getter_names:
+                    try:
+                        color = getattr(handle, getter_name)()
+                        if isinstance(color, np.ndarray):
+                            if (
+                                    color.shape[0] == 1
+                                    or np.isclose(color, color[0]).all()
+                            ):
+                                text.set_color(color[0])
+                            else:
+                                pass
+                        else:
+                            text.set_color(color)
+                        break
+                    except AttributeError:
+                        pass
+        elif cbook._str_equal(labelcolor, 'none'):
+            for text in self.texts:
+                text.set_color(labelcolor)
+        elif np.iterable(labelcolor):
+            for text, color in zip(self.texts,
+                                   itertools.cycle(
+                                       colors.to_rgba_array(labelcolor))):
+                text.set_color(color)
+        else:
+            raise ValueError(f"Invalid labelcolor: {labelcolor!r}")
+
+    def _set_artist_props(self, a):
+        """
+        Set the boilerplate props for artists added to Axes.
+        """
+        a.set_figure(self.get_figure(root=False))
+        if self.isaxes:
+            a.axes = self.axes
+
+        a.set_transform(self.get_transform())
+
+    @_docstring.interpd
+    def set_loc(self, loc=None):
+        """
+        Set the location of the legend.
+
+        .. versionadded:: 3.8
+
+        Parameters
+        ----------
+        %(_legend_kw_set_loc_doc)s
+        """
+        loc0 = loc
+        self._loc_used_default = loc is None
+        if loc is None:
+            loc = mpl.rcParams["legend.loc"]
+            if not self.isaxes and loc in [0, 'best']:
+                loc = 'upper right'
+
+        type_err_message = ("loc must be string, coordinate tuple, or"
+                            f" an integer 0-10, not {loc!r}")
+
+        # handle outside legends:
+        self._outside_loc = None
+        if isinstance(loc, str):
+            if loc.split()[0] == 'outside':
+                # strip outside:
+                loc = loc.split('outside ')[1]
+                # strip "center" at the beginning
+                self._outside_loc = loc.replace('center ', '')
+                # strip first
+                self._outside_loc = self._outside_loc.split()[0]
+                locs = loc.split()
+                if len(locs) > 1 and locs[0] in ('right', 'left'):
+                    # locs doesn't accept "left upper", etc, so swap
+                    if locs[0] != 'center':
+                        locs = locs[::-1]
+                    loc = locs[0] + ' ' + locs[1]
+            # check that loc is in acceptable strings
+            loc = _api.check_getitem(self.codes, loc=loc)
+        elif np.iterable(loc):
+            # coerce iterable into tuple
+            loc = tuple(loc)
+            # validate the tuple represents Real coordinates
+            if len(loc) != 2 or not all(isinstance(e, numbers.Real) for e in loc):
+                raise ValueError(type_err_message)
+        elif isinstance(loc, int):
+            # validate the integer represents a string numeric value
+            if loc < 0 or loc > 10:
+                raise ValueError(type_err_message)
+        else:
+            # all other cases are invalid values of loc
+            raise ValueError(type_err_message)
+
+        if self.isaxes and self._outside_loc:
+            raise ValueError(
+                f"'outside' option for loc='{loc0}' keyword argument only "
+                "works for figure legends")
+
+        if not self.isaxes and loc == 0:
+            raise ValueError(
+                "Automatic legend placement (loc='best') not implemented for "
+                "figure legend")
+
+        tmp = self._loc_used_default
+        self._set_loc(loc)
+        self._loc_used_default = tmp  # ignore changes done by _set_loc
+
+    def _set_loc(self, loc):
+        # find_offset function will be provided to _legend_box and
+        # _legend_box will draw itself at the location of the return
+        # value of the find_offset.
+        self._loc_used_default = False
+        self._loc_real = loc
+        self.stale = True
+        self._legend_box.set_offset(self._findoffset)
+
+    def set_ncols(self, ncols):
+        """Set the number of columns."""
+        self._ncols = ncols
+
+    def _get_loc(self):
+        return self._loc_real
+
+    _loc = property(_get_loc, _set_loc)
+
+    def _findoffset(self, width, height, xdescent, ydescent, renderer):
+        """Helper function to locate the legend."""
+
+        if self._loc == 0:  # "best".
+            x, y = self._find_best_position(width, height, renderer)
+        elif self._loc in Legend.codes.values():  # Fixed location.
+            bbox = Bbox.from_bounds(0, 0, width, height)
+            x, y = self._get_anchored_bbox(self._loc, bbox,
+                                           self.get_bbox_to_anchor(),
+                                           renderer)
+        else:  # Axes or figure coordinates.
+            fx, fy = self._loc
+            bbox = self.get_bbox_to_anchor()
+            x, y = bbox.x0 + bbox.width * fx, bbox.y0 + bbox.height * fy
+
+        return x + xdescent, y + ydescent
+
+    @allow_rasterization
+    def draw(self, renderer):
+        # docstring inherited
+        if not self.get_visible():
+            return
+
+        renderer.open_group('legend', gid=self.get_gid())
+
+        fontsize = renderer.points_to_pixels(self._fontsize)
+
+        # if mode == fill, set the width of the legend_box to the
+        # width of the parent (minus pads)
+        if self._mode in ["expand"]:
+            pad = 2 * (self.borderaxespad + self.borderpad) * fontsize
+            self._legend_box.set_width(self.get_bbox_to_anchor().width - pad)
+
+        # update the location and size of the legend. This needs to
+        # be done in any case to clip the figure right.
+        bbox = self._legend_box.get_window_extent(renderer)
+        self.legendPatch.set_bounds(bbox.bounds)
+        self.legendPatch.set_mutation_scale(fontsize)
+
+        # self.shadow is validated in __init__
+        # So by here it is a bool and self._shadow_props contains any configs
+
+        if self.shadow:
+            Shadow(self.legendPatch, **self._shadow_props).draw(renderer)
+
+        self.legendPatch.draw(renderer)
+        self._legend_box.draw(renderer)
+
+        renderer.close_group('legend')
+        self.stale = False
+
+    # _default_handler_map defines the default mapping between plot
+    # elements and the legend handlers.
+
+    _default_handler_map = {
+        StemContainer: legend_handler.HandlerStem(),
+        ErrorbarContainer: legend_handler.HandlerErrorbar(),
+        Line2D: legend_handler.HandlerLine2D(),
+        Patch: legend_handler.HandlerPatch(),
+        StepPatch: legend_handler.HandlerStepPatch(),
+        LineCollection: legend_handler.HandlerLineCollection(),
+        RegularPolyCollection: legend_handler.HandlerRegularPolyCollection(),
+        CircleCollection: legend_handler.HandlerCircleCollection(),
+        BarContainer: legend_handler.HandlerPatch(
+            update_func=legend_handler.update_from_first_child),
+        tuple: legend_handler.HandlerTuple(),
+        PathCollection: legend_handler.HandlerPathCollection(),
+        PolyCollection: legend_handler.HandlerPolyCollection()
+        }
+
+    # (get|set|update)_default_handler_maps are public interfaces to
+    # modify the default handler map.
+
+    @classmethod
+    def get_default_handler_map(cls):
+        """Return the global default handler map, shared by all legends."""
+        return cls._default_handler_map
+
+    @classmethod
+    def set_default_handler_map(cls, handler_map):
+        """Set the global default handler map, shared by all legends."""
+        cls._default_handler_map = handler_map
+
+    @classmethod
+    def update_default_handler_map(cls, handler_map):
+        """Update the global default handler map, shared by all legends."""
+        cls._default_handler_map.update(handler_map)
+
+    def get_legend_handler_map(self):
+        """Return this legend instance's handler map."""
+        default_handler_map = self.get_default_handler_map()
+        return ({**default_handler_map, **self._custom_handler_map}
+                if self._custom_handler_map else default_handler_map)
+
+    @staticmethod
+    def get_legend_handler(legend_handler_map, orig_handle):
+        """
+        Return a legend handler from *legend_handler_map* that
+        corresponds to *orig_handler*.
+
+        *legend_handler_map* should be a dictionary object (that is
+        returned by the get_legend_handler_map method).
+
+        It first checks if the *orig_handle* itself is a key in the
+        *legend_handler_map* and return the associated value.
+        Otherwise, it checks for each of the classes in its
+        method-resolution-order. If no matching key is found, it
+        returns ``None``.
+        """
+        try:
+            return legend_handler_map[orig_handle]
+        except (TypeError, KeyError):  # TypeError if unhashable.
+            pass
+        for handle_type in type(orig_handle).mro():
+            try:
+                return legend_handler_map[handle_type]
+            except KeyError:
+                pass
+        return None
+
+    def _init_legend_box(self, handles, labels, markerfirst=True):
+        """
+        Initialize the legend_box. The legend_box is an instance of
+        the OffsetBox, which is packed with legend handles and
+        texts. Once packed, their location is calculated during the
+        drawing time.
+        """
+
+        fontsize = self._fontsize
+
+        # legend_box is a HPacker, horizontally packed with columns.
+        # Each column is a VPacker, vertically packed with legend items.
+        # Each legend item is a HPacker packed with:
+        # - handlebox: a DrawingArea which contains the legend handle.
+        # - labelbox: a TextArea which contains the legend text.
+
+        text_list = []  # the list of text instances
+        handle_list = []  # the list of handle instances
+        handles_and_labels = []
+
+        # The approximate height and descent of text. These values are
+        # only used for plotting the legend handle.
+        descent = 0.35 * fontsize * (self.handleheight - 0.7)  # heuristic.
+        height = fontsize * self.handleheight - descent
+        # each handle needs to be drawn inside a box of (x, y, w, h) =
+        # (0, -descent, width, height).  And their coordinates should
+        # be given in the display coordinates.
+
+        # The transformation of each handle will be automatically set
+        # to self.get_transform(). If the artist does not use its
+        # default transform (e.g., Collections), you need to
+        # manually set their transform to the self.get_transform().
+        legend_handler_map = self.get_legend_handler_map()
+
+        for orig_handle, label in zip(handles, labels):
+            handler = self.get_legend_handler(legend_handler_map, orig_handle)
+            if handler is None:
+                _api.warn_external(
+                             "Legend does not support handles for "
+                             f"{type(orig_handle).__name__} "
+                             "instances.\nA proxy artist may be used "
+                             "instead.\nSee: https://matplotlib.org/"
+                             "stable/users/explain/axes/legend_guide.html"
+                             "#controlling-the-legend-entries")
+                # No handle for this artist, so we just defer to None.
+                handle_list.append(None)
+            else:
+                textbox = TextArea(label, multilinebaseline=True,
+                                   textprops=dict(
+                                       verticalalignment='baseline',
+                                       horizontalalignment='left',
+                                       fontproperties=self.prop))
+                handlebox = DrawingArea(width=self.handlelength * fontsize,
+                                        height=height,
+                                        xdescent=0., ydescent=descent)
+
+                text_list.append(textbox._text)
+                # Create the artist for the legend which represents the
+                # original artist/handle.
+                handle_list.append(handler.legend_artist(self, orig_handle,
+                                                         fontsize, handlebox))
+                handles_and_labels.append((handlebox, textbox))
+
+        columnbox = []
+        # array_split splits n handles_and_labels into ncols columns, with the
+        # first n%ncols columns having an extra entry.  filter(len, ...)
+        # handles the case where n < ncols: the last ncols-n columns are empty
+        # and get filtered out.
+        for handles_and_labels_column in filter(
+                len, np.array_split(handles_and_labels, self._ncols)):
+            # pack handlebox and labelbox into itembox
+            itemboxes = [HPacker(pad=0,
+                                 sep=self.handletextpad * fontsize,
+                                 children=[h, t] if markerfirst else [t, h],
+                                 align="baseline")
+                         for h, t in handles_and_labels_column]
+            # pack columnbox
+            alignment = "baseline" if markerfirst else "right"
+            columnbox.append(VPacker(pad=0,
+                                     sep=self.labelspacing * fontsize,
+                                     align=alignment,
+                                     children=itemboxes))
+
+        mode = "expand" if self._mode == "expand" else "fixed"
+        sep = self.columnspacing * fontsize
+        self._legend_handle_box = HPacker(pad=0,
+                                          sep=sep, align="baseline",
+                                          mode=mode,
+                                          children=columnbox)
+        self._legend_title_box = TextArea("")
+        self._legend_box = VPacker(pad=self.borderpad * fontsize,
+                                   sep=self.labelspacing * fontsize,
+                                   align=self._alignment,
+                                   children=[self._legend_title_box,
+                                             self._legend_handle_box])
+        self._legend_box.set_figure(self.get_figure(root=False))
+        self._legend_box.axes = self.axes
+        self.texts = text_list
+        self.legend_handles = handle_list
+
+    def _auto_legend_data(self, renderer):
+        """
+        Return display coordinates for hit testing for "best" positioning.
+
+        Returns
+        -------
+        bboxes
+            List of bounding boxes of all patches.
+        lines
+            List of `.Path` corresponding to each line.
+        offsets
+            List of (x, y) offsets of all collection.
+        """
+        assert self.isaxes  # always holds, as this is only called internally
+        bboxes = []
+        lines = []
+        offsets = []
+        for artist in self.parent._children:
+            if isinstance(artist, Line2D):
+                lines.append(
+                    artist.get_transform().transform_path(artist.get_path()))
+            elif isinstance(artist, Rectangle):
+                bboxes.append(
+                    artist.get_bbox().transformed(artist.get_data_transform()))
+            elif isinstance(artist, Patch):
+                lines.append(
+                    artist.get_transform().transform_path(artist.get_path()))
+            elif isinstance(artist, PolyCollection):
+                lines.extend(artist.get_transform().transform_path(path)
+                             for path in artist.get_paths())
+            elif isinstance(artist, Collection):
+                transform, transOffset, hoffsets, _ = artist._prepare_points()
+                if len(hoffsets):
+                    offsets.extend(transOffset.transform(hoffsets))
+            elif isinstance(artist, Text):
+                bboxes.append(artist.get_window_extent(renderer))
+
+        return bboxes, lines, offsets
+
+    def get_children(self):
+        # docstring inherited
+        return [self._legend_box, self.get_frame()]
+
+    def get_frame(self):
+        """Return the `~.patches.Rectangle` used to frame the legend."""
+        return self.legendPatch
+
+    def get_lines(self):
+        r"""Return the list of `~.lines.Line2D`\s in the legend."""
+        return [h for h in self.legend_handles if isinstance(h, Line2D)]
+
+    def get_patches(self):
+        r"""Return the list of `~.patches.Patch`\s in the legend."""
+        return silent_list('Patch',
+                           [h for h in self.legend_handles
+                            if isinstance(h, Patch)])
+
+    def get_texts(self):
+        r"""Return the list of `~.text.Text`\s in the legend."""
+        return silent_list('Text', self.texts)
+
+    def set_alignment(self, alignment):
+        """
+        Set the alignment of the legend title and the box of entries.
+
+        The entries are aligned as a single block, so that markers always
+        lined up.
+
+        Parameters
+        ----------
+        alignment : {'center', 'left', 'right'}.
+
+        """
+        _api.check_in_list(["center", "left", "right"], alignment=alignment)
+        self._alignment = alignment
+        self._legend_box.align = alignment
+
+    def get_alignment(self):
+        """Get the alignment value of the legend box"""
+        return self._legend_box.align
+
+    def set_title(self, title, prop=None):
+        """
+        Set legend title and title style.
+
+        Parameters
+        ----------
+        title : str
+            The legend title.
+
+        prop : `.font_manager.FontProperties` or `str` or `pathlib.Path`
+            The font properties of the legend title.
+            If a `str`, it is interpreted as a fontconfig pattern parsed by
+            `.FontProperties`.  If a `pathlib.Path`, it is interpreted as the
+            absolute path to a font file.
+
+        """
+        self._legend_title_box._text.set_text(title)
+        if title:
+            self._legend_title_box._text.set_visible(True)
+            self._legend_title_box.set_visible(True)
+        else:
+            self._legend_title_box._text.set_visible(False)
+            self._legend_title_box.set_visible(False)
+
+        if prop is not None:
+            self._legend_title_box._text.set_fontproperties(prop)
+
+        self.stale = True
+
+    def get_title(self):
+        """Return the `.Text` instance for the legend title."""
+        return self._legend_title_box._text
+
+    def get_window_extent(self, renderer=None):
+        # docstring inherited
+        if renderer is None:
+            renderer = self.get_figure(root=True)._get_renderer()
+        return self._legend_box.get_window_extent(renderer=renderer)
+
+    def get_tightbbox(self, renderer=None):
+        # docstring inherited
+        return self._legend_box.get_window_extent(renderer)
+
+    def get_frame_on(self):
+        """Get whether the legend box patch is drawn."""
+        return self.legendPatch.get_visible()
+
+    def set_frame_on(self, b):
+        """
+        Set whether the legend box patch is drawn.
+
+        Parameters
+        ----------
+        b : bool
+        """
+        self.legendPatch.set_visible(b)
+        self.stale = True
+
+    draw_frame = set_frame_on  # Backcompat alias.
+
+    def get_bbox_to_anchor(self):
+        """Return the bbox that the legend will be anchored to."""
+        if self._bbox_to_anchor is None:
+            return self.parent.bbox
+        else:
+            return self._bbox_to_anchor
+
+    def set_bbox_to_anchor(self, bbox, transform=None):
+        """
+        Set the bbox that the legend will be anchored to.
+
+        Parameters
+        ----------
+        bbox : `~matplotlib.transforms.BboxBase` or tuple
+            The bounding box can be specified in the following ways:
+
+            - A `.BboxBase` instance
+            - A tuple of ``(left, bottom, width, height)`` in the given
+              transform (normalized axes coordinate if None)
+            - A tuple of ``(left, bottom)`` where the width and height will be
+              assumed to be zero.
+            - *None*, to remove the bbox anchoring, and use the parent bbox.
+
+        transform : `~matplotlib.transforms.Transform`, optional
+            A transform to apply to the bounding box. If not specified, this
+            will use a transform to the bounding box of the parent.
+        """
+        if bbox is None:
+            self._bbox_to_anchor = None
+            return
+        elif isinstance(bbox, BboxBase):
+            self._bbox_to_anchor = bbox
+        else:
+            try:
+                l = len(bbox)
+            except TypeError as err:
+                raise ValueError(f"Invalid bbox: {bbox}") from err
+
+            if l == 2:
+                bbox = [bbox[0], bbox[1], 0, 0]
+
+            self._bbox_to_anchor = Bbox.from_bounds(*bbox)
+
+        if transform is None:
+            transform = BboxTransformTo(self.parent.bbox)
+
+        self._bbox_to_anchor = TransformedBbox(self._bbox_to_anchor,
+                                               transform)
+        self.stale = True
+
+    def _get_anchored_bbox(self, loc, bbox, parentbbox, renderer):
+        """
+        Place the *bbox* inside the *parentbbox* according to a given
+        location code. Return the (x, y) coordinate of the bbox.
+
+        Parameters
+        ----------
+        loc : int
+            A location code in range(1, 11). This corresponds to the possible
+            values for ``self._loc``, excluding "best".
+        bbox : `~matplotlib.transforms.Bbox`
+            bbox to be placed, in display coordinates.
+        parentbbox : `~matplotlib.transforms.Bbox`
+            A parent box which will contain the bbox, in display coordinates.
+        """
+        return offsetbox._get_anchored_bbox(
+            loc, bbox, parentbbox,
+            self.borderaxespad * renderer.points_to_pixels(self._fontsize))
+
+    def _find_best_position(self, width, height, renderer):
+        """Determine the best location to place the legend."""
+        assert self.isaxes  # always holds, as this is only called internally
+
+        start_time = time.perf_counter()
+
+        bboxes, lines, offsets = self._auto_legend_data(renderer)
+
+        bbox = Bbox.from_bounds(0, 0, width, height)
+
+        candidates = []
+        for idx in range(1, len(self.codes)):
+            l, b = self._get_anchored_bbox(idx, bbox,
+                                           self.get_bbox_to_anchor(),
+                                           renderer)
+            legendBox = Bbox.from_bounds(l, b, width, height)
+            # XXX TODO: If markers are present, it would be good to take them
+            # into account when checking vertex overlaps in the next line.
+            badness = (sum(legendBox.count_contains(line.vertices)
+                           for line in lines)
+                       + legendBox.count_contains(offsets)
+                       + legendBox.count_overlaps(bboxes)
+                       + sum(line.intersects_bbox(legendBox, filled=False)
+                             for line in lines))
+            # Include the index to favor lower codes in case of a tie.
+            candidates.append((badness, idx, (l, b)))
+            if badness == 0:
+                break
+
+        _, _, (l, b) = min(candidates)
+
+        if self._loc_used_default and time.perf_counter() - start_time > 1:
+            _api.warn_external(
+                'Creating legend with loc="best" can be slow with large '
+                'amounts of data.')
+
+        return l, b
+
+    def contains(self, mouseevent):
+        return self.legendPatch.contains(mouseevent)
+
+    def set_draggable(self, state, use_blit=False, update='loc'):
+        """
+        Enable or disable mouse dragging support of the legend.
+
+        Parameters
+        ----------
+        state : bool
+            Whether mouse dragging is enabled.
+        use_blit : bool, optional
+            Use blitting for faster image composition. For details see
+            :ref:`func-animation`.
+        update : {'loc', 'bbox'}, optional
+            The legend parameter to be changed when dragged:
+
+            - 'loc': update the *loc* parameter of the legend
+            - 'bbox': update the *bbox_to_anchor* parameter of the legend
+
+        Returns
+        -------
+        `.DraggableLegend` or *None*
+            If *state* is ``True`` this returns the `.DraggableLegend` helper
+            instance. Otherwise this returns *None*.
+        """
+        if state:
+            if self._draggable is None:
+                self._draggable = DraggableLegend(self,
+                                                  use_blit,
+                                                  update=update)
+        else:
+            if self._draggable is not None:
+                self._draggable.disconnect()
+            self._draggable = None
+        return self._draggable
+
+    def get_draggable(self):
+        """Return ``True`` if the legend is draggable, ``False`` otherwise."""
+        return self._draggable is not None
+
+
+# Helper functions to parse legend arguments for both `figure.legend` and
+# `axes.legend`:
+def _get_legend_handles(axs, legend_handler_map=None):
+    """Yield artists that can be used as handles in a legend."""
+    handles_original = []
+    for ax in axs:
+        handles_original += [
+            *(a for a in ax._children
+              if isinstance(a, (Line2D, Patch, Collection, Text))),
+            *ax.containers]
+        # support parasite Axes:
+        if hasattr(ax, 'parasites'):
+            for axx in ax.parasites:
+                handles_original += [
+                    *(a for a in axx._children
+                      if isinstance(a, (Line2D, Patch, Collection, Text))),
+                    *axx.containers]
+
+    handler_map = {**Legend.get_default_handler_map(),
+                   **(legend_handler_map or {})}
+    has_handler = Legend.get_legend_handler
+    for handle in handles_original:
+        label = handle.get_label()
+        if label != '_nolegend_' and has_handler(handler_map, handle):
+            yield handle
+        elif (label and not label.startswith('_') and
+                not has_handler(handler_map, handle)):
+            _api.warn_external(
+                             "Legend does not support handles for "
+                             f"{type(handle).__name__} "
+                             "instances.\nSee: https://matplotlib.org/stable/"
+                             "tutorials/intermediate/legend_guide.html"
+                             "#implementing-a-custom-legend-handler")
+            continue
+
+
+def _get_legend_handles_labels(axs, legend_handler_map=None):
+    """Return handles and labels for legend."""
+    handles = []
+    labels = []
+    for handle in _get_legend_handles(axs, legend_handler_map):
+        label = handle.get_label()
+        if label and not label.startswith('_'):
+            handles.append(handle)
+            labels.append(label)
+    return handles, labels
+
+
+def _parse_legend_args(axs, *args, handles=None, labels=None, **kwargs):
+    """
+    Get the handles and labels from the calls to either ``figure.legend``
+    or ``axes.legend``.
+
+    The parser is a bit involved because we support::
+
+        legend()
+        legend(labels)
+        legend(handles, labels)
+        legend(labels=labels)
+        legend(handles=handles)
+        legend(handles=handles, labels=labels)
+
+    The behavior for a mixture of positional and keyword handles and labels
+    is undefined and issues a warning; it will be an error in the future.
+
+    Parameters
+    ----------
+    axs : list of `.Axes`
+        If handles are not given explicitly, the artists in these Axes are
+        used as handles.
+    *args : tuple
+        Positional parameters passed to ``legend()``.
+    handles
+        The value of the keyword argument ``legend(handles=...)``, or *None*
+        if that keyword argument was not used.
+    labels
+        The value of the keyword argument ``legend(labels=...)``, or *None*
+        if that keyword argument was not used.
+    **kwargs
+        All other keyword arguments passed to ``legend()``.
+
+    Returns
+    -------
+    handles : list of (`.Artist` or tuple of `.Artist`)
+        The legend handles.
+    labels : list of str
+        The legend labels.
+    kwargs : dict
+        *kwargs* with keywords handles and labels removed.
+
+    """
+    log = logging.getLogger(__name__)
+
+    handlers = kwargs.get('handler_map')
+
+    if (handles is not None or labels is not None) and args:
+        _api.warn_deprecated("3.9", message=(
+            "You have mixed positional and keyword arguments, some input may "
+            "be discarded.  This is deprecated since %(since)s and will "
+            "become an error in %(removal)s."))
+
+    if (hasattr(handles, "__len__") and
+            hasattr(labels, "__len__") and
+            len(handles) != len(labels)):
+        _api.warn_external(f"Mismatched number of handles and labels: "
+                           f"len(handles) = {len(handles)} "
+                           f"len(labels) = {len(labels)}")
+    # if got both handles and labels as kwargs, make same length
+    if handles and labels:
+        handles, labels = zip(*zip(handles, labels))
+
+    elif handles is not None and labels is None:
+        labels = [handle.get_label() for handle in handles]
+
+    elif labels is not None and handles is None:
+        # Get as many handles as there are labels.
+        handles = [handle for handle, label
+                   in zip(_get_legend_handles(axs, handlers), labels)]
+
+    elif len(args) == 0:  # 0 args: automatically detect labels and handles.
+        handles, labels = _get_legend_handles_labels(axs, handlers)
+        if not handles:
+            _api.warn_external(
+                "No artists with labels found to put in legend.  Note that "
+                "artists whose label start with an underscore are ignored "
+                "when legend() is called with no argument.")
+
+    elif len(args) == 1:  # 1 arg: user defined labels, automatic handle detection.
+        labels, = args
+        if any(isinstance(l, Artist) for l in labels):
+            raise TypeError("A single argument passed to legend() must be a "
+                            "list of labels, but found an Artist in there.")
+
+        # Get as many handles as there are labels.
+        handles = [handle for handle, label
+                   in zip(_get_legend_handles(axs, handlers), labels)]
+
+    elif len(args) == 2:  # 2 args: user defined handles and labels.
+        handles, labels = args[:2]
+
+    else:
+        raise _api.nargs_error('legend', '0-2', len(args))
+
+    return handles, labels, kwargs

lib/python3.12/site-packages/matplotlib/lines.py

@@ -0,0 +1,1720 @@
+"""
+2D lines with support for a variety of line styles, markers, colors, etc.
+"""
+
+import copy
+
+from numbers import Integral, Number, Real
+import logging
+
+import numpy as np
+
+import matplotlib as mpl
+from . import _api, cbook, colors as mcolors, _docstring
+from .artist import Artist, allow_rasterization
+from .cbook import (
+    _to_unmasked_float_array, ls_mapper, ls_mapper_r, STEP_LOOKUP_MAP)
+from .markers import MarkerStyle
+from .path import Path
+from .transforms import Bbox, BboxTransformTo, TransformedPath
+from ._enums import JoinStyle, CapStyle
+
+# Imported here for backward compatibility, even though they don't
+# really belong.
+from . import _path
+from .markers import (  # noqa
+    CARETLEFT, CARETRIGHT, CARETUP, CARETDOWN,
+    CARETLEFTBASE, CARETRIGHTBASE, CARETUPBASE, CARETDOWNBASE,
+    TICKLEFT, TICKRIGHT, TICKUP, TICKDOWN)
+
+_log = logging.getLogger(__name__)
+
+
+def _get_dash_pattern(style):
+    """Convert linestyle to dash pattern."""
+    # go from short hand -> full strings
+    if isinstance(style, str):
+        style = ls_mapper.get(style, style)
+    # un-dashed styles
+    if style in ['solid', 'None']:
+        offset = 0
+        dashes = None
+    # dashed styles
+    elif style in ['dashed', 'dashdot', 'dotted']:
+        offset = 0
+        dashes = tuple(mpl.rcParams[f'lines.{style}_pattern'])
+    #
+    elif isinstance(style, tuple):
+        offset, dashes = style
+        if offset is None:
+            raise ValueError(f'Unrecognized linestyle: {style!r}')
+    else:
+        raise ValueError(f'Unrecognized linestyle: {style!r}')
+
+    # normalize offset to be positive and shorter than the dash cycle
+    if dashes is not None:
+        dsum = sum(dashes)
+        if dsum:
+            offset %= dsum
+
+    return offset, dashes
+
+
+def _get_dash_patterns(styles):
+    """Convert linestyle or sequence of linestyles to list of dash patterns."""
+    try:
+        patterns = [_get_dash_pattern(styles)]
+    except ValueError:
+        try:
+            patterns = [_get_dash_pattern(x) for x in styles]
+        except ValueError as err:
+            emsg = f'Do not know how to convert {styles!r} to dashes'
+            raise ValueError(emsg) from err
+
+    return patterns
+
+
+def _get_inverse_dash_pattern(offset, dashes):
+    """Return the inverse of the given dash pattern, for filling the gaps."""
+    # Define the inverse pattern by moving the last gap to the start of the
+    # sequence.
+    gaps = dashes[-1:] + dashes[:-1]
+    # Set the offset so that this new first segment is skipped
+    # (see backend_bases.GraphicsContextBase.set_dashes for offset definition).
+    offset_gaps = offset + dashes[-1]
+
+    return offset_gaps, gaps
+
+
+def _scale_dashes(offset, dashes, lw):
+    if not mpl.rcParams['lines.scale_dashes']:
+        return offset, dashes
+    scaled_offset = offset * lw
+    scaled_dashes = ([x * lw if x is not None else None for x in dashes]
+                     if dashes is not None else None)
+    return scaled_offset, scaled_dashes
+
+
+def segment_hits(cx, cy, x, y, radius):
+    """
+    Return the indices of the segments in the polyline with coordinates (*cx*,
+    *cy*) that are within a distance *radius* of the point (*x*, *y*).
+    """
+    # Process single points specially
+    if len(x) <= 1:
+        res, = np.nonzero((cx - x) ** 2 + (cy - y) ** 2 <= radius ** 2)
+        return res
+
+    # We need to lop the last element off a lot.
+    xr, yr = x[:-1], y[:-1]
+
+    # Only look at line segments whose nearest point to C on the line
+    # lies within the segment.
+    dx, dy = x[1:] - xr, y[1:] - yr
+    Lnorm_sq = dx ** 2 + dy ** 2  # Possibly want to eliminate Lnorm==0
+    u = ((cx - xr) * dx + (cy - yr) * dy) / Lnorm_sq
+    candidates = (u >= 0) & (u <= 1)
+
+    # Note that there is a little area near one side of each point
+    # which will be near neither segment, and another which will
+    # be near both, depending on the angle of the lines.  The
+    # following radius test eliminates these ambiguities.
+    point_hits = (cx - x) ** 2 + (cy - y) ** 2 <= radius ** 2
+    candidates = candidates & ~(point_hits[:-1] | point_hits[1:])
+
+    # For those candidates which remain, determine how far they lie away
+    # from the line.
+    px, py = xr + u * dx, yr + u * dy
+    line_hits = (cx - px) ** 2 + (cy - py) ** 2 <= radius ** 2
+    line_hits = line_hits & candidates
+    points, = point_hits.ravel().nonzero()
+    lines, = line_hits.ravel().nonzero()
+    return np.concatenate((points, lines))
+
+
+def _mark_every_path(markevery, tpath, affine, ax):
+    """
+    Helper function that sorts out how to deal the input
+    `markevery` and returns the points where markers should be drawn.
+
+    Takes in the `markevery` value and the line path and returns the
+    sub-sampled path.
+    """
+    # pull out the two bits of data we want from the path
+    codes, verts = tpath.codes, tpath.vertices
+
+    def _slice_or_none(in_v, slc):
+        """Helper function to cope with `codes` being an ndarray or `None`."""
+        if in_v is None:
+            return None
+        return in_v[slc]
+
+    # if just an int, assume starting at 0 and make a tuple
+    if isinstance(markevery, Integral):
+        markevery = (0, markevery)
+    # if just a float, assume starting at 0.0 and make a tuple
+    elif isinstance(markevery, Real):
+        markevery = (0.0, markevery)
+
+    if isinstance(markevery, tuple):
+        if len(markevery) != 2:
+            raise ValueError('`markevery` is a tuple but its len is not 2; '
+                             f'markevery={markevery}')
+        start, step = markevery
+        # if step is an int, old behavior
+        if isinstance(step, Integral):
+            # tuple of 2 int is for backwards compatibility,
+            if not isinstance(start, Integral):
+                raise ValueError(
+                    '`markevery` is a tuple with len 2 and second element is '
+                    'an int, but the first element is not an int; '
+                    f'markevery={markevery}')
+            # just return, we are done here
+
+            return Path(verts[slice(start, None, step)],
+                        _slice_or_none(codes, slice(start, None, step)))
+
+        elif isinstance(step, Real):
+            if not isinstance(start, Real):
+                raise ValueError(
+                    '`markevery` is a tuple with len 2 and second element is '
+                    'a float, but the first element is not a float or an int; '
+                    f'markevery={markevery}')
+            if ax is None:
+                raise ValueError(
+                    "markevery is specified relative to the Axes size, but "
+                    "the line does not have a Axes as parent")
+
+            # calc cumulative distance along path (in display coords):
+            fin = np.isfinite(verts).all(axis=1)
+            fverts = verts[fin]
+            disp_coords = affine.transform(fverts)
+
+            delta = np.empty((len(disp_coords), 2))
+            delta[0, :] = 0
+            delta[1:, :] = disp_coords[1:, :] - disp_coords[:-1, :]
+            delta = np.hypot(*delta.T).cumsum()
+            # calc distance between markers along path based on the Axes
+            # bounding box diagonal being a distance of unity:
+            (x0, y0), (x1, y1) = ax.transAxes.transform([[0, 0], [1, 1]])
+            scale = np.hypot(x1 - x0, y1 - y0)
+            marker_delta = np.arange(start * scale, delta[-1], step * scale)
+            # find closest actual data point that is closest to
+            # the theoretical distance along the path:
+            inds = np.abs(delta[np.newaxis, :] - marker_delta[:, np.newaxis])
+            inds = inds.argmin(axis=1)
+            inds = np.unique(inds)
+            # return, we are done here
+            return Path(fverts[inds], _slice_or_none(codes, inds))
+        else:
+            raise ValueError(
+                f"markevery={markevery!r} is a tuple with len 2, but its "
+                f"second element is not an int or a float")
+
+    elif isinstance(markevery, slice):
+        # mazol tov, it's already a slice, just return
+        return Path(verts[markevery], _slice_or_none(codes, markevery))
+
+    elif np.iterable(markevery):
+        # fancy indexing
+        try:
+            return Path(verts[markevery], _slice_or_none(codes, markevery))
+        except (ValueError, IndexError) as err:
+            raise ValueError(
+                f"markevery={markevery!r} is iterable but not a valid numpy "
+                f"fancy index") from err
+    else:
+        raise ValueError(f"markevery={markevery!r} is not a recognized value")
+
+
+@_docstring.interpd
+@_api.define_aliases({
+    "antialiased": ["aa"],
+    "color": ["c"],
+    "drawstyle": ["ds"],
+    "linestyle": ["ls"],
+    "linewidth": ["lw"],
+    "markeredgecolor": ["mec"],
+    "markeredgewidth": ["mew"],
+    "markerfacecolor": ["mfc"],
+    "markerfacecoloralt": ["mfcalt"],
+    "markersize": ["ms"],
+})
+class Line2D(Artist):
+    """
+    A line - the line can have both a solid linestyle connecting all
+    the vertices, and a marker at each vertex.  Additionally, the
+    drawing of the solid line is influenced by the drawstyle, e.g., one
+    can create "stepped" lines in various styles.
+    """
+
+    lineStyles = _lineStyles = {  # hidden names deprecated
+        '-':    '_draw_solid',
+        '--':   '_draw_dashed',
+        '-.':   '_draw_dash_dot',
+        ':':    '_draw_dotted',
+        'None': '_draw_nothing',
+        ' ':    '_draw_nothing',
+        '':     '_draw_nothing',
+    }
+
+    _drawStyles_l = {
+        'default':    '_draw_lines',
+        'steps-mid':  '_draw_steps_mid',
+        'steps-pre':  '_draw_steps_pre',
+        'steps-post': '_draw_steps_post',
+    }
+
+    _drawStyles_s = {
+        'steps': '_draw_steps_pre',
+    }
+
+    # drawStyles should now be deprecated.
+    drawStyles = {**_drawStyles_l, **_drawStyles_s}
+    # Need a list ordered with long names first:
+    drawStyleKeys = [*_drawStyles_l, *_drawStyles_s]
+
+    # Referenced here to maintain API.  These are defined in
+    # MarkerStyle
+    markers = MarkerStyle.markers
+    filled_markers = MarkerStyle.filled_markers
+    fillStyles = MarkerStyle.fillstyles
+
+    zorder = 2
+
+    _subslice_optim_min_size = 1000
+
+    def __str__(self):
+        if self._label != "":
+            return f"Line2D({self._label})"
+        elif self._x is None:
+            return "Line2D()"
+        elif len(self._x) > 3:
+            return "Line2D(({:g},{:g}),({:g},{:g}),...,({:g},{:g}))".format(
+                self._x[0], self._y[0],
+                self._x[1], self._y[1],
+                self._x[-1], self._y[-1])
+        else:
+            return "Line2D(%s)" % ",".join(
+                map("({:g},{:g})".format, self._x, self._y))
+
+    def __init__(self, xdata, ydata, *,
+                 linewidth=None,  # all Nones default to rc
+                 linestyle=None,
+                 color=None,
+                 gapcolor=None,
+                 marker=None,
+                 markersize=None,
+                 markeredgewidth=None,
+                 markeredgecolor=None,
+                 markerfacecolor=None,
+                 markerfacecoloralt='none',
+                 fillstyle=None,
+                 antialiased=None,
+                 dash_capstyle=None,
+                 solid_capstyle=None,
+                 dash_joinstyle=None,
+                 solid_joinstyle=None,
+                 pickradius=5,
+                 drawstyle=None,
+                 markevery=None,
+                 **kwargs
+                 ):
+        """
+        Create a `.Line2D` instance with *x* and *y* data in sequences of
+        *xdata*, *ydata*.
+
+        Additional keyword arguments are `.Line2D` properties:
+
+        %(Line2D:kwdoc)s
+
+        See :meth:`set_linestyle` for a description of the line styles,
+        :meth:`set_marker` for a description of the markers, and
+        :meth:`set_drawstyle` for a description of the draw styles.
+
+        """
+        super().__init__()
+
+        # Convert sequences to NumPy arrays.
+        if not np.iterable(xdata):
+            raise RuntimeError('xdata must be a sequence')
+        if not np.iterable(ydata):
+            raise RuntimeError('ydata must be a sequence')
+
+        if linewidth is None:
+            linewidth = mpl.rcParams['lines.linewidth']
+
+        if linestyle is None:
+            linestyle = mpl.rcParams['lines.linestyle']
+        if marker is None:
+            marker = mpl.rcParams['lines.marker']
+        if color is None:
+            color = mpl.rcParams['lines.color']
+
+        if markersize is None:
+            markersize = mpl.rcParams['lines.markersize']
+        if antialiased is None:
+            antialiased = mpl.rcParams['lines.antialiased']
+        if dash_capstyle is None:
+            dash_capstyle = mpl.rcParams['lines.dash_capstyle']
+        if dash_joinstyle is None:
+            dash_joinstyle = mpl.rcParams['lines.dash_joinstyle']
+        if solid_capstyle is None:
+            solid_capstyle = mpl.rcParams['lines.solid_capstyle']
+        if solid_joinstyle is None:
+            solid_joinstyle = mpl.rcParams['lines.solid_joinstyle']
+
+        if drawstyle is None:
+            drawstyle = 'default'
+
+        self._dashcapstyle = None
+        self._dashjoinstyle = None
+        self._solidjoinstyle = None
+        self._solidcapstyle = None
+        self.set_dash_capstyle(dash_capstyle)
+        self.set_dash_joinstyle(dash_joinstyle)
+        self.set_solid_capstyle(solid_capstyle)
+        self.set_solid_joinstyle(solid_joinstyle)
+
+        self._linestyles = None
+        self._drawstyle = None
+        self._linewidth = linewidth
+        self._unscaled_dash_pattern = (0, None)  # offset, dash
+        self._dash_pattern = (0, None)  # offset, dash (scaled by linewidth)
+
+        self.set_linewidth(linewidth)
+        self.set_linestyle(linestyle)
+        self.set_drawstyle(drawstyle)
+
+        self._color = None
+        self.set_color(color)
+        if marker is None:
+            marker = 'none'  # Default.
+        if not isinstance(marker, MarkerStyle):
+            self._marker = MarkerStyle(marker, fillstyle)
+        else:
+            self._marker = marker
+
+        self._gapcolor = None
+        self.set_gapcolor(gapcolor)
+
+        self._markevery = None
+        self._markersize = None
+        self._antialiased = None
+
+        self.set_markevery(markevery)
+        self.set_antialiased(antialiased)
+        self.set_markersize(markersize)
+
+        self._markeredgecolor = None
+        self._markeredgewidth = None
+        self._markerfacecolor = None
+        self._markerfacecoloralt = None
+
+        self.set_markerfacecolor(markerfacecolor)  # Normalizes None to rc.
+        self.set_markerfacecoloralt(markerfacecoloralt)
+        self.set_markeredgecolor(markeredgecolor)  # Normalizes None to rc.
+        self.set_markeredgewidth(markeredgewidth)
+
+        # update kwargs before updating data to give the caller a
+        # chance to init axes (and hence unit support)
+        self._internal_update(kwargs)
+        self.pickradius = pickradius
+        self.ind_offset = 0
+        if (isinstance(self._picker, Number) and
+                not isinstance(self._picker, bool)):
+            self._pickradius = self._picker
+
+        self._xorig = np.asarray([])
+        self._yorig = np.asarray([])
+        self._invalidx = True
+        self._invalidy = True
+        self._x = None
+        self._y = None
+        self._xy = None
+        self._path = None
+        self._transformed_path = None
+        self._subslice = False
+        self._x_filled = None  # used in subslicing; only x is needed
+
+        self.set_data(xdata, ydata)
+
+    def contains(self, mouseevent):
+        """
+        Test whether *mouseevent* occurred on the line.
+
+        An event is deemed to have occurred "on" the line if it is less
+        than ``self.pickradius`` (default: 5 points) away from it.  Use
+        `~.Line2D.get_pickradius` or `~.Line2D.set_pickradius` to get or set
+        the pick radius.
+
+        Parameters
+        ----------
+        mouseevent : `~matplotlib.backend_bases.MouseEvent`
+
+        Returns
+        -------
+        contains : bool
+            Whether any values are within the radius.
+        details : dict
+            A dictionary ``{'ind': pointlist}``, where *pointlist* is a
+            list of points of the line that are within the pickradius around
+            the event position.
+
+            TODO: sort returned indices by distance
+        """
+        if self._different_canvas(mouseevent):
+            return False, {}
+
+        # Make sure we have data to plot
+        if self._invalidy or self._invalidx:
+            self.recache()
+        if len(self._xy) == 0:
+            return False, {}
+
+        # Convert points to pixels
+        transformed_path = self._get_transformed_path()
+        path, affine = transformed_path.get_transformed_path_and_affine()
+        path = affine.transform_path(path)
+        xy = path.vertices
+        xt = xy[:, 0]
+        yt = xy[:, 1]
+
+        # Convert pick radius from points to pixels
+        fig = self.get_figure(root=True)
+        if fig is None:
+            _log.warning('no figure set when check if mouse is on line')
+            pixels = self._pickradius
+        else:
+            pixels = fig.dpi / 72. * self._pickradius
+
+        # The math involved in checking for containment (here and inside of
+        # segment_hits) assumes that it is OK to overflow, so temporarily set
+        # the error flags accordingly.
+        with np.errstate(all='ignore'):
+            # Check for collision
+            if self._linestyle in ['None', None]:
+                # If no line, return the nearby point(s)
+                ind, = np.nonzero(
+                    (xt - mouseevent.x) ** 2 + (yt - mouseevent.y) ** 2
+                    <= pixels ** 2)
+            else:
+                # If line, return the nearby segment(s)
+                ind = segment_hits(mouseevent.x, mouseevent.y, xt, yt, pixels)
+                if self._drawstyle.startswith("steps"):
+                    ind //= 2
+
+        ind += self.ind_offset
+
+        # Return the point(s) within radius
+        return len(ind) > 0, dict(ind=ind)
+
+    def get_pickradius(self):
+        """
+        Return the pick radius used for containment tests.
+
+        See `.contains` for more details.
+        """
+        return self._pickradius
+
+    def set_pickradius(self, pickradius):
+        """
+        Set the pick radius used for containment tests.
+
+        See `.contains` for more details.
+
+        Parameters
+        ----------
+        pickradius : float
+            Pick radius, in points.
+        """
+        if not isinstance(pickradius, Real) or pickradius < 0:
+            raise ValueError("pick radius should be a distance")
+        self._pickradius = pickradius
+
+    pickradius = property(get_pickradius, set_pickradius)
+
+    def get_fillstyle(self):
+        """
+        Return the marker fill style.
+
+        See also `~.Line2D.set_fillstyle`.
+        """
+        return self._marker.get_fillstyle()
+
+    def set_fillstyle(self, fs):
+        """
+        Set the marker fill style.
+
+        Parameters
+        ----------
+        fs : {'full', 'left', 'right', 'bottom', 'top', 'none'}
+            Possible values:
+
+            - 'full': Fill the whole marker with the *markerfacecolor*.
+            - 'left', 'right', 'bottom', 'top': Fill the marker half at
+              the given side with the *markerfacecolor*. The other
+              half of the marker is filled with *markerfacecoloralt*.
+            - 'none': No filling.
+
+            For examples see :ref:`marker_fill_styles`.
+        """
+        self.set_marker(MarkerStyle(self._marker.get_marker(), fs))
+        self.stale = True
+
+    def set_markevery(self, every):
+        """
+        Set the markevery property to subsample the plot when using markers.
+
+        e.g., if ``every=5``, every 5-th marker will be plotted.
+
+        Parameters
+        ----------
+        every : None or int or (int, int) or slice or list[int] or float or \
+(float, float) or list[bool]
+            Which markers to plot.
+
+            - ``every=None``: every point will be plotted.
+            - ``every=N``: every N-th marker will be plotted starting with
+              marker 0.
+            - ``every=(start, N)``: every N-th marker, starting at index
+              *start*, will be plotted.
+            - ``every=slice(start, end, N)``: every N-th marker, starting at
+              index *start*, up to but not including index *end*, will be
+              plotted.
+            - ``every=[i, j, m, ...]``: only markers at the given indices
+              will be plotted.
+            - ``every=[True, False, True, ...]``: only positions that are True
+              will be plotted. The list must have the same length as the data
+              points.
+            - ``every=0.1``, (i.e. a float): markers will be spaced at
+              approximately equal visual distances along the line; the distance
+              along the line between markers is determined by multiplying the
+              display-coordinate distance of the Axes bounding-box diagonal
+              by the value of *every*.
+            - ``every=(0.5, 0.1)`` (i.e. a length-2 tuple of float): similar
+              to ``every=0.1`` but the first marker will be offset along the
+              line by 0.5 multiplied by the
+              display-coordinate-diagonal-distance along the line.
+
+            For examples see
+            :doc:`/gallery/lines_bars_and_markers/markevery_demo`.
+
+        Notes
+        -----
+        Setting *markevery* will still only draw markers at actual data points.
+        While the float argument form aims for uniform visual spacing, it has
+        to coerce from the ideal spacing to the nearest available data point.
+        Depending on the number and distribution of data points, the result
+        may still not look evenly spaced.
+
+        When using a start offset to specify the first marker, the offset will
+        be from the first data point which may be different from the first
+        the visible data point if the plot is zoomed in.
+
+        If zooming in on a plot when using float arguments then the actual
+        data points that have markers will change because the distance between
+        markers is always determined from the display-coordinates
+        axes-bounding-box-diagonal regardless of the actual axes data limits.
+
+        """
+        self._markevery = every
+        self.stale = True
+
+    def get_markevery(self):
+        """
+        Return the markevery setting for marker subsampling.
+
+        See also `~.Line2D.set_markevery`.
+        """
+        return self._markevery
+
+    def set_picker(self, p):
+        """
+        Set the event picker details for the line.
+
+        Parameters
+        ----------
+        p : float or callable[[Artist, Event], tuple[bool, dict]]
+            If a float, it is used as the pick radius in points.
+        """
+        if not callable(p):
+            self.set_pickradius(p)
+        self._picker = p
+
+    def get_bbox(self):
+        """Get the bounding box of this line."""
+        bbox = Bbox([[0, 0], [0, 0]])
+        bbox.update_from_data_xy(self.get_xydata())
+        return bbox
+
+    def get_window_extent(self, renderer=None):
+        bbox = Bbox([[0, 0], [0, 0]])
+        trans_data_to_xy = self.get_transform().transform
+        bbox.update_from_data_xy(trans_data_to_xy(self.get_xydata()),
+                                 ignore=True)
+        # correct for marker size, if any
+        if self._marker:
+            ms = (self._markersize / 72.0 * self.get_figure(root=True).dpi) * 0.5
+            bbox = bbox.padded(ms)
+        return bbox
+
+    def set_data(self, *args):
+        """
+        Set the x and y data.
+
+        Parameters
+        ----------
+        *args : (2, N) array or two 1D arrays
+
+        See Also
+        --------
+        set_xdata
+        set_ydata
+        """
+        if len(args) == 1:
+            (x, y), = args
+        else:
+            x, y = args
+
+        self.set_xdata(x)
+        self.set_ydata(y)
+
+    def recache_always(self):
+        self.recache(always=True)
+
+    def recache(self, always=False):
+        if always or self._invalidx:
+            xconv = self.convert_xunits(self._xorig)
+            x = _to_unmasked_float_array(xconv).ravel()
+        else:
+            x = self._x
+        if always or self._invalidy:
+            yconv = self.convert_yunits(self._yorig)
+            y = _to_unmasked_float_array(yconv).ravel()
+        else:
+            y = self._y
+
+        self._xy = np.column_stack(np.broadcast_arrays(x, y)).astype(float)
+        self._x, self._y = self._xy.T  # views
+
+        self._subslice = False
+        if (self.axes
+                and len(x) > self._subslice_optim_min_size
+                and _path.is_sorted_and_has_non_nan(x)
+                and self.axes.name == 'rectilinear'
+                and self.axes.get_xscale() == 'linear'
+                and self._markevery is None
+                and self.get_clip_on()
+                and self.get_transform() == self.axes.transData):
+            self._subslice = True
+            nanmask = np.isnan(x)
+            if nanmask.any():
+                self._x_filled = self._x.copy()
+                indices = np.arange(len(x))
+                self._x_filled[nanmask] = np.interp(
+                    indices[nanmask], indices[~nanmask], self._x[~nanmask])
+            else:
+                self._x_filled = self._x
+
+        if self._path is not None:
+            interpolation_steps = self._path._interpolation_steps
+        else:
+            interpolation_steps = 1
+        xy = STEP_LOOKUP_MAP[self._drawstyle](*self._xy.T)
+        self._path = Path(np.asarray(xy).T,
+                          _interpolation_steps=interpolation_steps)
+        self._transformed_path = None
+        self._invalidx = False
+        self._invalidy = False
+
+    def _transform_path(self, subslice=None):
+        """
+        Put a TransformedPath instance at self._transformed_path;
+        all invalidation of the transform is then handled by the
+        TransformedPath instance.
+        """
+        # Masked arrays are now handled by the Path class itself
+        if subslice is not None:
+            xy = STEP_LOOKUP_MAP[self._drawstyle](*self._xy[subslice, :].T)
+            _path = Path(np.asarray(xy).T,
+                         _interpolation_steps=self._path._interpolation_steps)
+        else:
+            _path = self._path
+        self._transformed_path = TransformedPath(_path, self.get_transform())
+
+    def _get_transformed_path(self):
+        """Return this line's `~matplotlib.transforms.TransformedPath`."""
+        if self._transformed_path is None:
+            self._transform_path()
+        return self._transformed_path
+
+    def set_transform(self, t):
+        # docstring inherited
+        self._invalidx = True
+        self._invalidy = True
+        super().set_transform(t)
+
+    @allow_rasterization
+    def draw(self, renderer):
+        # docstring inherited
+
+        if not self.get_visible():
+            return
+
+        if self._invalidy or self._invalidx:
+            self.recache()
+        self.ind_offset = 0  # Needed for contains() method.
+        if self._subslice and self.axes:
+            x0, x1 = self.axes.get_xbound()
+            i0 = self._x_filled.searchsorted(x0, 'left')
+            i1 = self._x_filled.searchsorted(x1, 'right')
+            subslice = slice(max(i0 - 1, 0), i1 + 1)
+            self.ind_offset = subslice.start
+            self._transform_path(subslice)
+        else:
+            subslice = None
+
+        if self.get_path_effects():
+            from matplotlib.patheffects import PathEffectRenderer
+            renderer = PathEffectRenderer(self.get_path_effects(), renderer)
+
+        renderer.open_group('line2d', self.get_gid())
+        if self._lineStyles[self._linestyle] != '_draw_nothing':
+            tpath, affine = (self._get_transformed_path()
+                             .get_transformed_path_and_affine())
+            if len(tpath.vertices):
+                gc = renderer.new_gc()
+                self._set_gc_clip(gc)
+                gc.set_url(self.get_url())
+
+                gc.set_antialiased(self._antialiased)
+                gc.set_linewidth(self._linewidth)
+
+                if self.is_dashed():
+                    cap = self._dashcapstyle
+                    join = self._dashjoinstyle
+                else:
+                    cap = self._solidcapstyle
+                    join = self._solidjoinstyle
+                gc.set_joinstyle(join)
+                gc.set_capstyle(cap)
+                gc.set_snap(self.get_snap())
+                if self.get_sketch_params() is not None:
+                    gc.set_sketch_params(*self.get_sketch_params())
+
+                # We first draw a path within the gaps if needed.
+                if self.is_dashed() and self._gapcolor is not None:
+                    lc_rgba = mcolors.to_rgba(self._gapcolor, self._alpha)
+                    gc.set_foreground(lc_rgba, isRGBA=True)
+
+                    offset_gaps, gaps = _get_inverse_dash_pattern(
+                        *self._dash_pattern)
+
+                    gc.set_dashes(offset_gaps, gaps)
+                    renderer.draw_path(gc, tpath, affine.frozen())
+
+                lc_rgba = mcolors.to_rgba(self._color, self._alpha)
+                gc.set_foreground(lc_rgba, isRGBA=True)
+
+                gc.set_dashes(*self._dash_pattern)
+                renderer.draw_path(gc, tpath, affine.frozen())
+                gc.restore()
+
+        if self._marker and self._markersize > 0:
+            gc = renderer.new_gc()
+            self._set_gc_clip(gc)
+            gc.set_url(self.get_url())
+            gc.set_linewidth(self._markeredgewidth)
+            gc.set_antialiased(self._antialiased)
+
+            ec_rgba = mcolors.to_rgba(
+                self.get_markeredgecolor(), self._alpha)
+            fc_rgba = mcolors.to_rgba(
+                self._get_markerfacecolor(), self._alpha)
+            fcalt_rgba = mcolors.to_rgba(
+                self._get_markerfacecolor(alt=True), self._alpha)
+            # If the edgecolor is "auto", it is set according to the *line*
+            # color but inherits the alpha value of the *face* color, if any.
+            if (cbook._str_equal(self._markeredgecolor, "auto")
+                    and not cbook._str_lower_equal(
+                        self.get_markerfacecolor(), "none")):
+                ec_rgba = ec_rgba[:3] + (fc_rgba[3],)
+            gc.set_foreground(ec_rgba, isRGBA=True)
+            if self.get_sketch_params() is not None:
+                scale, length, randomness = self.get_sketch_params()
+                gc.set_sketch_params(scale/2, length/2, 2*randomness)
+
+            marker = self._marker
+
+            # Markers *must* be drawn ignoring the drawstyle (but don't pay the
+            # recaching if drawstyle is already "default").
+            if self.get_drawstyle() != "default":
+                with cbook._setattr_cm(
+                        self, _drawstyle="default", _transformed_path=None):
+                    self.recache()
+                    self._transform_path(subslice)
+                    tpath, affine = (self._get_transformed_path()
+                                     .get_transformed_points_and_affine())
+            else:
+                tpath, affine = (self._get_transformed_path()
+                                 .get_transformed_points_and_affine())
+
+            if len(tpath.vertices):
+                # subsample the markers if markevery is not None
+                markevery = self.get_markevery()
+                if markevery is not None:
+                    subsampled = _mark_every_path(
+                        markevery, tpath, affine, self.axes)
+                else:
+                    subsampled = tpath
+
+                snap = marker.get_snap_threshold()
+                if isinstance(snap, Real):
+                    snap = renderer.points_to_pixels(self._markersize) >= snap
+                gc.set_snap(snap)
+                gc.set_joinstyle(marker.get_joinstyle())
+                gc.set_capstyle(marker.get_capstyle())
+                marker_path = marker.get_path()
+                marker_trans = marker.get_transform()
+                w = renderer.points_to_pixels(self._markersize)
+
+                if cbook._str_equal(marker.get_marker(), ","):
+                    gc.set_linewidth(0)
+                else:
+                    # Don't scale for pixels, and don't stroke them
+                    marker_trans = marker_trans.scale(w)
+                renderer.draw_markers(gc, marker_path, marker_trans,
+                                      subsampled, affine.frozen(),
+                                      fc_rgba)
+
+                alt_marker_path = marker.get_alt_path()
+                if alt_marker_path:
+                    alt_marker_trans = marker.get_alt_transform()
+                    alt_marker_trans = alt_marker_trans.scale(w)
+                    renderer.draw_markers(
+                            gc, alt_marker_path, alt_marker_trans, subsampled,
+                            affine.frozen(), fcalt_rgba)
+
+            gc.restore()
+
+        renderer.close_group('line2d')
+        self.stale = False
+
+    def get_antialiased(self):
+        """Return whether antialiased rendering is used."""
+        return self._antialiased
+
+    def get_color(self):
+        """
+        Return the line color.
+
+        See also `~.Line2D.set_color`.
+        """
+        return self._color
+
+    def get_drawstyle(self):
+        """
+        Return the drawstyle.
+
+        See also `~.Line2D.set_drawstyle`.
+        """
+        return self._drawstyle
+
+    def get_gapcolor(self):
+        """
+        Return the line gapcolor.
+
+        See also `~.Line2D.set_gapcolor`.
+        """
+        return self._gapcolor
+
+    def get_linestyle(self):
+        """
+        Return the linestyle.
+
+        See also `~.Line2D.set_linestyle`.
+        """
+        return self._linestyle
+
+    def get_linewidth(self):
+        """
+        Return the linewidth in points.
+
+        See also `~.Line2D.set_linewidth`.
+        """
+        return self._linewidth
+
+    def get_marker(self):
+        """
+        Return the line marker.
+
+        See also `~.Line2D.set_marker`.
+        """
+        return self._marker.get_marker()
+
+    def get_markeredgecolor(self):
+        """
+        Return the marker edge color.
+
+        See also `~.Line2D.set_markeredgecolor`.
+        """
+        mec = self._markeredgecolor
+        if cbook._str_equal(mec, 'auto'):
+            if mpl.rcParams['_internal.classic_mode']:
+                if self._marker.get_marker() in ('.', ','):
+                    return self._color
+                if (self._marker.is_filled()
+                        and self._marker.get_fillstyle() != 'none'):
+                    return 'k'  # Bad hard-wired default...
+            return self._color
+        else:
+            return mec
+
+    def get_markeredgewidth(self):
+        """
+        Return the marker edge width in points.
+
+        See also `~.Line2D.set_markeredgewidth`.
+        """
+        return self._markeredgewidth
+
+    def _get_markerfacecolor(self, alt=False):
+        if self._marker.get_fillstyle() == 'none':
+            return 'none'
+        fc = self._markerfacecoloralt if alt else self._markerfacecolor
+        if cbook._str_lower_equal(fc, 'auto'):
+            return self._color
+        else:
+            return fc
+
+    def get_markerfacecolor(self):
+        """
+        Return the marker face color.
+
+        See also `~.Line2D.set_markerfacecolor`.
+        """
+        return self._get_markerfacecolor(alt=False)
+
+    def get_markerfacecoloralt(self):
+        """
+        Return the alternate marker face color.
+
+        See also `~.Line2D.set_markerfacecoloralt`.
+        """
+        return self._get_markerfacecolor(alt=True)
+
+    def get_markersize(self):
+        """
+        Return the marker size in points.
+
+        See also `~.Line2D.set_markersize`.
+        """
+        return self._markersize
+
+    def get_data(self, orig=True):
+        """
+        Return the line data as an ``(xdata, ydata)`` pair.
+
+        If *orig* is *True*, return the original data.
+        """
+        return self.get_xdata(orig=orig), self.get_ydata(orig=orig)
+
+    def get_xdata(self, orig=True):
+        """
+        Return the xdata.
+
+        If *orig* is *True*, return the original data, else the
+        processed data.
+        """
+        if orig:
+            return self._xorig
+        if self._invalidx:
+            self.recache()
+        return self._x
+
+    def get_ydata(self, orig=True):
+        """
+        Return the ydata.
+
+        If *orig* is *True*, return the original data, else the
+        processed data.
+        """
+        if orig:
+            return self._yorig
+        if self._invalidy:
+            self.recache()
+        return self._y
+
+    def get_path(self):
+        """Return the `~matplotlib.path.Path` associated with this line."""
+        if self._invalidy or self._invalidx:
+            self.recache()
+        return self._path
+
+    def get_xydata(self):
+        """Return the *xy* data as a (N, 2) array."""
+        if self._invalidy or self._invalidx:
+            self.recache()
+        return self._xy
+
+    def set_antialiased(self, b):
+        """
+        Set whether to use antialiased rendering.
+
+        Parameters
+        ----------
+        b : bool
+        """
+        if self._antialiased != b:
+            self.stale = True
+        self._antialiased = b
+
+    def set_color(self, color):
+        """
+        Set the color of the line.
+
+        Parameters
+        ----------
+        color : :mpltype:`color`
+        """
+        mcolors._check_color_like(color=color)
+        self._color = color
+        self.stale = True
+
+    def set_drawstyle(self, drawstyle):
+        """
+        Set the drawstyle of the plot.
+
+        The drawstyle determines how the points are connected.
+
+        Parameters
+        ----------
+        drawstyle : {'default', 'steps', 'steps-pre', 'steps-mid', \
+'steps-post'}, default: 'default'
+            For 'default', the points are connected with straight lines.
+
+            The steps variants connect the points with step-like lines,
+            i.e. horizontal lines with vertical steps. They differ in the
+            location of the step:
+
+            - 'steps-pre': The step is at the beginning of the line segment,
+              i.e. the line will be at the y-value of point to the right.
+            - 'steps-mid': The step is halfway between the points.
+            - 'steps-post: The step is at the end of the line segment,
+              i.e. the line will be at the y-value of the point to the left.
+            - 'steps' is equal to 'steps-pre' and is maintained for
+              backward-compatibility.
+
+            For examples see :doc:`/gallery/lines_bars_and_markers/step_demo`.
+        """
+        if drawstyle is None:
+            drawstyle = 'default'
+        _api.check_in_list(self.drawStyles, drawstyle=drawstyle)
+        if self._drawstyle != drawstyle:
+            self.stale = True
+            # invalidate to trigger a recache of the path
+            self._invalidx = True
+        self._drawstyle = drawstyle
+
+    def set_gapcolor(self, gapcolor):
+        """
+        Set a color to fill the gaps in the dashed line style.
+
+        .. note::
+
+            Striped lines are created by drawing two interleaved dashed lines.
+            There can be overlaps between those two, which may result in
+            artifacts when using transparency.
+
+            This functionality is experimental and may change.
+
+        Parameters
+        ----------
+        gapcolor : :mpltype:`color` or None
+            The color with which to fill the gaps. If None, the gaps are
+            unfilled.
+        """
+        if gapcolor is not None:
+            mcolors._check_color_like(color=gapcolor)
+        self._gapcolor = gapcolor
+        self.stale = True
+
+    def set_linewidth(self, w):
+        """
+        Set the line width in points.
+
+        Parameters
+        ----------
+        w : float
+            Line width, in points.
+        """
+        w = float(w)
+        if self._linewidth != w:
+            self.stale = True
+        self._linewidth = w
+        self._dash_pattern = _scale_dashes(*self._unscaled_dash_pattern, w)
+
+    def set_linestyle(self, ls):
+        """
+        Set the linestyle of the line.
+
+        Parameters
+        ----------
+        ls : {'-', '--', '-.', ':', '', (offset, on-off-seq), ...}
+            Possible values:
+
+            - A string:
+
+              ==========================================  =================
+              linestyle                                   description
+              ==========================================  =================
+              ``'-'`` or ``'solid'``                      solid line
+              ``'--'`` or  ``'dashed'``                   dashed line
+              ``'-.'`` or  ``'dashdot'``                  dash-dotted line
+              ``':'`` or ``'dotted'``                     dotted line
+              ``'none'``, ``'None'``, ``' '``, or ``''``  draw nothing
+              ==========================================  =================
+
+            - Alternatively a dash tuple of the following form can be
+              provided::
+
+                  (offset, onoffseq)
+
+              where ``onoffseq`` is an even length tuple of on and off ink
+              in points. See also :meth:`set_dashes`.
+
+            For examples see :doc:`/gallery/lines_bars_and_markers/linestyles`.
+        """
+        if isinstance(ls, str):
+            if ls in [' ', '', 'none']:
+                ls = 'None'
+            _api.check_in_list([*self._lineStyles, *ls_mapper_r], ls=ls)
+            if ls not in self._lineStyles:
+                ls = ls_mapper_r[ls]
+            self._linestyle = ls
+        else:
+            self._linestyle = '--'
+        self._unscaled_dash_pattern = _get_dash_pattern(ls)
+        self._dash_pattern = _scale_dashes(
+            *self._unscaled_dash_pattern, self._linewidth)
+        self.stale = True
+
+    @_docstring.interpd
+    def set_marker(self, marker):
+        """
+        Set the line marker.
+
+        Parameters
+        ----------
+        marker : marker style string, `~.path.Path` or `~.markers.MarkerStyle`
+            See `~matplotlib.markers` for full description of possible
+            arguments.
+        """
+        self._marker = MarkerStyle(marker, self._marker.get_fillstyle())
+        self.stale = True
+
+    def _set_markercolor(self, name, has_rcdefault, val):
+        if val is None:
+            val = mpl.rcParams[f"lines.{name}"] if has_rcdefault else "auto"
+        attr = f"_{name}"
+        current = getattr(self, attr)
+        if current is None:
+            self.stale = True
+        else:
+            neq = current != val
+            # Much faster than `np.any(current != val)` if no arrays are used.
+            if neq.any() if isinstance(neq, np.ndarray) else neq:
+                self.stale = True
+        setattr(self, attr, val)
+
+    def set_markeredgecolor(self, ec):
+        """
+        Set the marker edge color.
+
+        Parameters
+        ----------
+        ec : :mpltype:`color`
+        """
+        self._set_markercolor("markeredgecolor", True, ec)
+
+    def set_markerfacecolor(self, fc):
+        """
+        Set the marker face color.
+
+        Parameters
+        ----------
+        fc : :mpltype:`color`
+        """
+        self._set_markercolor("markerfacecolor", True, fc)
+
+    def set_markerfacecoloralt(self, fc):
+        """
+        Set the alternate marker face color.
+
+        Parameters
+        ----------
+        fc : :mpltype:`color`
+        """
+        self._set_markercolor("markerfacecoloralt", False, fc)
+
+    def set_markeredgewidth(self, ew):
+        """
+        Set the marker edge width in points.
+
+        Parameters
+        ----------
+        ew : float
+             Marker edge width, in points.
+        """
+        if ew is None:
+            ew = mpl.rcParams['lines.markeredgewidth']
+        if self._markeredgewidth != ew:
+            self.stale = True
+        self._markeredgewidth = ew
+
+    def set_markersize(self, sz):
+        """
+        Set the marker size in points.
+
+        Parameters
+        ----------
+        sz : float
+             Marker size, in points.
+        """
+        sz = float(sz)
+        if self._markersize != sz:
+            self.stale = True
+        self._markersize = sz
+
+    def set_xdata(self, x):
+        """
+        Set the data array for x.
+
+        Parameters
+        ----------
+        x : 1D array
+
+        See Also
+        --------
+        set_data
+        set_ydata
+        """
+        if not np.iterable(x):
+            raise RuntimeError('x must be a sequence')
+        self._xorig = copy.copy(x)
+        self._invalidx = True
+        self.stale = True
+
+    def set_ydata(self, y):
+        """
+        Set the data array for y.
+
+        Parameters
+        ----------
+        y : 1D array
+
+        See Also
+        --------
+        set_data
+        set_xdata
+        """
+        if not np.iterable(y):
+            raise RuntimeError('y must be a sequence')
+        self._yorig = copy.copy(y)
+        self._invalidy = True
+        self.stale = True
+
+    def set_dashes(self, seq):
+        """
+        Set the dash sequence.
+
+        The dash sequence is a sequence of floats of even length describing
+        the length of dashes and spaces in points.
+
+        For example, (5, 2, 1, 2) describes a sequence of 5 point and 1 point
+        dashes separated by 2 point spaces.
+
+        See also `~.Line2D.set_gapcolor`, which allows those spaces to be
+        filled with a color.
+
+        Parameters
+        ----------
+        seq : sequence of floats (on/off ink in points) or (None, None)
+            If *seq* is empty or ``(None, None)``, the linestyle will be set
+            to solid.
+        """
+        if seq == (None, None) or len(seq) == 0:
+            self.set_linestyle('-')
+        else:
+            self.set_linestyle((0, seq))
+
+    def update_from(self, other):
+        """Copy properties from *other* to self."""
+        super().update_from(other)
+        self._linestyle = other._linestyle
+        self._linewidth = other._linewidth
+        self._color = other._color
+        self._gapcolor = other._gapcolor
+        self._markersize = other._markersize
+        self._markerfacecolor = other._markerfacecolor
+        self._markerfacecoloralt = other._markerfacecoloralt
+        self._markeredgecolor = other._markeredgecolor
+        self._markeredgewidth = other._markeredgewidth
+        self._unscaled_dash_pattern = other._unscaled_dash_pattern
+        self._dash_pattern = other._dash_pattern
+        self._dashcapstyle = other._dashcapstyle
+        self._dashjoinstyle = other._dashjoinstyle
+        self._solidcapstyle = other._solidcapstyle
+        self._solidjoinstyle = other._solidjoinstyle
+
+        self._linestyle = other._linestyle
+        self._marker = MarkerStyle(marker=other._marker)
+        self._drawstyle = other._drawstyle
+
+    @_docstring.interpd
+    def set_dash_joinstyle(self, s):
+        """
+        How to join segments of the line if it `~Line2D.is_dashed`.
+
+        The default joinstyle is :rc:`lines.dash_joinstyle`.
+
+        Parameters
+        ----------
+        s : `.JoinStyle` or %(JoinStyle)s
+        """
+        js = JoinStyle(s)
+        if self._dashjoinstyle != js:
+            self.stale = True
+        self._dashjoinstyle = js
+
+    @_docstring.interpd
+    def set_solid_joinstyle(self, s):
+        """
+        How to join segments if the line is solid (not `~Line2D.is_dashed`).
+
+        The default joinstyle is :rc:`lines.solid_joinstyle`.
+
+        Parameters
+        ----------
+        s : `.JoinStyle` or %(JoinStyle)s
+        """
+        js = JoinStyle(s)
+        if self._solidjoinstyle != js:
+            self.stale = True
+        self._solidjoinstyle = js
+
+    def get_dash_joinstyle(self):
+        """
+        Return the `.JoinStyle` for dashed lines.
+
+        See also `~.Line2D.set_dash_joinstyle`.
+        """
+        return self._dashjoinstyle.name
+
+    def get_solid_joinstyle(self):
+        """
+        Return the `.JoinStyle` for solid lines.
+
+        See also `~.Line2D.set_solid_joinstyle`.
+        """
+        return self._solidjoinstyle.name
+
+    @_docstring.interpd
+    def set_dash_capstyle(self, s):
+        """
+        How to draw the end caps if the line is `~Line2D.is_dashed`.
+
+        The default capstyle is :rc:`lines.dash_capstyle`.
+
+        Parameters
+        ----------
+        s : `.CapStyle` or %(CapStyle)s
+        """
+        cs = CapStyle(s)
+        if self._dashcapstyle != cs:
+            self.stale = True
+        self._dashcapstyle = cs
+
+    @_docstring.interpd
+    def set_solid_capstyle(self, s):
+        """
+        How to draw the end caps if the line is solid (not `~Line2D.is_dashed`)
+
+        The default capstyle is :rc:`lines.solid_capstyle`.
+
+        Parameters
+        ----------
+        s : `.CapStyle` or %(CapStyle)s
+        """
+        cs = CapStyle(s)
+        if self._solidcapstyle != cs:
+            self.stale = True
+        self._solidcapstyle = cs
+
+    def get_dash_capstyle(self):
+        """
+        Return the `.CapStyle` for dashed lines.
+
+        See also `~.Line2D.set_dash_capstyle`.
+        """
+        return self._dashcapstyle.name
+
+    def get_solid_capstyle(self):
+        """
+        Return the `.CapStyle` for solid lines.
+
+        See also `~.Line2D.set_solid_capstyle`.
+        """
+        return self._solidcapstyle.name
+
+    def is_dashed(self):
+        """
+        Return whether line has a dashed linestyle.
+
+        A custom linestyle is assumed to be dashed, we do not inspect the
+        ``onoffseq`` directly.
+
+        See also `~.Line2D.set_linestyle`.
+        """
+        return self._linestyle in ('--', '-.', ':')
+
+
+class AxLine(Line2D):
+    """
+    A helper class that implements `~.Axes.axline`, by recomputing the artist
+    transform at draw time.
+    """
+
+    def __init__(self, xy1, xy2, slope, **kwargs):
+        """
+        Parameters
+        ----------
+        xy1 : (float, float)
+            The first set of (x, y) coordinates for the line to pass through.
+        xy2 : (float, float) or None
+            The second set of (x, y) coordinates for the line to pass through.
+            Both *xy2* and *slope* must be passed, but one of them must be None.
+        slope : float or None
+            The slope of the line. Both *xy2* and *slope* must be passed, but one of
+            them must be None.
+        """
+        super().__init__([0, 1], [0, 1], **kwargs)
+
+        if (xy2 is None and slope is None or
+                xy2 is not None and slope is not None):
+            raise TypeError(
+                "Exactly one of 'xy2' and 'slope' must be given")
+
+        self._slope = slope
+        self._xy1 = xy1
+        self._xy2 = xy2
+
+    def get_transform(self):
+        ax = self.axes
+        points_transform = self._transform - ax.transData + ax.transScale
+
+        if self._xy2 is not None:
+            # two points were given
+            (x1, y1), (x2, y2) = \
+                points_transform.transform([self._xy1, self._xy2])
+            dx = x2 - x1
+            dy = y2 - y1
+            if dx == 0:
+                if dy == 0:
+                    raise ValueError(
+                        f"Cannot draw a line through two identical points "
+                        f"(x={(x1, x2)}, y={(y1, y2)})")
+                slope = np.inf
+            else:
+                slope = dy / dx
+        else:
+            # one point and a slope were given
+            x1, y1 = points_transform.transform(self._xy1)
+            slope = self._slope
+        (vxlo, vylo), (vxhi, vyhi) = ax.transScale.transform(ax.viewLim)
+        # General case: find intersections with view limits in either
+        # direction, and draw between the middle two points.
+        if slope == 0:
+            start = vxlo, y1
+            stop = vxhi, y1
+        elif np.isinf(slope):
+            start = x1, vylo
+            stop = x1, vyhi
+        else:
+            _, start, stop, _ = sorted([
+                (vxlo, y1 + (vxlo - x1) * slope),
+                (vxhi, y1 + (vxhi - x1) * slope),
+                (x1 + (vylo - y1) / slope, vylo),
+                (x1 + (vyhi - y1) / slope, vyhi),
+            ])
+        return (BboxTransformTo(Bbox([start, stop]))
+                + ax.transLimits + ax.transAxes)
+
+    def draw(self, renderer):
+        self._transformed_path = None  # Force regen.
+        super().draw(renderer)
+
+    def get_xy1(self):
+        """Return the *xy1* value of the line."""
+        return self._xy1
+
+    def get_xy2(self):
+        """Return the *xy2* value of the line."""
+        return self._xy2
+
+    def get_slope(self):
+        """Return the *slope* value of the line."""
+        return self._slope
+
+    def set_xy1(self, *args, **kwargs):
+        """
+        Set the *xy1* value of the line.
+
+        Parameters
+        ----------
+        xy1 : tuple[float, float]
+            Points for the line to pass through.
+        """
+        params = _api.select_matching_signature([
+            lambda self, x, y: locals(), lambda self, xy1: locals(),
+        ], self, *args, **kwargs)
+        if "x" in params:
+            _api.warn_deprecated("3.10", message=(
+                "Passing x and y separately to AxLine.set_xy1 is deprecated since "
+                "%(since)s; pass them as a single tuple instead."))
+            xy1 = params["x"], params["y"]
+        else:
+            xy1 = params["xy1"]
+        self._xy1 = xy1
+
+    def set_xy2(self, *args, **kwargs):
+        """
+        Set the *xy2* value of the line.
+
+        .. note::
+
+            You can only set *xy2* if the line was created using the *xy2*
+            parameter. If the line was created using *slope*, please use
+            `~.AxLine.set_slope`.
+
+        Parameters
+        ----------
+        xy2 : tuple[float, float]
+            Points for the line to pass through.
+        """
+        if self._slope is None:
+            params = _api.select_matching_signature([
+                lambda self, x, y: locals(), lambda self, xy2: locals(),
+            ], self, *args, **kwargs)
+            if "x" in params:
+                _api.warn_deprecated("3.10", message=(
+                    "Passing x and y separately to AxLine.set_xy2 is deprecated since "
+                    "%(since)s; pass them as a single tuple instead."))
+                xy2 = params["x"], params["y"]
+            else:
+                xy2 = params["xy2"]
+            self._xy2 = xy2
+        else:
+            raise ValueError("Cannot set an 'xy2' value while 'slope' is set;"
+                             " they differ but their functionalities overlap")
+
+    def set_slope(self, slope):
+        """
+        Set the *slope* value of the line.
+
+        .. note::
+
+            You can only set *slope* if the line was created using the *slope*
+            parameter. If the line was created using *xy2*, please use
+            `~.AxLine.set_xy2`.
+
+        Parameters
+        ----------
+        slope : float
+            The slope of the line.
+        """
+        if self._xy2 is None:
+            self._slope = slope
+        else:
+            raise ValueError("Cannot set a 'slope' value while 'xy2' is set;"
+                             " they differ but their functionalities overlap")
+
+
+class VertexSelector:
+    """
+    Manage the callbacks to maintain a list of selected vertices for `.Line2D`.
+    Derived classes should override the `process_selected` method to do
+    something with the picks.
+
+    Here is an example which highlights the selected verts with red circles::
+
+        import numpy as np
+        import matplotlib.pyplot as plt
+        import matplotlib.lines as lines
+
+        class HighlightSelected(lines.VertexSelector):
+            def __init__(self, line, fmt='ro', **kwargs):
+                super().__init__(line)
+                self.markers, = self.axes.plot([], [], fmt, **kwargs)
+
+            def process_selected(self, ind, xs, ys):
+                self.markers.set_data(xs, ys)
+                self.canvas.draw()
+
+        fig, ax = plt.subplots()
+        x, y = np.random.rand(2, 30)
+        line, = ax.plot(x, y, 'bs-', picker=5)
+
+        selector = HighlightSelected(line)
+        plt.show()
+    """
+
+    def __init__(self, line):
+        """
+        Parameters
+        ----------
+        line : `~matplotlib.lines.Line2D`
+            The line must already have been added to an `~.axes.Axes` and must
+            have its picker property set.
+        """
+        if line.axes is None:
+            raise RuntimeError('You must first add the line to the Axes')
+        if line.get_picker() is None:
+            raise RuntimeError('You must first set the picker property '
+                               'of the line')
+        self.axes = line.axes
+        self.line = line
+        self.cid = self.canvas.callbacks._connect_picklable(
+            'pick_event', self.onpick)
+        self.ind = set()
+
+    canvas = property(lambda self: self.axes.get_figure(root=True).canvas)
+
+    def process_selected(self, ind, xs, ys):
+        """
+        Default "do nothing" implementation of the `process_selected` method.
+
+        Parameters
+        ----------
+        ind : list of int
+            The indices of the selected vertices.
+        xs, ys : array-like
+            The coordinates of the selected vertices.
+        """
+        pass
+
+    def onpick(self, event):
+        """When the line is picked, update the set of selected indices."""
+        if event.artist is not self.line:
+            return
+        self.ind ^= set(event.ind)
+        ind = sorted(self.ind)
+        xdata, ydata = self.line.get_data()
+        self.process_selected(ind, xdata[ind], ydata[ind])
+
+
+lineStyles = Line2D._lineStyles
+lineMarkers = MarkerStyle.markers
+drawStyles = Line2D.drawStyles
+fillStyles = MarkerStyle.fillstyles

lib/python3.12/site-packages/matplotlib/mathtext.py

@@ -0,0 +1,140 @@
+r"""
+A module for parsing a subset of the TeX math syntax and rendering it to a
+Matplotlib backend.
+
+For a tutorial of its usage, see :ref:`mathtext`.  This
+document is primarily concerned with implementation details.
+
+The module uses pyparsing_ to parse the TeX expression.
+
+.. _pyparsing: https://pypi.org/project/pyparsing/
+
+The Bakoma distribution of the TeX Computer Modern fonts, and STIX
+fonts are supported.  There is experimental support for using
+arbitrary fonts, but results may vary without proper tweaking and
+metrics for those fonts.
+"""
+
+import functools
+import logging
+
+import matplotlib as mpl
+from matplotlib import _api, _mathtext
+from matplotlib.ft2font import LoadFlags
+from matplotlib.font_manager import FontProperties
+from ._mathtext import (  # noqa: F401, reexported API
+    RasterParse, VectorParse, get_unicode_index)
+
+_log = logging.getLogger(__name__)
+
+
+get_unicode_index.__module__ = __name__
+
+##############################################################################
+# MAIN
+
+
+class MathTextParser:
+    _parser = None
+    _font_type_mapping = {
+        'cm':          _mathtext.BakomaFonts,
+        'dejavuserif': _mathtext.DejaVuSerifFonts,
+        'dejavusans':  _mathtext.DejaVuSansFonts,
+        'stix':        _mathtext.StixFonts,
+        'stixsans':    _mathtext.StixSansFonts,
+        'custom':      _mathtext.UnicodeFonts,
+    }
+
+    def __init__(self, output):
+        """
+        Create a MathTextParser for the given backend *output*.
+
+        Parameters
+        ----------
+        output : {"path", "agg"}
+            Whether to return a `VectorParse` ("path") or a
+            `RasterParse` ("agg", or its synonym "macosx").
+        """
+        self._output_type = _api.check_getitem(
+            {"path": "vector", "agg": "raster", "macosx": "raster"},
+            output=output.lower())
+
+    def parse(self, s, dpi=72, prop=None, *, antialiased=None):
+        """
+        Parse the given math expression *s* at the given *dpi*.  If *prop* is
+        provided, it is a `.FontProperties` object specifying the "default"
+        font to use in the math expression, used for all non-math text.
+
+        The results are cached, so multiple calls to `parse`
+        with the same expression should be fast.
+
+        Depending on the *output* type, this returns either a `VectorParse` or
+        a `RasterParse`.
+        """
+        # lru_cache can't decorate parse() directly because prop is
+        # mutable, so we key the cache using an internal copy (see
+        # Text._get_text_metrics_with_cache for a similar case); likewise,
+        # we need to check the mutable state of the text.antialiased and
+        # text.hinting rcParams.
+        prop = prop.copy() if prop is not None else None
+        antialiased = mpl._val_or_rc(antialiased, 'text.antialiased')
+        from matplotlib.backends import backend_agg
+        load_glyph_flags = {
+            "vector": LoadFlags.NO_HINTING,
+            "raster": backend_agg.get_hinting_flag(),
+        }[self._output_type]
+        return self._parse_cached(s, dpi, prop, antialiased, load_glyph_flags)
+
+    @functools.lru_cache(50)
+    def _parse_cached(self, s, dpi, prop, antialiased, load_glyph_flags):
+        if prop is None:
+            prop = FontProperties()
+        fontset_class = _api.check_getitem(
+            self._font_type_mapping, fontset=prop.get_math_fontfamily())
+        fontset = fontset_class(prop, load_glyph_flags)
+        fontsize = prop.get_size_in_points()
+
+        if self._parser is None:  # Cache the parser globally.
+            self.__class__._parser = _mathtext.Parser()
+
+        box = self._parser.parse(s, fontset, fontsize, dpi)
+        output = _mathtext.ship(box)
+        if self._output_type == "vector":
+            return output.to_vector()
+        elif self._output_type == "raster":
+            return output.to_raster(antialiased=antialiased)
+
+
+def math_to_image(s, filename_or_obj, prop=None, dpi=None, format=None,
+                  *, color=None):
+    """
+    Given a math expression, renders it in a closely-clipped bounding
+    box to an image file.
+
+    Parameters
+    ----------
+    s : str
+        A math expression.  The math portion must be enclosed in dollar signs.
+    filename_or_obj : str or path-like or file-like
+        Where to write the image data.
+    prop : `.FontProperties`, optional
+        The size and style of the text.
+    dpi : float, optional
+        The output dpi.  If not set, the dpi is determined as for
+        `.Figure.savefig`.
+    format : str, optional
+        The output format, e.g., 'svg', 'pdf', 'ps' or 'png'.  If not set, the
+        format is determined as for `.Figure.savefig`.
+    color : str, optional
+        Foreground color, defaults to :rc:`text.color`.
+    """
+    from matplotlib import figure
+
+    parser = MathTextParser('path')
+    width, height, depth, _, _ = parser.parse(s, dpi=72, prop=prop)
+
+    fig = figure.Figure(figsize=(width / 72.0, height / 72.0))
+    fig.text(0, depth/height, s, fontproperties=prop, color=color)
+    fig.savefig(filename_or_obj, dpi=dpi, format=format)
+
+    return depth

lib/python3.12/site-packages/matplotlib/offsetbox.pyi

@@ -0,0 +1,298 @@
+import matplotlib.artist as martist
+from matplotlib.backend_bases import RendererBase, Event, FigureCanvasBase
+from matplotlib.colors import Colormap, Normalize
+import matplotlib.text as mtext
+from matplotlib.figure import Figure, SubFigure
+from matplotlib.font_manager import FontProperties
+from matplotlib.image import BboxImage
+from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
+from matplotlib.transforms import Bbox, BboxBase, Transform
+from matplotlib.typing import CoordsType
+
+import numpy as np
+from numpy.typing import ArrayLike
+from collections.abc import Callable, Sequence
+from typing import Any, Literal, overload
+
+DEBUG: bool
+
+def _get_packed_offsets(
+    widths: Sequence[float],
+    total: float | None,
+    sep: float | None,
+    mode: Literal["fixed", "expand", "equal"] = ...,
+) -> tuple[float, np.ndarray]: ...
+
+class OffsetBox(martist.Artist):
+    width: float | None
+    height: float | None
+    def __init__(self, *args, **kwargs) -> None: ...
+    def set_figure(self, fig: Figure | SubFigure) -> None: ...
+    def set_offset(
+        self,
+        xy: tuple[float, float]
+        | Callable[[float, float, float, float, RendererBase], tuple[float, float]],
+    ) -> None: ...
+
+    @overload
+    def get_offset(self, bbox: Bbox, renderer: RendererBase) -> tuple[float, float]: ...
+    @overload
+    def get_offset(
+        self,
+        width: float,
+        height: float,
+        xdescent: float,
+        ydescent: float,
+        renderer: RendererBase
+    ) -> tuple[float, float]: ...
+
+    def set_width(self, width: float) -> None: ...
+    def set_height(self, height: float) -> None: ...
+    def get_visible_children(self) -> list[martist.Artist]: ...
+    def get_children(self) -> list[martist.Artist]: ...
+    def get_bbox(self, renderer: RendererBase) -> Bbox: ...
+    def get_window_extent(self, renderer: RendererBase | None = ...) -> Bbox: ...
+
+class PackerBase(OffsetBox):
+    height: float | None
+    width: float | None
+    sep: float | None
+    pad: float | None
+    mode: Literal["fixed", "expand", "equal"]
+    align: Literal["top", "bottom", "left", "right", "center", "baseline"]
+    def __init__(
+        self,
+        pad: float | None = ...,
+        sep: float | None = ...,
+        width: float | None = ...,
+        height: float | None = ...,
+        align: Literal["top", "bottom", "left", "right", "center", "baseline"] = ...,
+        mode: Literal["fixed", "expand", "equal"] = ...,
+        children: list[martist.Artist] | None = ...,
+    ) -> None: ...
+
+class VPacker(PackerBase): ...
+class HPacker(PackerBase): ...
+
+class PaddedBox(OffsetBox):
+    pad: float | None
+    patch: FancyBboxPatch
+    def __init__(
+        self,
+        child: martist.Artist,
+        pad: float | None = ...,
+        *,
+        draw_frame: bool = ...,
+        patch_attrs: dict[str, Any] | None = ...,
+    ) -> None: ...
+    def update_frame(self, bbox: Bbox, fontsize: float | None = ...) -> None: ...
+    def draw_frame(self, renderer: RendererBase) -> None: ...
+
+class DrawingArea(OffsetBox):
+    width: float
+    height: float
+    xdescent: float
+    ydescent: float
+    offset_transform: Transform
+    dpi_transform: Transform
+    def __init__(
+        self,
+        width: float,
+        height: float,
+        xdescent: float = ...,
+        ydescent: float = ...,
+        clip: bool = ...,
+    ) -> None: ...
+    @property
+    def clip_children(self) -> bool: ...
+    @clip_children.setter
+    def clip_children(self, val: bool) -> None: ...
+    def get_transform(self) -> Transform: ...
+
+    # does not accept all options of superclass
+    def set_offset(self, xy: tuple[float, float]) -> None: ...  # type: ignore[override]
+    def get_offset(self) -> tuple[float, float]: ...  # type: ignore[override]
+    def add_artist(self, a: martist.Artist) -> None: ...
+
+class TextArea(OffsetBox):
+    offset_transform: Transform
+    def __init__(
+        self,
+        s: str,
+        *,
+        textprops: dict[str, Any] | None = ...,
+        multilinebaseline: bool = ...,
+    ) -> None: ...
+    def set_text(self, s: str) -> None: ...
+    def get_text(self) -> str: ...
+    def set_multilinebaseline(self, t: bool) -> None: ...
+    def get_multilinebaseline(self) -> bool: ...
+
+    # does not accept all options of superclass
+    def set_offset(self, xy: tuple[float, float]) -> None: ...  # type: ignore[override]
+    def get_offset(self) -> tuple[float, float]: ...  # type: ignore[override]
+
+class AuxTransformBox(OffsetBox):
+    aux_transform: Transform
+    offset_transform: Transform
+    ref_offset_transform: Transform
+    def __init__(self, aux_transform: Transform) -> None: ...
+    def add_artist(self, a: martist.Artist) -> None: ...
+    def get_transform(self) -> Transform: ...
+
+    # does not accept all options of superclass
+    def set_offset(self, xy: tuple[float, float]) -> None: ...  # type: ignore[override]
+    def get_offset(self) -> tuple[float, float]: ...  # type: ignore[override]
+
+class AnchoredOffsetbox(OffsetBox):
+    zorder: float
+    codes: dict[str, int]
+    loc: int
+    borderpad: float
+    pad: float
+    prop: FontProperties
+    patch: FancyBboxPatch
+    def __init__(
+        self,
+        loc: str,
+        *,
+        pad: float = ...,
+        borderpad: float = ...,
+        child: OffsetBox | None = ...,
+        prop: FontProperties | None = ...,
+        frameon: bool = ...,
+        bbox_to_anchor: BboxBase
+        | tuple[float, float]
+        | tuple[float, float, float, float]
+        | None = ...,
+        bbox_transform: Transform | None = ...,
+        **kwargs
+    ) -> None: ...
+    def set_child(self, child: OffsetBox | None) -> None: ...
+    def get_child(self) -> OffsetBox | None: ...
+    def get_children(self) -> list[martist.Artist]: ...
+    def get_bbox_to_anchor(self) -> Bbox: ...
+    def set_bbox_to_anchor(
+        self, bbox: BboxBase, transform: Transform | None = ...
+    ) -> None: ...
+    def update_frame(self, bbox: Bbox, fontsize: float | None = ...) -> None: ...
+
+class AnchoredText(AnchoredOffsetbox):
+    txt: TextArea
+    def __init__(
+        self,
+        s: str,
+        loc: str,
+        *,
+        pad: float = ...,
+        borderpad: float = ...,
+        prop: dict[str, Any] | None = ...,
+        **kwargs
+    ) -> None: ...
+
+class OffsetImage(OffsetBox):
+    image: BboxImage
+    def __init__(
+        self,
+        arr: ArrayLike,
+        *,
+        zoom: float = ...,
+        cmap: Colormap | str | None = ...,
+        norm: Normalize | str | None = ...,
+        interpolation: str | None = ...,
+        origin: Literal["upper", "lower"] | None = ...,
+        filternorm: bool = ...,
+        filterrad: float = ...,
+        resample: bool = ...,
+        dpi_cor: bool = ...,
+        **kwargs
+    ) -> None: ...
+    stale: bool
+    def set_data(self, arr: ArrayLike | None) -> None: ...
+    def get_data(self) -> ArrayLike | None: ...
+    def set_zoom(self, zoom: float) -> None: ...
+    def get_zoom(self) -> float: ...
+    def get_children(self) -> list[martist.Artist]: ...
+    def get_offset(self) -> tuple[float, float]: ...  # type: ignore[override]
+
+class AnnotationBbox(martist.Artist, mtext._AnnotationBase):
+    zorder: float
+    offsetbox: OffsetBox
+    arrowprops: dict[str, Any] | None
+    xybox: tuple[float, float]
+    boxcoords: CoordsType
+    arrow_patch: FancyArrowPatch | None
+    patch: FancyBboxPatch
+    prop: FontProperties
+    def __init__(
+        self,
+        offsetbox: OffsetBox,
+        xy: tuple[float, float],
+        xybox: tuple[float, float] | None = ...,
+        xycoords: CoordsType = ...,
+        boxcoords: CoordsType | None = ...,
+        *,
+        frameon: bool = ...,
+        pad: float = ...,
+        annotation_clip: bool | None = ...,
+        box_alignment: tuple[float, float] = ...,
+        bboxprops: dict[str, Any] | None = ...,
+        arrowprops: dict[str, Any] | None = ...,
+        fontsize: float | str | None = ...,
+        **kwargs
+    ) -> None: ...
+    @property
+    def xyann(self) -> tuple[float, float]: ...
+    @xyann.setter
+    def xyann(self, xyann: tuple[float, float]) -> None: ...
+    @property
+    def anncoords(
+        self,
+    ) -> CoordsType: ...
+    @anncoords.setter
+    def anncoords(
+        self,
+        coords: CoordsType,
+    ) -> None: ...
+    def get_children(self) -> list[martist.Artist]: ...
+    def set_figure(self, fig: Figure | SubFigure) -> None: ...
+    def set_fontsize(self, s: str | float | None = ...) -> None: ...
+    def get_fontsize(self) -> float: ...
+    def get_tightbbox(self, renderer: RendererBase | None = ...) -> Bbox: ...
+    def update_positions(self, renderer: RendererBase) -> None: ...
+
+class DraggableBase:
+    ref_artist: martist.Artist
+    got_artist: bool
+    mouse_x: int
+    mouse_y: int
+    background: Any
+
+    @property
+    def canvas(self) -> FigureCanvasBase: ...
+    @property
+    def cids(self) -> list[int]: ...
+
+    def __init__(self, ref_artist: martist.Artist, use_blit: bool = ...) -> None: ...
+    def on_motion(self, evt: Event) -> None: ...
+    def on_pick(self, evt: Event) -> None: ...
+    def on_release(self, event: Event) -> None: ...
+    def disconnect(self) -> None: ...
+    def save_offset(self) -> None: ...
+    def update_offset(self, dx: float, dy: float) -> None: ...
+    def finalize_offset(self) -> None: ...
+
+class DraggableOffsetBox(DraggableBase):
+    offsetbox: OffsetBox
+    def __init__(
+        self, ref_artist: martist.Artist, offsetbox: OffsetBox, use_blit: bool = ...
+    ) -> None: ...
+    def save_offset(self) -> None: ...
+    def update_offset(self, dx: float, dy: float) -> None: ...
+    def get_loc_in_canvas(self) -> tuple[float, float]: ...
+
+class DraggableAnnotation(DraggableBase):
+    annotation: mtext.Annotation
+    def __init__(self, annotation: mtext.Annotation, use_blit: bool = ...) -> None: ...
+    def save_offset(self) -> None: ...
+    def update_offset(self, dx: float, dy: float) -> None: ...

lib/python3.12/site-packages/matplotlib/sankey.pyi

@@ -0,0 +1,61 @@
+from matplotlib.axes import Axes
+
+from collections.abc import Callable, Iterable
+from typing import Any
+from typing_extensions import Self  # < Py 3.11
+
+import numpy as np
+
+__license__: str
+__credits__: list[str]
+__author__: str
+__version__: str
+
+RIGHT: int
+UP: int
+DOWN: int
+
+# TODO typing units
+class Sankey:
+    diagrams: list[Any]
+    ax: Axes
+    unit: Any
+    format: str | Callable[[float], str]
+    scale: float
+    gap: float
+    radius: float
+    shoulder: float
+    offset: float
+    margin: float
+    pitch: float
+    tolerance: float
+    extent: np.ndarray
+    def __init__(
+        self,
+        ax: Axes | None = ...,
+        scale: float = ...,
+        unit: Any = ...,
+        format: str | Callable[[float], str] = ...,
+        gap: float = ...,
+        radius: float = ...,
+        shoulder: float = ...,
+        offset: float = ...,
+        head_angle: float = ...,
+        margin: float = ...,
+        tolerance: float = ...,
+        **kwargs
+    ) -> None: ...
+    def add(
+        self,
+        patchlabel: str = ...,
+        flows: Iterable[float] | None = ...,
+        orientations: Iterable[int] | None = ...,
+        labels: str | Iterable[str | None] = ...,
+        trunklength: float = ...,
+        pathlengths: float | Iterable[float] = ...,
+        prior: int | None = ...,
+        connect: tuple[int, int] = ...,
+        rotation: float = ...,
+        **kwargs
+    ) -> Self: ...
+    def finish(self) -> list[Any]: ...

lib/python3.12/site-packages/matplotlib/scale.py

@@ -0,0 +1,767 @@
+"""
+Scales define the distribution of data values on an axis, e.g. a log scaling.
+
+The mapping is implemented through `.Transform` subclasses.
+
+The following scales are built-in:
+
+.. _builtin_scales:
+
+============= ===================== ================================ =================================
+Name          Class                 Transform                        Inverted transform
+============= ===================== ================================ =================================
+"asinh"       `AsinhScale`          `AsinhTransform`                 `InvertedAsinhTransform`
+"function"    `FuncScale`           `FuncTransform`                  `FuncTransform`
+"functionlog" `FuncScaleLog`        `FuncTransform` + `LogTransform` `InvertedLogTransform` + `FuncTransform`
+"linear"      `LinearScale`         `.IdentityTransform`             `.IdentityTransform`
+"log"         `LogScale`            `LogTransform`                   `InvertedLogTransform`
+"logit"       `LogitScale`          `LogitTransform`                 `LogisticTransform`
+"symlog"      `SymmetricalLogScale` `SymmetricalLogTransform`        `InvertedSymmetricalLogTransform`
+============= ===================== ================================ =================================
+
+A user will often only use the scale name, e.g. when setting the scale through
+`~.Axes.set_xscale`: ``ax.set_xscale("log")``.
+
+See also the :ref:`scales examples <sphx_glr_gallery_scales>` in the documentation.
+
+Custom scaling can be achieved through `FuncScale`, or by creating your own
+`ScaleBase` subclass and corresponding transforms (see :doc:`/gallery/scales/custom_scale`).
+Third parties can register their scales by name through `register_scale`.
+"""  # noqa: E501
+
+import inspect
+import textwrap
+
+import numpy as np
+
+import matplotlib as mpl
+from matplotlib import _api, _docstring
+from matplotlib.ticker import (
+    NullFormatter, ScalarFormatter, LogFormatterSciNotation, LogitFormatter,
+    NullLocator, LogLocator, AutoLocator, AutoMinorLocator,
+    SymmetricalLogLocator, AsinhLocator, LogitLocator)
+from matplotlib.transforms import Transform, IdentityTransform
+
+
+class ScaleBase:
+    """
+    The base class for all scales.
+
+    Scales are separable transformations, working on a single dimension.
+
+    Subclasses should override
+
+    :attr:`name`
+        The scale's name.
+    :meth:`get_transform`
+        A method returning a `.Transform`, which converts data coordinates to
+        scaled coordinates.  This transform should be invertible, so that e.g.
+        mouse positions can be converted back to data coordinates.
+    :meth:`set_default_locators_and_formatters`
+        A method that sets default locators and formatters for an `~.axis.Axis`
+        that uses this scale.
+    :meth:`limit_range_for_scale`
+        An optional method that "fixes" the axis range to acceptable values,
+        e.g. restricting log-scaled axes to positive values.
+    """
+
+    def __init__(self, axis):
+        r"""
+        Construct a new scale.
+
+        Notes
+        -----
+        The following note is for scale implementers.
+
+        For back-compatibility reasons, scales take an `~matplotlib.axis.Axis`
+        object as first argument.  However, this argument should not
+        be used: a single scale object should be usable by multiple
+        `~matplotlib.axis.Axis`\es at the same time.
+        """
+
+    def get_transform(self):
+        """
+        Return the `.Transform` object associated with this scale.
+        """
+        raise NotImplementedError()
+
+    def set_default_locators_and_formatters(self, axis):
+        """
+        Set the locators and formatters of *axis* to instances suitable for
+        this scale.
+        """
+        raise NotImplementedError()
+
+    def limit_range_for_scale(self, vmin, vmax, minpos):
+        """
+        Return the range *vmin*, *vmax*, restricted to the
+        domain supported by this scale (if any).
+
+        *minpos* should be the minimum positive value in the data.
+        This is used by log scales to determine a minimum value.
+        """
+        return vmin, vmax
+
+
+class LinearScale(ScaleBase):
+    """
+    The default linear scale.
+    """
+
+    name = 'linear'
+
+    def __init__(self, axis):
+        # This method is present only to prevent inheritance of the base class'
+        # constructor docstring, which would otherwise end up interpolated into
+        # the docstring of Axis.set_scale.
+        """
+        """  # noqa: D419
+
+    def set_default_locators_and_formatters(self, axis):
+        # docstring inherited
+        axis.set_major_locator(AutoLocator())
+        axis.set_major_formatter(ScalarFormatter())
+        axis.set_minor_formatter(NullFormatter())
+        # update the minor locator for x and y axis based on rcParams
+        if (axis.axis_name == 'x' and mpl.rcParams['xtick.minor.visible'] or
+                axis.axis_name == 'y' and mpl.rcParams['ytick.minor.visible']):
+            axis.set_minor_locator(AutoMinorLocator())
+        else:
+            axis.set_minor_locator(NullLocator())
+
+    def get_transform(self):
+        """
+        Return the transform for linear scaling, which is just the
+        `~matplotlib.transforms.IdentityTransform`.
+        """
+        return IdentityTransform()
+
+
+class FuncTransform(Transform):
+    """
+    A simple transform that takes and arbitrary function for the
+    forward and inverse transform.
+    """
+
+    input_dims = output_dims = 1
+
+    def __init__(self, forward, inverse):
+        """
+        Parameters
+        ----------
+        forward : callable
+            The forward function for the transform.  This function must have
+            an inverse and, for best behavior, be monotonic.
+            It must have the signature::
+
+               def forward(values: array-like) -> array-like
+
+        inverse : callable
+            The inverse of the forward function.  Signature as ``forward``.
+        """
+        super().__init__()
+        if callable(forward) and callable(inverse):
+            self._forward = forward
+            self._inverse = inverse
+        else:
+            raise ValueError('arguments to FuncTransform must be functions')
+
+    def transform_non_affine(self, values):
+        return self._forward(values)
+
+    def inverted(self):
+        return FuncTransform(self._inverse, self._forward)
+
+
+class FuncScale(ScaleBase):
+    """
+    Provide an arbitrary scale with user-supplied function for the axis.
+    """
+
+    name = 'function'
+
+    def __init__(self, axis, functions):
+        """
+        Parameters
+        ----------
+        axis : `~matplotlib.axis.Axis`
+            The axis for the scale.
+        functions : (callable, callable)
+            two-tuple of the forward and inverse functions for the scale.
+            The forward function must be monotonic.
+
+            Both functions must have the signature::
+
+               def forward(values: array-like) -> array-like
+        """
+        forward, inverse = functions
+        transform = FuncTransform(forward, inverse)
+        self._transform = transform
+
+    def get_transform(self):
+        """Return the `.FuncTransform` associated with this scale."""
+        return self._transform
+
+    def set_default_locators_and_formatters(self, axis):
+        # docstring inherited
+        axis.set_major_locator(AutoLocator())
+        axis.set_major_formatter(ScalarFormatter())
+        axis.set_minor_formatter(NullFormatter())
+        # update the minor locator for x and y axis based on rcParams
+        if (axis.axis_name == 'x' and mpl.rcParams['xtick.minor.visible'] or
+                axis.axis_name == 'y' and mpl.rcParams['ytick.minor.visible']):
+            axis.set_minor_locator(AutoMinorLocator())
+        else:
+            axis.set_minor_locator(NullLocator())
+
+
+class LogTransform(Transform):
+    input_dims = output_dims = 1
+
+    def __init__(self, base, nonpositive='clip'):
+        super().__init__()
+        if base <= 0 or base == 1:
+            raise ValueError('The log base cannot be <= 0 or == 1')
+        self.base = base
+        self._clip = _api.check_getitem(
+            {"clip": True, "mask": False}, nonpositive=nonpositive)
+
+    def __str__(self):
+        return "{}(base={}, nonpositive={!r})".format(
+            type(self).__name__, self.base, "clip" if self._clip else "mask")
+
+    def transform_non_affine(self, values):
+        # Ignore invalid values due to nans being passed to the transform.
+        with np.errstate(divide="ignore", invalid="ignore"):
+            log = {np.e: np.log, 2: np.log2, 10: np.log10}.get(self.base)
+            if log:  # If possible, do everything in a single call to NumPy.
+                out = log(values)
+            else:
+                out = np.log(values)
+                out /= np.log(self.base)
+            if self._clip:
+                # SVG spec says that conforming viewers must support values up
+                # to 3.4e38 (C float); however experiments suggest that
+                # Inkscape (which uses cairo for rendering) runs into cairo's
+                # 24-bit limit (which is apparently shared by Agg).
+                # Ghostscript (used for pdf rendering appears to overflow even
+                # earlier, with the max value around 2 ** 15 for the tests to
+                # pass. On the other hand, in practice, we want to clip beyond
+                #     np.log10(np.nextafter(0, 1)) ~ -323
+                # so 1000 seems safe.
+                out[values <= 0] = -1000
+        return out
+
+    def inverted(self):
+        return InvertedLogTransform(self.base)
+
+
+class InvertedLogTransform(Transform):
+    input_dims = output_dims = 1
+
+    def __init__(self, base):
+        super().__init__()
+        self.base = base
+
+    def __str__(self):
+        return f"{type(self).__name__}(base={self.base})"
+
+    def transform_non_affine(self, values):
+        return np.power(self.base, values)
+
+    def inverted(self):
+        return LogTransform(self.base)
+
+
+class LogScale(ScaleBase):
+    """
+    A standard logarithmic scale.  Care is taken to only plot positive values.
+    """
+    name = 'log'
+
+    def __init__(self, axis, *, base=10, subs=None, nonpositive="clip"):
+        """
+        Parameters
+        ----------
+        axis : `~matplotlib.axis.Axis`
+            The axis for the scale.
+        base : float, default: 10
+            The base of the logarithm.
+        nonpositive : {'clip', 'mask'}, default: 'clip'
+            Determines the behavior for non-positive values. They can either
+            be masked as invalid, or clipped to a very small positive number.
+        subs : sequence of int, default: None
+            Where to place the subticks between each major tick.  For example,
+            in a log10 scale, ``[2, 3, 4, 5, 6, 7, 8, 9]`` will place 8
+            logarithmically spaced minor ticks between each major tick.
+        """
+        self._transform = LogTransform(base, nonpositive)
+        self.subs = subs
+
+    base = property(lambda self: self._transform.base)
+
+    def set_default_locators_and_formatters(self, axis):
+        # docstring inherited
+        axis.set_major_locator(LogLocator(self.base))
+        axis.set_major_formatter(LogFormatterSciNotation(self.base))
+        axis.set_minor_locator(LogLocator(self.base, self.subs))
+        axis.set_minor_formatter(
+            LogFormatterSciNotation(self.base,
+                                    labelOnlyBase=(self.subs is not None)))
+
+    def get_transform(self):
+        """Return the `.LogTransform` associated with this scale."""
+        return self._transform
+
+    def limit_range_for_scale(self, vmin, vmax, minpos):
+        """Limit the domain to positive values."""
+        if not np.isfinite(minpos):
+            minpos = 1e-300  # Should rarely (if ever) have a visible effect.
+
+        return (minpos if vmin <= 0 else vmin,
+                minpos if vmax <= 0 else vmax)
+
+
+class FuncScaleLog(LogScale):
+    """
+    Provide an arbitrary scale with user-supplied function for the axis and
+    then put on a logarithmic axes.
+    """
+
+    name = 'functionlog'
+
+    def __init__(self, axis, functions, base=10):
+        """
+        Parameters
+        ----------
+        axis : `~matplotlib.axis.Axis`
+            The axis for the scale.
+        functions : (callable, callable)
+            two-tuple of the forward and inverse functions for the scale.
+            The forward function must be monotonic.
+
+            Both functions must have the signature::
+
+                def forward(values: array-like) -> array-like
+
+        base : float, default: 10
+            Logarithmic base of the scale.
+        """
+        forward, inverse = functions
+        self.subs = None
+        self._transform = FuncTransform(forward, inverse) + LogTransform(base)
+
+    @property
+    def base(self):
+        return self._transform._b.base  # Base of the LogTransform.
+
+    def get_transform(self):
+        """Return the `.Transform` associated with this scale."""
+        return self._transform
+
+
+class SymmetricalLogTransform(Transform):
+    input_dims = output_dims = 1
+
+    def __init__(self, base, linthresh, linscale):
+        super().__init__()
+        if base <= 1.0:
+            raise ValueError("'base' must be larger than 1")
+        if linthresh <= 0.0:
+            raise ValueError("'linthresh' must be positive")
+        if linscale <= 0.0:
+            raise ValueError("'linscale' must be positive")
+        self.base = base
+        self.linthresh = linthresh
+        self.linscale = linscale
+        self._linscale_adj = (linscale / (1.0 - self.base ** -1))
+        self._log_base = np.log(base)
+
+    def transform_non_affine(self, values):
+        abs_a = np.abs(values)
+        with np.errstate(divide="ignore", invalid="ignore"):
+            out = np.sign(values) * self.linthresh * (
+                self._linscale_adj +
+                np.log(abs_a / self.linthresh) / self._log_base)
+            inside = abs_a <= self.linthresh
+        out[inside] = values[inside] * self._linscale_adj
+        return out
+
+    def inverted(self):
+        return InvertedSymmetricalLogTransform(self.base, self.linthresh,
+                                               self.linscale)
+
+
+class InvertedSymmetricalLogTransform(Transform):
+    input_dims = output_dims = 1
+
+    def __init__(self, base, linthresh, linscale):
+        super().__init__()
+        symlog = SymmetricalLogTransform(base, linthresh, linscale)
+        self.base = base
+        self.linthresh = linthresh
+        self.invlinthresh = symlog.transform(linthresh)
+        self.linscale = linscale
+        self._linscale_adj = (linscale / (1.0 - self.base ** -1))
+
+    def transform_non_affine(self, values):
+        abs_a = np.abs(values)
+        with np.errstate(divide="ignore", invalid="ignore"):
+            out = np.sign(values) * self.linthresh * (
+                np.power(self.base,
+                         abs_a / self.linthresh - self._linscale_adj))
+            inside = abs_a <= self.invlinthresh
+        out[inside] = values[inside] / self._linscale_adj
+        return out
+
+    def inverted(self):
+        return SymmetricalLogTransform(self.base,
+                                       self.linthresh, self.linscale)
+
+
+class SymmetricalLogScale(ScaleBase):
+    """
+    The symmetrical logarithmic scale is logarithmic in both the
+    positive and negative directions from the origin.
+
+    Since the values close to zero tend toward infinity, there is a
+    need to have a range around zero that is linear.  The parameter
+    *linthresh* allows the user to specify the size of this range
+    (-*linthresh*, *linthresh*).
+
+    See :doc:`/gallery/scales/symlog_demo` for a detailed description.
+
+    Parameters
+    ----------
+    base : float, default: 10
+        The base of the logarithm.
+
+    linthresh : float, default: 2
+        Defines the range ``(-x, x)``, within which the plot is linear.
+        This avoids having the plot go to infinity around zero.
+
+    subs : sequence of int
+        Where to place the subticks between each major tick.
+        For example, in a log10 scale: ``[2, 3, 4, 5, 6, 7, 8, 9]`` will place
+        8 logarithmically spaced minor ticks between each major tick.
+
+    linscale : float, optional
+        This allows the linear range ``(-linthresh, linthresh)`` to be
+        stretched relative to the logarithmic range. Its value is the number of
+        decades to use for each half of the linear range. For example, when
+        *linscale* == 1.0 (the default), the space used for the positive and
+        negative halves of the linear range will be equal to one decade in
+        the logarithmic range.
+    """
+    name = 'symlog'
+
+    def __init__(self, axis, *, base=10, linthresh=2, subs=None, linscale=1):
+        self._transform = SymmetricalLogTransform(base, linthresh, linscale)
+        self.subs = subs
+
+    base = property(lambda self: self._transform.base)
+    linthresh = property(lambda self: self._transform.linthresh)
+    linscale = property(lambda self: self._transform.linscale)
+
+    def set_default_locators_and_formatters(self, axis):
+        # docstring inherited
+        axis.set_major_locator(SymmetricalLogLocator(self.get_transform()))
+        axis.set_major_formatter(LogFormatterSciNotation(self.base))
+        axis.set_minor_locator(SymmetricalLogLocator(self.get_transform(),
+                                                     self.subs))
+        axis.set_minor_formatter(NullFormatter())
+
+    def get_transform(self):
+        """Return the `.SymmetricalLogTransform` associated with this scale."""
+        return self._transform
+
+
+class AsinhTransform(Transform):
+    """Inverse hyperbolic-sine transformation used by `.AsinhScale`"""
+    input_dims = output_dims = 1
+
+    def __init__(self, linear_width):
+        super().__init__()
+        if linear_width <= 0.0:
+            raise ValueError("Scale parameter 'linear_width' " +
+                             "must be strictly positive")
+        self.linear_width = linear_width
+
+    def transform_non_affine(self, values):
+        return self.linear_width * np.arcsinh(values / self.linear_width)
+
+    def inverted(self):
+        return InvertedAsinhTransform(self.linear_width)
+
+
+class InvertedAsinhTransform(Transform):
+    """Hyperbolic sine transformation used by `.AsinhScale`"""
+    input_dims = output_dims = 1
+
+    def __init__(self, linear_width):
+        super().__init__()
+        self.linear_width = linear_width
+
+    def transform_non_affine(self, values):
+        return self.linear_width * np.sinh(values / self.linear_width)
+
+    def inverted(self):
+        return AsinhTransform(self.linear_width)
+
+
+class AsinhScale(ScaleBase):
+    """
+    A quasi-logarithmic scale based on the inverse hyperbolic sine (asinh)
+
+    For values close to zero, this is essentially a linear scale,
+    but for large magnitude values (either positive or negative)
+    it is asymptotically logarithmic. The transition between these
+    linear and logarithmic regimes is smooth, and has no discontinuities
+    in the function gradient in contrast to
+    the `.SymmetricalLogScale` ("symlog") scale.
+
+    Specifically, the transformation of an axis coordinate :math:`a` is
+    :math:`a \\rightarrow a_0 \\sinh^{-1} (a / a_0)` where :math:`a_0`
+    is the effective width of the linear region of the transformation.
+    In that region, the transformation is
+    :math:`a \\rightarrow a + \\mathcal{O}(a^3)`.
+    For large values of :math:`a` the transformation behaves as
+    :math:`a \\rightarrow a_0 \\, \\mathrm{sgn}(a) \\ln |a| + \\mathcal{O}(1)`.
+
+    .. note::
+
+       This API is provisional and may be revised in the future
+       based on early user feedback.
+    """
+
+    name = 'asinh'
+
+    auto_tick_multipliers = {
+        3: (2, ),
+        4: (2, ),
+        5: (2, ),
+        8: (2, 4),
+        10: (2, 5),
+        16: (2, 4, 8),
+        64: (4, 16),
+        1024: (256, 512)
+    }
+
+    def __init__(self, axis, *, linear_width=1.0,
+                 base=10, subs='auto', **kwargs):
+        """
+        Parameters
+        ----------
+        linear_width : float, default: 1
+            The scale parameter (elsewhere referred to as :math:`a_0`)
+            defining the extent of the quasi-linear region,
+            and the coordinate values beyond which the transformation
+            becomes asymptotically logarithmic.
+        base : int, default: 10
+            The number base used for rounding tick locations
+            on a logarithmic scale. If this is less than one,
+            then rounding is to the nearest integer multiple
+            of powers of ten.
+        subs : sequence of int
+            Multiples of the number base used for minor ticks.
+            If set to 'auto', this will use built-in defaults,
+            e.g. (2, 5) for base=10.
+        """
+        super().__init__(axis)
+        self._transform = AsinhTransform(linear_width)
+        self._base = int(base)
+        if subs == 'auto':
+            self._subs = self.auto_tick_multipliers.get(self._base)
+        else:
+            self._subs = subs
+
+    linear_width = property(lambda self: self._transform.linear_width)
+
+    def get_transform(self):
+        return self._transform
+
+    def set_default_locators_and_formatters(self, axis):
+        axis.set(major_locator=AsinhLocator(self.linear_width,
+                                            base=self._base),
+                 minor_locator=AsinhLocator(self.linear_width,
+                                            base=self._base,
+                                            subs=self._subs),
+                 minor_formatter=NullFormatter())
+        if self._base > 1:
+            axis.set_major_formatter(LogFormatterSciNotation(self._base))
+        else:
+            axis.set_major_formatter('{x:.3g}')
+
+
+class LogitTransform(Transform):
+    input_dims = output_dims = 1
+
+    def __init__(self, nonpositive='mask'):
+        super().__init__()
+        _api.check_in_list(['mask', 'clip'], nonpositive=nonpositive)
+        self._nonpositive = nonpositive
+        self._clip = {"clip": True, "mask": False}[nonpositive]
+
+    def transform_non_affine(self, values):
+        """logit transform (base 10), masked or clipped"""
+        with np.errstate(divide="ignore", invalid="ignore"):
+            out = np.log10(values / (1 - values))
+        if self._clip:  # See LogTransform for choice of clip value.
+            out[values <= 0] = -1000
+            out[1 <= values] = 1000
+        return out
+
+    def inverted(self):
+        return LogisticTransform(self._nonpositive)
+
+    def __str__(self):
+        return f"{type(self).__name__}({self._nonpositive!r})"
+
+
+class LogisticTransform(Transform):
+    input_dims = output_dims = 1
+
+    def __init__(self, nonpositive='mask'):
+        super().__init__()
+        self._nonpositive = nonpositive
+
+    def transform_non_affine(self, values):
+        """logistic transform (base 10)"""
+        return 1.0 / (1 + 10**(-values))
+
+    def inverted(self):
+        return LogitTransform(self._nonpositive)
+
+    def __str__(self):
+        return f"{type(self).__name__}({self._nonpositive!r})"
+
+
+class LogitScale(ScaleBase):
+    """
+    Logit scale for data between zero and one, both excluded.
+
+    This scale is similar to a log scale close to zero and to one, and almost
+    linear around 0.5. It maps the interval ]0, 1[ onto ]-infty, +infty[.
+    """
+    name = 'logit'
+
+    def __init__(self, axis, nonpositive='mask', *,
+                 one_half=r"\frac{1}{2}", use_overline=False):
+        r"""
+        Parameters
+        ----------
+        axis : `~matplotlib.axis.Axis`
+            Currently unused.
+        nonpositive : {'mask', 'clip'}
+            Determines the behavior for values beyond the open interval ]0, 1[.
+            They can either be masked as invalid, or clipped to a number very
+            close to 0 or 1.
+        use_overline : bool, default: False
+            Indicate the usage of survival notation (\overline{x}) in place of
+            standard notation (1-x) for probability close to one.
+        one_half : str, default: r"\frac{1}{2}"
+            The string used for ticks formatter to represent 1/2.
+        """
+        self._transform = LogitTransform(nonpositive)
+        self._use_overline = use_overline
+        self._one_half = one_half
+
+    def get_transform(self):
+        """Return the `.LogitTransform` associated with this scale."""
+        return self._transform
+
+    def set_default_locators_and_formatters(self, axis):
+        # docstring inherited
+        # ..., 0.01, 0.1, 0.5, 0.9, 0.99, ...
+        axis.set_major_locator(LogitLocator())
+        axis.set_major_formatter(
+            LogitFormatter(
+                one_half=self._one_half,
+                use_overline=self._use_overline
+            )
+        )
+        axis.set_minor_locator(LogitLocator(minor=True))
+        axis.set_minor_formatter(
+            LogitFormatter(
+                minor=True,
+                one_half=self._one_half,
+                use_overline=self._use_overline
+            )
+        )
+
+    def limit_range_for_scale(self, vmin, vmax, minpos):
+        """
+        Limit the domain to values between 0 and 1 (excluded).
+        """
+        if not np.isfinite(minpos):
+            minpos = 1e-7  # Should rarely (if ever) have a visible effect.
+        return (minpos if vmin <= 0 else vmin,
+                1 - minpos if vmax >= 1 else vmax)
+
+
+_scale_mapping = {
+    'linear': LinearScale,
+    'log':    LogScale,
+    'symlog': SymmetricalLogScale,
+    'asinh':  AsinhScale,
+    'logit':  LogitScale,
+    'function': FuncScale,
+    'functionlog': FuncScaleLog,
+    }
+
+
+def get_scale_names():
+    """Return the names of the available scales."""
+    return sorted(_scale_mapping)
+
+
+def scale_factory(scale, axis, **kwargs):
+    """
+    Return a scale class by name.
+
+    Parameters
+    ----------
+    scale : {%(names)s}
+    axis : `~matplotlib.axis.Axis`
+    """
+    scale_cls = _api.check_getitem(_scale_mapping, scale=scale)
+    return scale_cls(axis, **kwargs)
+
+
+if scale_factory.__doc__:
+    scale_factory.__doc__ = scale_factory.__doc__ % {
+        "names": ", ".join(map(repr, get_scale_names()))}
+
+
+def register_scale(scale_class):
+    """
+    Register a new kind of scale.
+
+    Parameters
+    ----------
+    scale_class : subclass of `ScaleBase`
+        The scale to register.
+    """
+    _scale_mapping[scale_class.name] = scale_class
+
+
+def _get_scale_docs():
+    """
+    Helper function for generating docstrings related to scales.
+    """
+    docs = []
+    for name, scale_class in _scale_mapping.items():
+        docstring = inspect.getdoc(scale_class.__init__) or ""
+        docs.extend([
+            f"    {name!r}",
+            "",
+            textwrap.indent(docstring, " " * 8),
+            ""
+        ])
+    return "\n".join(docs)
+
+
+_docstring.interpd.register(
+    scale_type='{%s}' % ', '.join([repr(x) for x in get_scale_names()]),
+    scale_docs=_get_scale_docs().rstrip(),
+    )

lib/python3.12/site-packages/matplotlib/texmanager.py

@@ -0,0 +1,368 @@
+r"""
+Support for embedded TeX expressions in Matplotlib.
+
+Requirements:
+
+* LaTeX.
+* \*Agg backends: dvipng>=1.6.
+* PS backend: PSfrag, dvips, and Ghostscript>=9.0.
+* PDF and SVG backends: if LuaTeX is present, it will be used to speed up some
+  post-processing steps, but note that it is not used to parse the TeX string
+  itself (only LaTeX is supported).
+
+To enable TeX rendering of all text in your Matplotlib figure, set
+:rc:`text.usetex` to True.
+
+TeX and dvipng/dvips processing results are cached
+in ~/.matplotlib/tex.cache for reuse between sessions.
+
+`TexManager.get_rgba` can also be used to directly obtain raster output as RGBA
+NumPy arrays.
+"""
+
+import functools
+import hashlib
+import logging
+import os
+from pathlib import Path
+import subprocess
+from tempfile import TemporaryDirectory
+
+import numpy as np
+
+import matplotlib as mpl
+from matplotlib import cbook, dviread
+
+_log = logging.getLogger(__name__)
+
+
+def _usepackage_if_not_loaded(package, *, option=None):
+    """
+    Output LaTeX code that loads a package (possibly with an option) if it
+    hasn't been loaded yet.
+
+    LaTeX cannot load twice a package with different options, so this helper
+    can be used to protect against users loading arbitrary packages/options in
+    their custom preamble.
+    """
+    option = f"[{option}]" if option is not None else ""
+    return (
+        r"\makeatletter"
+        r"\@ifpackageloaded{%(package)s}{}{\usepackage%(option)s{%(package)s}}"
+        r"\makeatother"
+    ) % {"package": package, "option": option}
+
+
+class TexManager:
+    """
+    Convert strings to dvi files using TeX, caching the results to a directory.
+
+    The cache directory is called ``tex.cache`` and is located in the directory
+    returned by `.get_cachedir`.
+
+    Repeated calls to this constructor always return the same instance.
+    """
+
+    _texcache = os.path.join(mpl.get_cachedir(), 'tex.cache')
+    _grey_arrayd = {}
+
+    _font_families = ('serif', 'sans-serif', 'cursive', 'monospace')
+    _font_preambles = {
+        'new century schoolbook': r'\renewcommand{\rmdefault}{pnc}',
+        'bookman': r'\renewcommand{\rmdefault}{pbk}',
+        'times': r'\usepackage{mathptmx}',
+        'palatino': r'\usepackage{mathpazo}',
+        'zapf chancery': r'\usepackage{chancery}',
+        'cursive': r'\usepackage{chancery}',
+        'charter': r'\usepackage{charter}',
+        'serif': '',
+        'sans-serif': '',
+        'helvetica': r'\usepackage{helvet}',
+        'avant garde': r'\usepackage{avant}',
+        'courier': r'\usepackage{courier}',
+        # Loading the type1ec package ensures that cm-super is installed, which
+        # is necessary for Unicode computer modern.  (It also allows the use of
+        # computer modern at arbitrary sizes, but that's just a side effect.)
+        'monospace': r'\usepackage{type1ec}',
+        'computer modern roman': r'\usepackage{type1ec}',
+        'computer modern sans serif': r'\usepackage{type1ec}',
+        'computer modern typewriter': r'\usepackage{type1ec}',
+    }
+    _font_types = {
+        'new century schoolbook': 'serif',
+        'bookman': 'serif',
+        'times': 'serif',
+        'palatino': 'serif',
+        'zapf chancery': 'cursive',
+        'charter': 'serif',
+        'helvetica': 'sans-serif',
+        'avant garde': 'sans-serif',
+        'courier': 'monospace',
+        'computer modern roman': 'serif',
+        'computer modern sans serif': 'sans-serif',
+        'computer modern typewriter': 'monospace',
+    }
+
+    @functools.lru_cache  # Always return the same instance.
+    def __new__(cls):
+        Path(cls._texcache).mkdir(parents=True, exist_ok=True)
+        return object.__new__(cls)
+
+    @classmethod
+    def _get_font_family_and_reduced(cls):
+        """Return the font family name and whether the font is reduced."""
+        ff = mpl.rcParams['font.family']
+        ff_val = ff[0].lower() if len(ff) == 1 else None
+        if len(ff) == 1 and ff_val in cls._font_families:
+            return ff_val, False
+        elif len(ff) == 1 and ff_val in cls._font_preambles:
+            return cls._font_types[ff_val], True
+        else:
+            _log.info('font.family must be one of (%s) when text.usetex is '
+                      'True. serif will be used by default.',
+                      ', '.join(cls._font_families))
+            return 'serif', False
+
+    @classmethod
+    def _get_font_preamble_and_command(cls):
+        requested_family, is_reduced_font = cls._get_font_family_and_reduced()
+
+        preambles = {}
+        for font_family in cls._font_families:
+            if is_reduced_font and font_family == requested_family:
+                preambles[font_family] = cls._font_preambles[
+                    mpl.rcParams['font.family'][0].lower()]
+            else:
+                rcfonts = mpl.rcParams[f"font.{font_family}"]
+                for i, font in enumerate(map(str.lower, rcfonts)):
+                    if font in cls._font_preambles:
+                        preambles[font_family] = cls._font_preambles[font]
+                        _log.debug(
+                            'family: %s, package: %s, font: %s, skipped: %s',
+                            font_family, cls._font_preambles[font], rcfonts[i],
+                            ', '.join(rcfonts[:i]),
+                        )
+                        break
+                else:
+                    _log.info('No LaTeX-compatible font found for the %s font'
+                              'family in rcParams. Using default.',
+                              font_family)
+                    preambles[font_family] = cls._font_preambles[font_family]
+
+        # The following packages and commands need to be included in the latex
+        # file's preamble:
+        cmd = {preambles[family]
+               for family in ['serif', 'sans-serif', 'monospace']}
+        if requested_family == 'cursive':
+            cmd.add(preambles['cursive'])
+        cmd.add(r'\usepackage{type1cm}')
+        preamble = '\n'.join(sorted(cmd))
+        fontcmd = (r'\sffamily' if requested_family == 'sans-serif' else
+                   r'\ttfamily' if requested_family == 'monospace' else
+                   r'\rmfamily')
+        return preamble, fontcmd
+
+    @classmethod
+    def get_basefile(cls, tex, fontsize, dpi=None):
+        """
+        Return a filename based on a hash of the string, fontsize, and dpi.
+        """
+        src = cls._get_tex_source(tex, fontsize) + str(dpi)
+        filehash = hashlib.sha256(
+            src.encode('utf-8'),
+            usedforsecurity=False
+        ).hexdigest()
+        filepath = Path(cls._texcache)
+
+        num_letters, num_levels = 2, 2
+        for i in range(0, num_letters*num_levels, num_letters):
+            filepath = filepath / Path(filehash[i:i+2])
+
+        filepath.mkdir(parents=True, exist_ok=True)
+        return os.path.join(filepath, filehash)
+
+    @classmethod
+    def get_font_preamble(cls):
+        """
+        Return a string containing font configuration for the tex preamble.
+        """
+        font_preamble, command = cls._get_font_preamble_and_command()
+        return font_preamble
+
+    @classmethod
+    def get_custom_preamble(cls):
+        """Return a string containing user additions to the tex preamble."""
+        return mpl.rcParams['text.latex.preamble']
+
+    @classmethod
+    def _get_tex_source(cls, tex, fontsize):
+        """Return the complete TeX source for processing a TeX string."""
+        font_preamble, fontcmd = cls._get_font_preamble_and_command()
+        baselineskip = 1.25 * fontsize
+        return "\n".join([
+            r"\documentclass{article}",
+            r"% Pass-through \mathdefault, which is used in non-usetex mode",
+            r"% to use the default text font but was historically suppressed",
+            r"% in usetex mode.",
+            r"\newcommand{\mathdefault}[1]{#1}",
+            font_preamble,
+            r"\usepackage[utf8]{inputenc}",
+            r"\DeclareUnicodeCharacter{2212}{\ensuremath{-}}",
+            r"% geometry is loaded before the custom preamble as ",
+            r"% convert_psfrags relies on a custom preamble to change the ",
+            r"% geometry.",
+            r"\usepackage[papersize=72in, margin=1in]{geometry}",
+            cls.get_custom_preamble(),
+            r"% Use `underscore` package to take care of underscores in text.",
+            r"% The [strings] option allows to use underscores in file names.",
+            _usepackage_if_not_loaded("underscore", option="strings"),
+            r"% Custom packages (e.g. newtxtext) may already have loaded ",
+            r"% textcomp with different options.",
+            _usepackage_if_not_loaded("textcomp"),
+            r"\pagestyle{empty}",
+            r"\begin{document}",
+            r"% The empty hbox ensures that a page is printed even for empty",
+            r"% inputs, except when using psfrag which gets confused by it.",
+            r"% matplotlibbaselinemarker is used by dviread to detect the",
+            r"% last line's baseline.",
+            rf"\fontsize{{{fontsize}}}{{{baselineskip}}}%",
+            r"\ifdefined\psfrag\else\hbox{}\fi%",
+            rf"{{{fontcmd} {tex}}}%",
+            r"\end{document}",
+        ])
+
+    @classmethod
+    def make_tex(cls, tex, fontsize):
+        """
+        Generate a tex file to render the tex string at a specific font size.
+
+        Return the file name.
+        """
+        texfile = cls.get_basefile(tex, fontsize) + ".tex"
+        Path(texfile).write_text(cls._get_tex_source(tex, fontsize),
+                                 encoding='utf-8')
+        return texfile
+
+    @classmethod
+    def _run_checked_subprocess(cls, command, tex, *, cwd=None):
+        _log.debug(cbook._pformat_subprocess(command))
+        try:
+            report = subprocess.check_output(
+                command, cwd=cwd if cwd is not None else cls._texcache,
+                stderr=subprocess.STDOUT)
+        except FileNotFoundError as exc:
+            raise RuntimeError(
+                f'Failed to process string with tex because {command[0]} '
+                'could not be found') from exc
+        except subprocess.CalledProcessError as exc:
+            raise RuntimeError(
+                '{prog} was not able to process the following string:\n'
+                '{tex!r}\n\n'
+                'Here is the full command invocation and its output:\n\n'
+                '{format_command}\n\n'
+                '{exc}\n\n'.format(
+                    prog=command[0],
+                    format_command=cbook._pformat_subprocess(command),
+                    tex=tex.encode('unicode_escape'),
+                    exc=exc.output.decode('utf-8', 'backslashreplace'))
+                ) from None
+        _log.debug(report)
+        return report
+
+    @classmethod
+    def make_dvi(cls, tex, fontsize):
+        """
+        Generate a dvi file containing latex's layout of tex string.
+
+        Return the file name.
+        """
+        basefile = cls.get_basefile(tex, fontsize)
+        dvifile = '%s.dvi' % basefile
+        if not os.path.exists(dvifile):
+            texfile = Path(cls.make_tex(tex, fontsize))
+            # Generate the dvi in a temporary directory to avoid race
+            # conditions e.g. if multiple processes try to process the same tex
+            # string at the same time.  Having tmpdir be a subdirectory of the
+            # final output dir ensures that they are on the same filesystem,
+            # and thus replace() works atomically.  It also allows referring to
+            # the texfile with a relative path (for pathological MPLCONFIGDIRs,
+            # the absolute path may contain characters (e.g. ~) that TeX does
+            # not support; n.b. relative paths cannot traverse parents, or it
+            # will be blocked when `openin_any = p` in texmf.cnf).
+            cwd = Path(dvifile).parent
+            with TemporaryDirectory(dir=cwd) as tmpdir:
+                tmppath = Path(tmpdir)
+                cls._run_checked_subprocess(
+                    ["latex", "-interaction=nonstopmode", "--halt-on-error",
+                     f"--output-directory={tmppath.name}",
+                     f"{texfile.name}"], tex, cwd=cwd)
+                (tmppath / Path(dvifile).name).replace(dvifile)
+        return dvifile
+
+    @classmethod
+    def make_png(cls, tex, fontsize, dpi):
+        """
+        Generate a png file containing latex's rendering of tex string.
+
+        Return the file name.
+        """
+        basefile = cls.get_basefile(tex, fontsize, dpi)
+        pngfile = '%s.png' % basefile
+        # see get_rgba for a discussion of the background
+        if not os.path.exists(pngfile):
+            dvifile = cls.make_dvi(tex, fontsize)
+            cmd = ["dvipng", "-bg", "Transparent", "-D", str(dpi),
+                   "-T", "tight", "-o", pngfile, dvifile]
+            # When testing, disable FreeType rendering for reproducibility; but
+            # dvipng 1.16 has a bug (fixed in f3ff241) that breaks --freetype0
+            # mode, so for it we keep FreeType enabled; the image will be
+            # slightly off.
+            if (getattr(mpl, "_called_from_pytest", False) and
+                    mpl._get_executable_info("dvipng").raw_version != "1.16"):
+                cmd.insert(1, "--freetype0")
+            cls._run_checked_subprocess(cmd, tex)
+        return pngfile
+
+    @classmethod
+    def get_grey(cls, tex, fontsize=None, dpi=None):
+        """Return the alpha channel."""
+        if not fontsize:
+            fontsize = mpl.rcParams['font.size']
+        if not dpi:
+            dpi = mpl.rcParams['savefig.dpi']
+        key = cls._get_tex_source(tex, fontsize), dpi
+        alpha = cls._grey_arrayd.get(key)
+        if alpha is None:
+            pngfile = cls.make_png(tex, fontsize, dpi)
+            rgba = mpl.image.imread(os.path.join(cls._texcache, pngfile))
+            cls._grey_arrayd[key] = alpha = rgba[:, :, -1]
+        return alpha
+
+    @classmethod
+    def get_rgba(cls, tex, fontsize=None, dpi=None, rgb=(0, 0, 0)):
+        r"""
+        Return latex's rendering of the tex string as an RGBA array.
+
+        Examples
+        --------
+        >>> texmanager = TexManager()
+        >>> s = r"\TeX\ is $\displaystyle\sum_n\frac{-e^{i\pi}}{2^n}$!"
+        >>> Z = texmanager.get_rgba(s, fontsize=12, dpi=80, rgb=(1, 0, 0))
+        """
+        alpha = cls.get_grey(tex, fontsize, dpi)
+        rgba = np.empty((*alpha.shape, 4))
+        rgba[..., :3] = mpl.colors.to_rgb(rgb)
+        rgba[..., -1] = alpha
+        return rgba
+
+    @classmethod
+    def get_text_width_height_descent(cls, tex, fontsize, renderer=None):
+        """Return width, height and descent of the text."""
+        if tex.strip() == '':
+            return 0, 0, 0
+        dvifile = cls.make_dvi(tex, fontsize)
+        dpi_fraction = renderer.points_to_pixels(1.) if renderer else 1
+        with dviread.Dvi(dvifile, 72 * dpi_fraction) as dvi:
+            page, = dvi
+        # A total height (including the descent) needs to be returned.
+        return page.width, page.height + page.descent, page.descent

lib/python3.12/site-packages/sglang/srt/layers/attention/base_attn_backend.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+    from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+
+
+class AttentionBackend(ABC):
+    """The base class of attention backends"""
+
+    @abstractmethod
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        """Init the metadata for a forward pass."""
+        raise NotImplementedError()
+
+    def init_cuda_graph_state(self, max_bs: int):
+        """Init the global shared states for cuda graph."""
+        raise NotImplementedError()
+
+    def init_forward_metadata_capture_cuda_graph(
+        self,
+        bs: int,
+        num_tokens: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        """Init the metadata for a forward pass for capturing a cuda graph."""
+        raise NotImplementedError()
+
+    def init_forward_metadata_replay_cuda_graph(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        seq_lens_cpu: Optional[torch.Tensor],
+    ):
+        """Init the metadata for a forward pass for replaying a cuda graph."""
+        raise NotImplementedError()
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        """Get the fill value for padded seq lens. Typically, it is 0 or 1."""
+        raise NotImplementedError()
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+        **kwargs,
+    ):
+        """Run forward on an attention layer."""
+        if forward_batch.forward_mode.is_decode():
+            return self.forward_decode(
+                q,
+                k,
+                v,
+                layer,
+                forward_batch,
+                save_kv_cache=save_kv_cache,
+                **kwargs,
+            )
+        else:
+            return self.forward_extend(
+                q,
+                k,
+                v,
+                layer,
+                forward_batch,
+                save_kv_cache=save_kv_cache,
+                **kwargs,
+            )
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+    ):
+        """Run a forward for decode."""
+        raise NotImplementedError()
+
+    def forward_extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+    ):
+        """Run a forward for extend."""
+        raise NotImplementedError()

lib/python3.12/site-packages/sglang/srt/layers/attention/cutlass_mla_backend.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+"""
+Support attention backend for Cutlass MLA.
+
+"""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+import triton
+
+from sglang.global_config import global_config
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend
+from sglang.srt.layers.attention.utils import create_flashmla_kv_indices_triton
+from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.utils import is_cuda
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+    from sglang.srt.speculative.spec_info import SpecInfo
+
+_is_cuda = is_cuda()
+if _is_cuda:
+    from sgl_kernel import cutlass_mla_decode, cutlass_mla_get_workspace_size
+
+
+# Cutlass MLA only supports pagesize=128
+PAGE_SIZE = 128
+
+
+@dataclass
+class CutlassMLADecodeMetadata:
+    workspace: Optional[torch.Tensor] = None
+    block_kv_indices: Optional[torch.Tensor] = None
+
+    def __init__(
+        self,
+        workspace: Optional[torch.Tensor] = None,
+        block_kv_indices: Optional[torch.Tensor] = None,
+    ):
+        self.workspace = workspace
+        self.block_kv_indices = block_kv_indices
+
+
+class CutlassMLABackend(FlashInferMLAAttnBackend):
+    """Cutlass attention kernels."""
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        skip_prefill: bool = False,
+        kv_indptr_buf: Optional[torch.Tensor] = None,
+        kv_last_page_len_buf: Optional[torch.Tensor] = None,
+    ):
+        super().__init__(
+            model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf
+        )
+
+        self.num_q_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
+            get_attention_tp_size()
+        )
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.num_local_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.forward_metadata: Union[CutlassMLADecodeMetadata] = None
+        self.kv_lora_rank = model_runner.model_config.kv_lora_rank
+        self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
+        self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
+        self.v_head_dim = model_runner.model_config.v_head_dim
+        self.scaling = model_runner.model_config.scaling
+        self.data_type = model_runner.kv_cache_dtype
+        self.q_data_type = model_runner.dtype
+        self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+
+        bs = forward_batch.batch_size
+        spec_info = forward_batch.spec_info
+        if forward_batch.forward_mode.is_decode_or_idle():
+            if spec_info is None:
+                max_seqlen_pad = triton.cdiv(
+                    forward_batch.seq_lens_cpu.max().item(), PAGE_SIZE
+                )
+                block_kv_indices = torch.full(
+                    (bs, max_seqlen_pad),
+                    -1,
+                    dtype=torch.int32,
+                    device=forward_batch.seq_lens.device,
+                )
+                create_flashmla_kv_indices_triton[(bs,)](
+                    self.req_to_token,
+                    forward_batch.req_pool_indices,
+                    forward_batch.seq_lens,
+                    None,
+                    block_kv_indices,
+                    self.req_to_token.stride(0),
+                    max_seqlen_pad,
+                    PAGE_SIZE,
+                )
+                workspace_size = cutlass_mla_get_workspace_size(
+                    max_seqlen_pad * PAGE_SIZE, bs
+                )
+                workspace = torch.empty(
+                    workspace_size, device="cuda", dtype=torch.uint8
+                )
+                self.forward_metadata = CutlassMLADecodeMetadata(
+                    workspace,
+                    block_kv_indices,
+                )
+            else:
+                super().init_forward_metadata(forward_batch)
+        else:
+            super().init_forward_metadata(forward_batch)
+
+    def init_cuda_graph_state(
+        self,
+        max_bs: int,
+        block_kv_indices: Optional[torch.Tensor] = None,
+    ):
+        if block_kv_indices is None:
+            cuda_graph_kv_indices = torch.full(
+                (max_bs, (self.max_context_len + PAGE_SIZE) // PAGE_SIZE),
+                1,
+                dtype=torch.int32,
+                device="cuda",
+            )
+        else:
+            cuda_graph_kv_indices = block_kv_indices
+
+        workspace_size = cutlass_mla_get_workspace_size(
+            cuda_graph_kv_indices.shape[1] * PAGE_SIZE, max_bs
+        )
+        self.cuda_graph_mla_workspace = torch.empty(
+            workspace_size, device="cuda", dtype=torch.uint8
+        )
+        self.cuda_graph_kv_indices = cuda_graph_kv_indices
+
+    def init_forward_metadata_capture_cuda_graph(
+        self,
+        bs: int,
+        num_tokens: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[SpecInfo],
+    ):
+        if forward_mode.is_decode_or_idle():
+            if spec_info is None:
+                max_seqlen_pad = triton.cdiv(seq_lens.max().item(), PAGE_SIZE)
+
+                create_flashmla_kv_indices_triton[(bs,)](
+                    self.req_to_token,
+                    req_pool_indices,
+                    seq_lens,
+                    None,
+                    self.cuda_graph_kv_indices,
+                    self.req_to_token.stride(0),
+                    self.cuda_graph_kv_indices.stride(0),
+                    PAGE_SIZE,
+                )
+                workspace_size = cutlass_mla_get_workspace_size(
+                    max_seqlen_pad * PAGE_SIZE, bs
+                )
+                self.cuda_graph_mla_workspace = torch.empty(
+                    workspace_size, device="cuda", dtype=torch.uint8
+                )
+                self.forward_metadata = CutlassMLADecodeMetadata(
+                    self.cuda_graph_mla_workspace,
+                    self.cuda_graph_kv_indices[:bs, :max_seqlen_pad],
+                )
+        else:
+            super().init_forward_metadata_capture_cuda_graph(
+                bs,
+                num_tokens,
+                req_pool_indices,
+                seq_lens,
+                encoder_lens,
+                forward_mode,
+                spec_info,
+            )
+
+    def init_forward_metadata_replay_cuda_graph(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[SpecInfo],
+        seq_lens_cpu: Optional[torch.Tensor],
+    ):
+
+        if forward_mode.is_decode_or_idle():
+            assert seq_lens_cpu is not None
+            seq_lens = seq_lens[:bs]
+            seq_lens_cpu = seq_lens_cpu[:bs]
+            max_seqlen_pad = triton.cdiv(seq_lens_cpu.max().item(), PAGE_SIZE)
+            create_flashmla_kv_indices_triton[(bs,)](
+                self.req_to_token,
+                req_pool_indices[:bs],
+                seq_lens,
+                None,
+                self.cuda_graph_kv_indices,
+                self.req_to_token.stride(0),
+                self.cuda_graph_kv_indices.stride(0),
+                PAGE_SIZE,
+            )
+            workspace_size = cutlass_mla_get_workspace_size(
+                max_seqlen_pad * PAGE_SIZE, bs
+            )
+            self.cuda_graph_mla_workspace = torch.empty(
+                workspace_size, device="cuda", dtype=torch.uint8
+            )
+            self.forward_metadata.workspace = self.cuda_graph_mla_workspace
+            self.forward_metadata.block_kv_indices = self.cuda_graph_kv_indices[
+                :bs, :max_seqlen_pad
+            ]
+        else:
+            super().init_forward_metadata_replay_cuda_graph(
+                bs,
+                req_pool_indices,
+                seq_lens,
+                seq_lens_sum,
+                encoder_lens,
+                forward_mode,
+                spec_info,
+                seq_lens_cpu,
+            )
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        return 1
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+    ):
+        cache_loc = forward_batch.out_cache_loc
+
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer,
+                    cache_loc,
+                    k,
+                    v,
+                )
+        bs = forward_batch.batch_size
+        k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
+
+        reshape_q = q.view(-1, layer.tp_q_head_num, layer.head_dim)
+
+        o = cutlass_mla_decode(
+            q_nope_and_q_pe=reshape_q,
+            kv_c_and_k_pe_cache=k_cache.view(-1, PAGE_SIZE, self.kv_cache_dim),
+            seq_lens=forward_batch.seq_lens.to(torch.int32),
+            page_table=self.forward_metadata.block_kv_indices,
+            workspace=self.forward_metadata.workspace,
+        )
+
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)

lib/python3.12/site-packages/sglang/srt/layers/attention/double_sparsity_backend.py

@@ -0,0 +1,257 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.managers.schedule_batch import global_server_args_dict
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+
+class DoubleSparseAttnBackend(AttentionBackend):
+    def __init__(self, model_runner: ModelRunner):
+        # Lazy import to avoid the initialization of cuda context
+        from sglang.srt.layers.attention.triton_ops.double_sparsity_attention import (
+            extend_attention_fwd,
+            flash_decode_attention_fwd,
+            flash_decode_sparse_attention_fwd,
+        )
+
+        super().__init__()
+
+        self.decode_attention_fwd = flash_decode_attention_fwd
+        self.decode_sparse_attention_fwd = flash_decode_sparse_attention_fwd
+        self.extend_attention_fwd = extend_attention_fwd
+        self.num_head = model_runner.model_config.num_attention_heads
+        self.head_dim = model_runner.model_config.hidden_size // self.num_head
+        self.heavy_token_num = model_runner.server_args.ds_heavy_token_num
+
+        self.sorted_channels = model_runner.sorted_channels
+        self.sparse_decode_thresold = (
+            model_runner.server_args.ds_sparse_decode_threshold
+        )
+        self.att_out_approx: torch.Tensor = None
+        self.mid_out: torch.Tensor = None
+        self.mid_o_logexpsum: torch.Tensor = None
+
+        # TODO: Change the hard-coded block_seq_num
+        self.BLOCK_SEQ = 128
+
+        if global_server_args_dict.get("triton_attention_reduce_in_fp32", False):
+            self.reduce_dtype = torch.float32
+        else:
+            self.reduce_dtype = torch.float16
+
+        self.forward_metadata = None
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        """Init auxiliary variables for triton attention backend."""
+
+        if forward_batch.forward_mode.is_decode():
+            start_loc = torch.zeros_like(forward_batch.seq_lens, dtype=torch.int32)
+            start_loc[1:] = torch.cumsum(forward_batch.seq_lens[:-1], dim=0)
+
+            total_num_tokens = torch.sum(forward_batch.seq_lens).item()
+            attn_logits = torch.empty(
+                (self.num_head, total_num_tokens),
+                dtype=self.reduce_dtype,
+                device="cuda",
+            )
+
+            max_seq_len = torch.max(forward_batch.seq_lens).item()
+            min_seq_len = torch.min(forward_batch.seq_lens).item()
+            max_extend_len = None
+            # NOTE: Align sequence order with req_to_token order
+            ds_req_to_token = forward_batch.req_to_token_pool.req_to_token[
+                forward_batch.req_pool_indices
+            ]
+
+            bsz = forward_batch.seq_lens.shape[0]
+
+            att_out_approx = torch.empty(
+                [self.num_head, bsz, max_seq_len],
+                dtype=self.reduce_dtype,
+                device="cuda",
+            )
+
+            block_seq_num = (
+                self.heavy_token_num + self.BLOCK_SEQ - 1
+            ) // self.BLOCK_SEQ
+
+            mid_out = torch.empty(
+                [bsz, self.num_head, block_seq_num, self.head_dim],
+                dtype=torch.float32,
+                device="cuda",
+            )
+            mid_o_logexpsum = torch.empty(
+                [bsz, self.num_head, block_seq_num], dtype=torch.float32, device="cuda"
+            )
+            self.att_out_approx = att_out_approx
+            self.mid_out = mid_out
+            self.mid_o_logexpsum = mid_o_logexpsum
+
+        else:
+            start_loc = attn_logits = max_seq_len = min_seq_len = None
+            prefix_lens = forward_batch.extend_prefix_lens
+            max_extend_len = torch.max(forward_batch.seq_lens - prefix_lens).item()
+            ds_req_to_token = None
+
+        self.forward_metadata = (
+            start_loc,
+            attn_logits,
+            max_seq_len,
+            min_seq_len,
+            max_extend_len,
+            ds_req_to_token,
+        )
+
+    def forward_extend(
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+    ):
+        # TODO: reuse the buffer across layers
+        if layer.qk_head_dim != layer.v_head_dim:
+            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
+        else:
+            o = torch.empty_like(q)
+
+        k_label = torch.gather(
+            k,
+            2,
+            self.sorted_channels[layer.layer_id]
+            .unsqueeze(0)
+            .expand(k.shape[0], -1, -1),
+        )
+
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v, k_label
+            )
+
+        (
+            start_loc,
+            attn_logits,
+            max_seq_len,
+            min_seq_len,
+            max_extend_len,
+            ds_req_to_token,
+        ) = self.forward_metadata
+        self.extend_attention_fwd(
+            q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+            k.contiguous(),
+            v.contiguous(),
+            o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
+            forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
+            forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
+            forward_batch.req_to_token_pool.req_to_token,
+            forward_batch.req_pool_indices,
+            forward_batch.seq_lens,
+            forward_batch.extend_seq_lens,
+            forward_batch.extend_start_loc,
+            max_extend_len,
+            layer.scaling,
+            layer.logit_cap,
+        )
+        return o
+
+    def forward_decode(
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+    ):
+        # During torch.compile, there is a bug in rotary_emb that causes the
+        # output value to have a 3D tensor shape. This reshapes the output correctly.
+        q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
+
+        # TODO: reuse the buffer across layers
+        if layer.qk_head_dim != layer.v_head_dim:
+            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
+        else:
+            o = torch.empty_like(q)
+
+        # TODO: Add min seqlen
+        (
+            start_loc,
+            attn_logits,
+            max_seq_len,
+            min_seq_len,
+            max_extend_len,
+            ds_req_to_token,
+        ) = self.forward_metadata
+
+        k_label = torch.gather(
+            k,
+            2,
+            self.sorted_channels[layer.layer_id]
+            .unsqueeze(0)
+            .expand(k.shape[0], -1, -1),
+        )
+
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v, k_label
+            )
+
+        # NOTE(Andy) shouldn't be used when max_len_in_batch < heavy_token_num
+        #            and set a minimum value for sparse_decode
+        if (
+            min_seq_len < self.heavy_token_num
+            or max_seq_len < self.sparse_decode_thresold
+        ):
+            self.decode_attention_fwd(
+                q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+                forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
+                forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
+                o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
+                forward_batch.req_to_token_pool.req_to_token,
+                forward_batch.req_pool_indices,
+                start_loc,
+                forward_batch.seq_lens,
+                attn_logits,
+                max_seq_len,
+                layer.scaling,
+                layer.logit_cap,
+            )
+        else:
+            # TODO(Andy): indexing with torch.gather or torch.index_select or customized kernel
+            q_label = torch.gather(
+                q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+                2,
+                self.sorted_channels[layer.layer_id]
+                .unsqueeze(0)
+                .expand(q.shape[0], -1, -1),
+            )
+            self.decode_sparse_attention_fwd(
+                q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+                forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
+                forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
+                o.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+                q_label,
+                forward_batch.token_to_kv_pool.get_label_buffer(layer.layer_id),
+                ds_req_to_token,
+                forward_batch.seq_lens,
+                max_seq_len,
+                layer.scaling,
+                layer.logit_cap,
+                self.heavy_token_num,
+                self.att_out_approx,
+                self.mid_out,
+                self.mid_o_logexpsum,
+                self.BLOCK_SEQ,
+            )
+
+        return o

lib/python3.12/site-packages/sglang/srt/layers/attention/flashattention_backend.py

@@ -0,0 +1,1851 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional, Union
+
+import numpy as np
+import torch
+
+from sglang.srt.configs.model_config import AttentionArch
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.managers.schedule_batch import global_server_args_dict
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+from sgl_kernel import merge_state_v2
+from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
+
+
+@dataclass
+class FlashAttentionMetadata:
+    """Metadata to be init once in the model forward pass,
+    each layer's forward pass can reuse the metadata.
+
+    For each init metadata function, we will try set up them in below order
+    """
+
+    # Sequence lengths for the forward batch
+    cache_seqlens_int32: torch.Tensor = None
+    # Maximum sequence length for query
+    max_seq_len_q: int = 1
+    # Maximum sequence length for key
+    max_seq_len_k: int = 0
+    # Cumulative sequence lengths for query
+    cu_seqlens_q: torch.Tensor = None
+    # Cumulative sequence lengths for key
+    cu_seqlens_k: torch.Tensor = None
+    # Window size (typically used by Gemma)
+    window_size: tuple = (-1, -1)
+    # Page table, the index of KV Cache Tables/Blocks
+    page_table: torch.Tensor = None
+
+    # Encoder metadata
+    # Cumulative sequence lengths for encoder key
+    encoder_cu_seqlens_k: torch.Tensor = None
+    # Maximum sequence length for encoder key
+    encoder_max_seq_len_k: int = 0
+    # Sequence lengths for the forward batch
+    encoder_lens_int32: torch.Tensor = None
+    # Page table for the encoder
+    encoder_page_table: torch.Tensor = None
+
+    @dataclass
+    class LocalAttentionMetadata:
+        local_query_start_loc: torch.Tensor = None  # cu_seqlens_q for local attention
+        local_seqused_k: torch.Tensor = None  # sequence lengths for local attention
+        local_block_table: torch.Tensor = None  # block table for local attention
+        local_max_query_len: int = 0  # max query length for local attention
+        local_max_seq_len: int = 0  # max sequence length for local attention
+
+    local_attn_metadata: Optional[LocalAttentionMetadata] = None
+
+
+# Copied from:
+# https://github.com/houseroad/vllm/blob/4e45bfcaf928bdb9bd952b4ac922a3c205589ae8/vllm/v1/attention/backends/flash_attn.py
+#
+# Take in `query_start_loc_np` and `seq_lens_np` and break the sequences into
+# local attention blocks, where each block is passed to the attention kernel
+# as an independent local ("virtual") batch item.
+#
+# For example, if are performing a chunked prefill a batch of 3 sequences:
+#   q_seqlens  = [4, 10, 5]
+#   kv_seqlens = [6, 17, 9]
+# Then normally for regular attention we would compute with an attention mask
+#  for batch idx 0 (q_seqlens = 4, kv_seqlens = 6) like:
+#   batch idx: 0 (q_seqlens = 4, kv_seqlens = 6)
+#        k_toks >   0 1 2 3 4 5
+#        q_toks v  _____________
+#               0 | 1 1 1
+#               1 | 1 1 1 1
+#               2 | 1 1 1 1 1
+#               3 | 1 1 1 1 1 1
+#
+# for local attention (with attn_chunk_size = 4) we would compute with an
+#  attention mask like:
+#   batch idx: 0  (q_seqlens = 4, kv_seqlens = 6, attn_chunk_size = 4)
+#        k_toks >   0 1 2 3 4 5
+#        q_toks v  _____________
+#               0 | 1 1 1
+#               1 | 1 1 1 1
+#               2 |         1
+#               3 |         1 1
+#
+# We can simulate this mask using standard flash-attention by breaking the
+#  sequences into local ("virtual") batches, where each local batch item is a
+#  local attention block, so in this case batch idx 0 would be broken up into:
+#
+#   local-batch idx: 0 (q_seqlens = 2, kv_seqlens = 4)  (batch 0)
+#        k_toks >   0 1 2 3
+#        q_toks v  _____________
+#               0 | 1 1 1
+#               1 | 1 1 1 1
+#   local-batch idx: 1 (q_seqlens = 2, kv_seqlens = 2) (batch 0)
+#        k_toks >   4 5
+#        q_toks v  _____________
+#               2 | 1
+#               3 | 1 1
+#
+# e.g. if we have:
+#   attn_chunk_size = 4
+#   query_start_loc_np = [0, 4, 14, 19] (q_seqlens = [4, 10, 5])
+# Then this function would return:
+#                           __b0__  ______b1______  __b2__ < orig batch indices
+#   q_seqlens_local    = [   2,  2,  1,  4,  4,  1,  4,  1]
+#   cu_seqlens_q_local = [0, 4,  6, 10, 14, 18, 19, 23, 24]
+#   seqlens_k_local    = [   4,  2,  4,  4,  4,  1,  4,  1]
+#   block_table_local  : shape[local_virtual_batches, pages_per_local_batch]
+def make_local_attention_virtual_batches(
+    attn_chunk_size: int,
+    query_start_loc_np: np.ndarray,
+    seq_lens_np: np.ndarray,
+    block_table: torch.Tensor,
+    page_size: int = 0,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, torch.Tensor]:
+    """
+    Take in `query_start_loc_np` and `seq_lens_np` and break the sequences into
+    local attention blocks, where each block is passed to the attention kernel
+    as an independent local ("virtual") batch item.
+
+    Args:
+        attn_chunk_size: Size of local attention chunks
+        query_start_loc_np: Cumulative sum of query lengths (numpy array)
+        seq_lens_np: Sequence lengths (numpy array)
+        block_table: Block table for KV cache
+        page_size: Size of each page in the KV cache
+
+    Returns:
+        seqlens_q_local: Query sequence lengths for local attention
+        cu_seqlens_q_local: Cumulative sum of query sequence lengths for local attention
+        seqlens_k_local: Key sequence lengths for local attention
+        block_table_local: Block table for local attention
+    """
+    # Adjust attention_chunk_size based on the actual sequence length
+    # to avoid index out of bounds errors
+    max_seq_len = seq_lens_np.max()
+    effective_chunk_size = min(attn_chunk_size, max_seq_len)
+    # Make sure effective_chunk_size is divisible by page_size
+    effective_chunk_size = (effective_chunk_size // page_size) * page_size
+    if effective_chunk_size < page_size:
+        effective_chunk_size = page_size
+    attn_chunk_size = effective_chunk_size
+
+    q_seqlens = query_start_loc_np[1:] - query_start_loc_np[:-1]
+    actual_batch_size = seq_lens_np.shape[0]
+
+    # Handle if we are starting in the middle of a local attention block,
+    #  we assume q_seqlens > 0 (for all elements), for each batch idx we compute
+    #  the number of tokens that are not in the first local attention block and
+    #  then we can simply use a cdiv for the rest.
+    # For example if we have:
+    #   attn_chunk_size = 4
+    #   q_seqlens = [4, 10, 5]
+    #   k_seqlens = [6, 17, 9]
+    # Then we would get:
+    #   new_tokens_in_first_block = [2, 1, 4]
+    #   local_blocks = [2, 4, 2]
+    q_tokens_in_first_block = np.minimum(
+        attn_chunk_size - ((seq_lens_np - q_seqlens) % attn_chunk_size), q_seqlens
+    ).astype(np.int32)
+    tokens_in_last_block = attn_chunk_size + (seq_lens_np % -attn_chunk_size)
+    local_blocks = 1 + cdiv(q_seqlens - q_tokens_in_first_block, attn_chunk_size)
+
+    # Once we know the number of local blocks we can compute the request spans
+    #  for each batch idx, we can figure out the number of "virtual" requests we
+    #  have to make,
+    # For the above example we would get:
+    #   seqlens_q_local = [2, 2, 1, 4, 4, 1, 4, 1]
+    #
+    # First Get batched arange. (E.g., [2, 4, 2] -> [0, 1, 0, 1, 2, 3, 0, 1])
+    #   (TODO: max a utility to share this code with _prepare_inputs)
+    # arange step 1. [2, 4, 2] -> [2, 6, 8]
+    cu_num_blocks = np.cumsum(local_blocks)
+    virtual_batches = cu_num_blocks[-1]
+    # arange step 2. [2, 6, 8] -> [0, 0, 2, 2, 2, 2, 6, 6]
+    block_offsets = np.repeat(cu_num_blocks - local_blocks, local_blocks)
+    # arange step 3. [0, 1, 0, 1, 2, 3, 0, 1]
+    arange = np.arange(virtual_batches, dtype=np.int32) - block_offsets
+    # also compute reverse arange (i.e. [1, 0, 3, 2, 1, 0, 1, 0])
+    rarange = np.repeat(local_blocks, local_blocks) - arange - 1
+    # Then we can compute the seqlens_q_local, handling the fact that the
+    #  first and last blocks could be partial
+    seqlens_q_local = np.repeat(q_seqlens - q_tokens_in_first_block, local_blocks)
+    # set the first block since this may be a partial block
+    seqlens_q_local[arange == 0] = q_tokens_in_first_block
+    # set the remaining blocks
+    seqlens_q_local[arange > 0] = np.minimum(
+        seqlens_q_local - attn_chunk_size * (arange - 1), attn_chunk_size
+    )[arange > 0]
+
+    # convert from q_seqlens to cu_seqlens_q
+    cu_seqlens_q_local = np.pad(np.cumsum(seqlens_q_local), (1, 0)).astype(np.int32)
+
+    # compute the seqlens_k_local,
+    #  basically a full local attention block for all but the last block in each
+    #  batch
+    # For our example this will be:
+    #   seqlens_k_local = [4, 2, 4, 4, 4, 1, 4, 1]
+    seqlens_k_local = np.full(cu_num_blocks[-1], attn_chunk_size, dtype=np.int32)
+    seqlens_k_local[cu_num_blocks - 1] = tokens_in_last_block
+
+    k_seqstarts_absolute = np.repeat(seq_lens_np, local_blocks) - (
+        rarange * attn_chunk_size + np.repeat(tokens_in_last_block, local_blocks)
+    )
+    # For the example the local attention blocks start at:
+    #                           _b0_  _____b1_____  _b2_
+    #   k_seqstarts_absolute = [0, 4, 4, 8, 12, 16, 4, 8]
+    block_starts = k_seqstarts_absolute // page_size
+
+    assert attn_chunk_size % page_size == 0, (
+        f"attn_chunk_size {attn_chunk_size} is not "
+        f"divisible by page_size {page_size}"
+    )
+    pages_per_local_batch = attn_chunk_size // page_size
+
+    # Create a block_table for the local attention blocks
+    # For out example if we have a block-table like (assuming page_size=2):
+    #   block_table = [
+    #     [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],  < batch 0
+    #     [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],  < batch 1
+    #     [20, 21, 22, 23, 24, 25, 26, 27, 28, 29],  < batch 2
+    #   ]
+    # Then for the local batches we would want a block-table like
+    #   block_table_local = [
+    #     [  0,  1 ], < local-batch 0, (batch 0, starting from k[0])
+    #     [  2,  3 ], < local-batch 1, (batch 0, starting from k[4])
+    #     [ 12, 13 ], < local-batch 2, (batch 1, starting from k[4])
+    #     [ 14, 15 ], < local-batch 3, (batch 1, starting from k[8])
+    #     [ 16, 17 ], < local-batch 4, (batch 1, starting from k[12])
+    #     [ 18, 19 ], < local-batch 5, (batch 1, starting from k[16])
+    #     [ 22, 23 ], < local-batch 6, (batch 2, starting from k[4])
+    #     [ 24, 25 ], < local-batch 7, (batch 2, starting from k[8])
+    #   ]
+    block_indices = np.broadcast_to(
+        np.arange(pages_per_local_batch, dtype=np.int32),
+        (virtual_batches, pages_per_local_batch),
+    ) + np.expand_dims(block_starts, axis=1)
+    # Ensure block_indices doesn't exceed block_table dimensions
+    # This is a critical safety check that prevents index out of bounds errors
+    # when dealing with large sequences (>8192 tokens) or when the block_table
+    # dimensions are smaller than what would be needed for the full attention chunk size.
+    block_indices = block_indices.flatten().clip(max=block_table.shape[1] - 1)
+    batch_indices = np.repeat(
+        np.arange(actual_batch_size, dtype=np.int32),
+        local_blocks * pages_per_local_batch,
+    )
+    block_table_local = block_table[batch_indices, block_indices].view(
+        virtual_batches, -1
+    )
+
+    return seqlens_q_local, cu_seqlens_q_local, seqlens_k_local, block_table_local
+
+
+def cdiv(a: int, b: int) -> int:
+    """Ceiling division."""
+    return -(a // -b)
+
+
+# TODO(hebiao064): remove this once we have a better way to handle the merge_state_v2 torch.compile issue
+@torch._dynamo.disable()
+def merge_state_v2_wrapper(o, s_a, o_exp, s_b):
+    return merge_state_v2(o, s_a, o_exp, s_b)
+
+
+class FlashAttentionBackend(AttentionBackend):
+    """FlashAttention backend implementation.
+
+    Note about the init:
+    - If no spec decoding
+        - FlashAttentionBackend will be init once when the server starts.
+    - If spec decoding
+        - FlashAttentionBackend will be init once for the target worker
+        - FlashAttentionMultiStepBackend will be once for the draft worker
+            - It will spawn num_steps FlashAttentionBackend for the draft worker
+
+    Note about CUDA Graph:
+    - We only support CUDA Graph for Decode (Normal Decode and Draft Decode) and Target Verify.
+    - We don't support CUDA Graph for Extend and Draft Extend.
+    - When server init, init_cuda_graph_state will be called first and then init_cuda_graph_capture will be called.
+    - For each forward batch, init_replay_cuda_graph will be called first and then replay the graph.
+    """
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        skip_prefill: bool = False,
+        speculative_step_id=0,
+        topk=0,
+        speculative_num_steps=0,
+    ):
+        super().__init__()
+
+        assert not (
+            model_runner.sliding_window_size is not None
+            and model_runner.model_config.is_encoder_decoder
+        ), "Sliding window and cross attention are not supported together"
+
+        self.forward_metadata: FlashAttentionMetadata = None
+        # extra metdata for handling speculative decoding topk > 1, extended draft decode and verify
+        self.forward_metadata_spec_decode_expand: FlashAttentionMetadata = None
+        self.max_context_len = model_runner.model_config.context_len
+        self.device = model_runner.device
+        self.decode_cuda_graph_metadata = {}
+        self.target_verify_metadata = {}
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.kv_cache_dtype = model_runner.kv_cache_dtype
+        self.kv_cache_dtype_str = model_runner.server_args.kv_cache_dtype
+        self.page_size = model_runner.page_size
+        self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
+        self.skip_prefill = skip_prefill
+        self.topk = model_runner.server_args.speculative_eagle_topk or 0
+        self.speculative_num_steps = speculative_num_steps
+        self.speculative_num_draft_tokens = (
+            model_runner.server_args.speculative_num_draft_tokens
+        )
+        self.speculative_step_id = speculative_step_id
+
+        # Local attention settings
+        self.attention_chunk_size = (
+            model_runner.attention_chunk_size
+            if hasattr(model_runner, "attention_chunk_size")
+            else None
+        )
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        """Initialize forward metadata hence all layers in the forward pass can reuse it."""
+        metadata = FlashAttentionMetadata()
+        seqlens_in_batch = forward_batch.seq_lens
+        batch_size = len(seqlens_in_batch)
+        device = seqlens_in_batch.device
+
+        if forward_batch.forward_mode.is_decode_or_idle():
+            # Draft Decode
+            if forward_batch.spec_info is not None:
+                if self.topk <= 1:
+                    metadata.cache_seqlens_int32 = (
+                        seqlens_in_batch + (self.speculative_step_id + 1)
+                    ).to(torch.int32)
+                    metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + (
+                        self.speculative_step_id + 1
+                    )
+                    metadata.cu_seqlens_q = torch.arange(
+                        0, batch_size + 1, dtype=torch.int32, device=device
+                    )
+                    metadata.cu_seqlens_k = torch.nn.functional.pad(
+                        torch.cumsum(
+                            metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+                        ),
+                        (1, 0),
+                    )
+                    metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
+                        forward_batch.req_pool_indices, : metadata.max_seq_len_k
+                    ]
+                else:
+                    metadata.cache_seqlens_int32 = (seqlens_in_batch).to(torch.int32)
+                    metadata.max_seq_len_q = self.topk
+                    metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item()
+                    metadata.cu_seqlens_q = torch.arange(
+                        0,
+                        batch_size * self.topk + 1,
+                        step=self.topk,
+                        dtype=torch.int32,
+                        device=device,
+                    )
+                    metadata.cu_seqlens_k = torch.nn.functional.pad(
+                        torch.cumsum(
+                            metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+                        ),
+                        (1, 0),
+                    )
+                    metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
+                        forward_batch.req_pool_indices, : metadata.max_seq_len_k
+                    ]
+
+                    metadata_expand = FlashAttentionMetadata()
+                    decode_length = self.speculative_step_id + 1
+                    metadata_expand.cache_seqlens_int32 = torch.full(
+                        (seqlens_in_batch.numel() * self.topk,),
+                        decode_length,
+                        device=device,
+                        dtype=torch.int32,
+                    )
+                    metadata_expand.max_seq_len_q = 1
+                    metadata_expand.max_seq_len_k = self.speculative_step_id + 1
+                    metadata_expand.cu_seqlens_q = torch.arange(
+                        0,
+                        metadata_expand.cache_seqlens_int32.numel() + 1,
+                        dtype=torch.int32,
+                        device=device,
+                    )
+                    metadata_expand.cu_seqlens_k = torch.arange(
+                        0,
+                        metadata_expand.cache_seqlens_int32.numel() * decode_length + 1,
+                        step=decode_length,
+                        dtype=torch.int32,
+                        device=device,
+                    )
+                    cache_loc = forward_batch.out_cache_loc.view(
+                        self.speculative_num_steps, -1
+                    ).T.contiguous()
+                    metadata_expand.page_table = (
+                        cache_loc[:, :decode_length].contiguous().to(torch.int32)
+                    )
+                    self.forward_metadata_spec_decode_expand = metadata_expand
+            else:
+                # Normal Decode
+                metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32)
+                metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item()
+                metadata.cu_seqlens_q = torch.arange(
+                    0, batch_size + 1, dtype=torch.int32, device=device
+                )
+                metadata.cu_seqlens_k = torch.nn.functional.pad(
+                    torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)
+                )
+                metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
+                    forward_batch.req_pool_indices, : metadata.max_seq_len_k
+                ]
+            # TODO: we need to test this part for llama 4 eagle case
+            self._init_local_attn_metadata(metadata, device)
+        elif forward_batch.forward_mode.is_target_verify():
+            if self.topk <= 1:
+                metadata.cache_seqlens_int32 = (
+                    forward_batch.seq_lens + self.speculative_num_draft_tokens
+                ).to(torch.int32)
+                metadata.max_seq_len_q = self.speculative_num_draft_tokens
+                metadata.max_seq_len_k = (
+                    forward_batch.seq_lens_cpu.max().item()
+                    + self.speculative_num_draft_tokens
+                )
+                metadata.cu_seqlens_q = torch.arange(
+                    0,
+                    batch_size * self.speculative_num_draft_tokens + 1,
+                    self.speculative_num_draft_tokens,
+                    dtype=torch.int32,
+                    device=device,
+                )
+                metadata.cu_seqlens_k = torch.nn.functional.pad(
+                    torch.cumsum(
+                        metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+                    ),
+                    (1, 0),
+                )
+                metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
+                    forward_batch.req_pool_indices, : metadata.max_seq_len_k
+                ]
+
+                self._init_local_attn_metadata(metadata, device)
+            else:
+                metadata.cache_seqlens_int32 = forward_batch.seq_lens.to(torch.int32)
+                metadata.max_seq_len_q = self.speculative_num_draft_tokens
+                metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item()
+                metadata.cu_seqlens_q = torch.arange(
+                    0,
+                    batch_size * self.speculative_num_draft_tokens + 1,
+                    step=self.speculative_num_draft_tokens,
+                    dtype=torch.int32,
+                    device=device,
+                )
+                metadata.cu_seqlens_k = torch.nn.functional.pad(
+                    torch.cumsum(
+                        metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+                    ),
+                    (1, 0),
+                )
+                metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
+                    forward_batch.req_pool_indices, : metadata.max_seq_len_k
+                ]
+
+                metadata_expand = FlashAttentionMetadata()
+
+                metadata_expand.max_seq_len_q = 1
+                metadata_expand.cu_seqlens_q = torch.arange(
+                    0,
+                    forward_batch.seq_lens.numel() * self.speculative_num_draft_tokens
+                    + 1,
+                    dtype=torch.int32,
+                    device=device,
+                )
+
+                # create expand page table
+                offsets = torch.arange(
+                    self.speculative_num_draft_tokens, device=device
+                ).unsqueeze(
+                    0
+                )  # shape: (1, self.speculative_num_draft_tokens)
+                cols = offsets.expand(
+                    forward_batch.seq_lens.numel(), -1
+                ) + forward_batch.seq_lens.unsqueeze(1)
+                cum_len = torch.nn.functional.pad(
+                    torch.cumsum(
+                        (
+                            forward_batch.seq_lens + self.speculative_num_draft_tokens
+                        ).repeat_interleave(self.speculative_num_draft_tokens),
+                        dim=0,
+                    ),
+                    (1, 0),
+                )[:-1]
+                mask_extraction_indices = (
+                    cols.repeat_interleave(self.speculative_num_draft_tokens, dim=0)
+                    + cum_len[:, None]
+                ).view(1, -1)
+                mask = forward_batch.spec_info.custom_mask[
+                    mask_extraction_indices
+                ].view(
+                    -1, self.speculative_num_draft_tokens
+                )  # (bsz * draft_num, draft_num)
+
+                # shift table indices to avoid padding
+                # non_masked_page_table [[8, 9, 10],   mask (display with int format) [[1, 0, 0],
+                #                        [8, 9, 10],                                   [1, 1, 0],
+                #                        [8, 9, 10]]                                   [1, 0, 1]]
+                # if masked with padding [[8, 0, 0],   our mask without padding       [[8, 9, 10],
+                #                        [8, 9, 0],                                    [8, 9, 10],
+                #                        [8, 0, 10]]                                   [8, 10, 9]]
+                # note here cache_seqlens_int32 is [1, 2, 2] so extra page indices will be ignored in each row
+                col_indices = offsets.expand(
+                    mask.shape[0], self.speculative_num_draft_tokens
+                )
+                # Build keys: if an entry is valid (mask==True), keep its original index;
+                # if not, add self.speculative_num_draft_tokens so that it sorts after all valid entries.
+                keys = torch.where(
+                    mask, col_indices, col_indices + self.speculative_num_draft_tokens
+                )
+                _, sort_order = torch.sort(keys, dim=1)
+                non_masked_page_table = (
+                    forward_batch.req_to_token_pool.req_to_token[
+                        forward_batch.req_pool_indices, :
+                    ]
+                    .gather(1, cols)
+                    .repeat_interleave(self.speculative_num_draft_tokens, dim=0)
+                )  # (bsz, draft_num)
+                metadata_expand.page_table = non_masked_page_table.gather(1, sort_order)
+                metadata_expand.cache_seqlens_int32 = mask.sum(dim=1).to(torch.int32)
+                metadata_expand.cu_seqlens_k = torch.nn.functional.pad(
+                    torch.cumsum(
+                        metadata_expand.cache_seqlens_int32, dim=0, dtype=torch.int32
+                    ),
+                    (1, 0),
+                )
+                metadata_expand.max_seq_len_k = (
+                    metadata_expand.cache_seqlens_int32.max().item()
+                )
+                self.forward_metadata_spec_decode_expand = metadata_expand
+        elif forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed():
+            metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32)
+            metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item()
+            metadata.cu_seqlens_k = torch.nn.functional.pad(
+                torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)
+            )
+            metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
+                forward_batch.req_pool_indices, : metadata.max_seq_len_k
+            ]
+
+            if (
+                any(forward_batch.extend_prefix_lens_cpu)
+                or forward_batch.forward_mode == ForwardMode.DRAFT_EXTEND
+            ):
+                extend_seq_lens = forward_batch.extend_seq_lens
+                metadata.max_seq_len_q = max(forward_batch.extend_seq_lens_cpu)
+                metadata.cu_seqlens_q = torch.nn.functional.pad(
+                    torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), (1, 0)
+                )
+            else:
+                metadata.max_seq_len_q = metadata.max_seq_len_k
+                metadata.cu_seqlens_q = metadata.cu_seqlens_k
+
+            # Setup local attention if enabled
+            if forward_batch.forward_mode == ForwardMode.EXTEND:
+                self._init_local_attn_metadata(metadata, device)
+
+        # Encoder metadata for cross attention
+        if forward_batch.encoder_lens is not None:
+            assert (
+                forward_batch.encoder_lens.numel() == 1
+            ), "Only encoder size 1 is supported for now"
+
+            metadata.encoder_lens_int32 = forward_batch.encoder_lens.to(torch.int32)
+            metadata.encoder_cu_seqlens_k = torch.nn.functional.pad(
+                torch.cumsum(metadata.encoder_lens_int32, dim=0, dtype=torch.int32),
+                (1, 0),
+            )
+            metadata.encoder_max_seq_len_k = metadata.encoder_lens_int32.max().item()
+            metadata.encoder_page_table = forward_batch.req_to_token_pool.req_to_token[
+                forward_batch.req_pool_indices, : metadata.encoder_max_seq_len_k
+            ]
+
+            # Currently only support forward_batch.encoder_lens.numel() == 1
+            metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
+                forward_batch.req_pool_indices,
+                metadata.encoder_max_seq_len_k : (
+                    metadata.encoder_max_seq_len_k + metadata.max_seq_len_k
+                ),
+            ]
+
+        # Convert the page table to a strided format which is needed by FA3 API
+        if self.page_size > 1:
+            self.strided_indices = torch.arange(
+                0, metadata.page_table.shape[1], self.page_size, device=self.device
+            )
+            metadata.page_table = (
+                metadata.page_table[:, self.strided_indices] // self.page_size
+            )
+
+        self.forward_metadata = metadata
+
+    def forward_extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+        # For multi-head latent attention
+        q_rope: Optional[torch.Tensor] = None,
+        k_rope: Optional[torch.Tensor] = None,
+    ):
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                cache_loc = (
+                    forward_batch.out_cache_loc
+                    if not layer.is_cross_attention
+                    else forward_batch.encoder_out_cache_loc
+                )
+                if not self.use_mla:
+                    forward_batch.token_to_kv_pool.set_kv_buffer(
+                        layer, cache_loc, k, v, layer.k_scale, layer.v_scale
+                    )
+                else:
+                    forward_batch.token_to_kv_pool.set_mla_kv_buffer(
+                        layer,
+                        cache_loc,
+                        k,
+                        k_rope,
+                    )
+
+        # Use precomputed metadata across all layers
+        metadata = self.forward_metadata
+
+        # Calculate window size (can be moved to metadata if layer properties don't change)
+        # we don't do layer.sliding_window_size - 1 since in model.get_attention_sliding_window_size() we already - 1
+        # here is two side inclusive
+        window_size = (
+            (layer.sliding_window_size, 0)
+            if layer.sliding_window_size is not None and layer.sliding_window_size > -1
+            else (-1, -1)
+        )
+        k_descale, v_descale = None, None
+        # only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention
+        # has corresponding quantization method so that layer.k_scale is not None
+        if self.kv_cache_dtype_str != "auto" and layer.k_scale is not None:
+            descale_shape = (forward_batch.batch_size, layer.tp_k_head_num)
+            k_descale = layer.k_scale.expand(descale_shape)
+            v_descale = layer.v_scale.expand(descale_shape)
+            q = q.to(self.kv_cache_dtype)
+        causal = not layer.is_cross_attention
+
+        # Check if we should use local attention
+        use_local_attn = (
+            self.attention_chunk_size is not None
+            and metadata.local_attn_metadata is not None
+            and (hasattr(layer, "use_irope") and layer.use_irope)
+        )
+
+        # We do cascade attention for Target Verify with topk > 1
+        use_cascade_attn = (
+            forward_batch.forward_mode.is_target_verify() and self.topk > 1
+        )
+
+        # Get the appropriate page table based on whether we're using local attention
+        if use_local_attn:
+            local_metadata = metadata.local_attn_metadata
+            page_table = local_metadata.local_block_table
+            cu_seqlens_q = local_metadata.local_query_start_loc
+            cache_seqlens = local_metadata.local_seqused_k
+            max_seqlen_q = local_metadata.local_max_query_len
+            max_seqlen_k = local_metadata.local_max_seq_len
+        else:
+            page_table = metadata.page_table
+            cu_seqlens_q = metadata.cu_seqlens_q
+            cache_seqlens = metadata.cache_seqlens_int32
+            max_seqlen_q = metadata.max_seq_len_q
+            max_seqlen_k = metadata.max_seq_len_k
+            cu_seqlens_k = metadata.cu_seqlens_k
+
+        # Use Flash Attention for prefill
+        if not self.use_mla:
+            # Do multi-head attention
+            key_cache, value_cache = forward_batch.token_to_kv_pool.get_kv_buffer(
+                layer.layer_id
+            )
+            key_cache = key_cache.view(
+                -1, self.page_size, layer.tp_k_head_num, layer.head_dim
+            )
+            value_cache = value_cache.view(
+                -1, self.page_size, layer.tp_v_head_num, layer.head_dim
+            )
+            if layer.is_cross_attention:
+                page_table = metadata.encoder_page_table
+                cache_seqlens = metadata.encoder_lens_int32
+                cu_seqlens_k = metadata.encoder_cu_seqlens_k
+                window_size = (-1, -1)
+
+            result = flash_attn_with_kvcache(
+                q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                k_cache=key_cache,
+                v_cache=value_cache,
+                page_table=page_table,
+                cache_seqlens=cache_seqlens,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
+                max_seqlen_q=max_seqlen_q,
+                softmax_scale=layer.scaling,
+                causal=False if use_cascade_attn else causal,
+                window_size=window_size,
+                softcap=layer.logit_cap,
+                k_descale=k_descale,
+                v_descale=v_descale,
+                return_softmax_lse=use_cascade_attn,
+            )
+
+            if use_cascade_attn:
+                o, softmax_lse, *rest = result
+                o_expand, softmax_lse_expand, *rest_expand = flash_attn_with_kvcache(
+                    q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                    k_cache=key_cache,
+                    v_cache=value_cache,
+                    page_table=self.forward_metadata_spec_decode_expand.page_table,
+                    cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
+                    cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
+                    cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k,
+                    max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
+                    softmax_scale=layer.scaling,
+                    causal=False,
+                    window_size=window_size,
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                    return_softmax_lse=True,
+                )
+                o, _ = merge_state_v2_wrapper(
+                    o,
+                    softmax_lse.T.contiguous(),
+                    o_expand,
+                    softmax_lse_expand.T.contiguous(),
+                )
+            else:
+                o = result
+        else:
+            if (
+                not global_server_args_dict["disable_chunked_prefix_cache"]
+                and forward_batch.attn_attend_prefix_cache is not None
+                and not forward_batch.forward_mode.is_target_verify()
+                and not forward_batch.forward_mode.is_draft_extend()
+            ):
+                # Do multi-head attention with chunked prefix cache
+
+                if forward_batch.attn_attend_prefix_cache:
+                    # MHA for chunked prefix kv cache when running model with MLA
+                    assert forward_batch.prefix_chunk_idx is not None
+                    assert forward_batch.prefix_chunk_cu_seq_lens is not None
+                    assert forward_batch.prefix_chunk_max_seq_lens is not None
+
+                    chunk_idx = forward_batch.prefix_chunk_idx
+                    assert chunk_idx >= 0
+
+                    output, lse, *rest = flash_attn_varlen_func(
+                        q=q.view(-1, layer.tp_q_head_num, layer.head_dim),
+                        k=k.view(-1, layer.tp_k_head_num, layer.head_dim),
+                        v=v.view(-1, layer.tp_k_head_num, layer.v_head_dim),
+                        cu_seqlens_q=metadata.cu_seqlens_q,
+                        cu_seqlens_k=forward_batch.prefix_chunk_cu_seq_lens[chunk_idx],
+                        max_seqlen_q=metadata.max_seq_len_q,
+                        max_seqlen_k=forward_batch.prefix_chunk_max_seq_lens[chunk_idx],
+                        softmax_scale=layer.scaling,
+                        causal=False,
+                        return_softmax_lse=True,
+                    )
+                else:
+                    # MHA for extend part of sequence without attending prefix kv cache
+                    output, lse, *rest = flash_attn_varlen_func(
+                        q=q.view(-1, layer.tp_q_head_num, layer.head_dim),
+                        k=k.view(-1, layer.tp_k_head_num, layer.head_dim),
+                        v=v.view(-1, layer.tp_k_head_num, layer.v_head_dim),
+                        cu_seqlens_q=metadata.cu_seqlens_q,
+                        cu_seqlens_k=metadata.cu_seqlens_q,
+                        max_seqlen_q=metadata.max_seq_len_q,
+                        max_seqlen_k=metadata.max_seq_len_q,
+                        softmax_scale=layer.scaling,
+                        causal=True,
+                        return_softmax_lse=True,
+                    )
+                return output, lse
+            else:
+                # Do absorbed multi-latent attention
+                kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
+                k_rope = kv_cache[:, :, layer.v_head_dim :]
+                c_kv = kv_cache[:, :, : layer.v_head_dim]
+                k_rope_cache = k_rope.view(
+                    -1,
+                    self.page_size,
+                    layer.tp_k_head_num,
+                    layer.head_dim - layer.v_head_dim,
+                )
+                c_kv_cache = c_kv.view(
+                    -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
+                )
+                if q_rope is not None:
+                    q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+                    q_rope = q_rope.view(
+                        -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+                    )
+                else:
+                    q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
+                    q_nope = q_all[:, :, : layer.v_head_dim]
+                    q_rope = q_all[:, :, layer.v_head_dim :]
+
+                result = flash_attn_with_kvcache(
+                    q=q_rope,
+                    k_cache=k_rope_cache,
+                    v_cache=c_kv_cache,
+                    qv=q_nope,
+                    page_table=page_table,
+                    cache_seqlens=cache_seqlens,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
+                    max_seqlen_q=max_seqlen_q,
+                    softmax_scale=layer.scaling,
+                    causal=False if use_cascade_attn else causal,
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                    return_softmax_lse=use_cascade_attn,
+                )
+                if use_cascade_attn:
+                    o, softmax_lse, *rest = result
+                    o_expand, softmax_lse_expand, *rest_expand = (
+                        flash_attn_with_kvcache(
+                            q=q_rope,
+                            k_cache=k_rope_cache,
+                            v_cache=c_kv_cache,
+                            qv=q_nope,
+                            page_table=self.forward_metadata_spec_decode_expand.page_table,
+                            cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
+                            cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
+                            cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k,
+                            max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
+                            softmax_scale=layer.scaling,
+                            causal=False,
+                            window_size=window_size,
+                            softcap=layer.logit_cap,
+                            k_descale=k_descale,
+                            v_descale=v_descale,
+                            return_softmax_lse=True,
+                        )
+                    )
+                    o, _ = merge_state_v2_wrapper(
+                        o,
+                        softmax_lse.T.contiguous(),
+                        o_expand,
+                        softmax_lse_expand.T.contiguous(),
+                    )
+                else:
+                    o = result
+
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+        # For multi-head latent attention
+        q_rope: Optional[torch.Tensor] = None,
+        k_rope: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                cache_loc = (
+                    forward_batch.out_cache_loc
+                    if not layer.is_cross_attention
+                    else forward_batch.encoder_out_cache_loc
+                )
+                if not self.use_mla:
+                    forward_batch.token_to_kv_pool.set_kv_buffer(
+                        layer, cache_loc, k, v, layer.k_scale, layer.v_scale
+                    )
+                else:
+                    forward_batch.token_to_kv_pool.set_mla_kv_buffer(
+                        layer,
+                        cache_loc,
+                        k,
+                        k_rope,
+                    )
+
+        # Use precomputed metadata across all layers
+        metadata = self.forward_metadata
+        local_attn_metadata = getattr(metadata, "local_attn_metadata", None)
+        use_local_attention = (
+            self.attention_chunk_size is not None and local_attn_metadata is not None
+        )
+        # We do cascade attention for Draft Decode with topk > 1
+        use_cascade_attn = self.topk > 1
+
+        # Calculate window size (can be moved to metadata if layer properties don't change)
+        # we don't do layer.sliding_window_size - 1 since in model.get_attention_sliding_window_size() we already - 1
+        # here is two side inclusive
+        window_size = (
+            (layer.sliding_window_size, 0)
+            if layer.sliding_window_size is not None and layer.sliding_window_size > -1
+            else (-1, -1)
+        )
+        causal = not layer.is_cross_attention
+
+        k_descale, v_descale = None, None
+        # only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention
+        # has corresponding quantization method so that layer.k_scale is not None
+        if self.kv_cache_dtype_str != "auto":
+            if layer.k_scale is not None:
+                descale_shape = (forward_batch.batch_size, layer.tp_k_head_num)
+                k_descale = layer.k_scale.expand(descale_shape)
+                v_descale = layer.v_scale.expand(descale_shape)
+            q = q.to(self.kv_cache_dtype)
+
+        if not self.use_mla:
+            # Do multi-head attention
+
+            key_cache, value_cache = forward_batch.token_to_kv_pool.get_kv_buffer(
+                layer.layer_id
+            )
+            key_cache = key_cache.view(
+                -1, self.page_size, layer.tp_k_head_num, layer.head_dim
+            )
+            value_cache = value_cache.view(
+                -1, self.page_size, layer.tp_v_head_num, layer.head_dim
+            )
+
+            if layer.is_cross_attention:
+                # Always use non-chunked logic for cross-attention
+                o = flash_attn_with_kvcache(
+                    q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                    k_cache=key_cache,
+                    v_cache=value_cache,
+                    page_table=metadata.encoder_page_table,
+                    cache_seqlens=metadata.encoder_lens_int32,
+                    cu_seqlens_q=metadata.cu_seqlens_q,
+                    cu_seqlens_k_new=metadata.encoder_cu_seqlens_k,
+                    max_seqlen_q=1,
+                    softmax_scale=layer.scaling,
+                    causal=False,
+                    window_size=(-1, -1),
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                )
+            elif use_local_attention:
+                # Use chunked (local) attention batching for self-attention
+                o = flash_attn_with_kvcache(
+                    q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                    k_cache=key_cache,
+                    v_cache=value_cache,
+                    page_table=local_attn_metadata.local_block_table,
+                    cache_seqlens=local_attn_metadata.local_seqused_k,
+                    cu_seqlens_q=local_attn_metadata.local_query_start_loc,
+                    cu_seqlens_k_new=metadata.cu_seqlens_k,
+                    max_seqlen_q=local_attn_metadata.local_max_query_len,
+                    softmax_scale=layer.scaling,
+                    causal=True,
+                    window_size=(-1, -1),
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                )
+            else:
+                page_table = metadata.page_table
+                cache_seqlens = metadata.cache_seqlens_int32
+                cu_seqlens_k = metadata.cu_seqlens_k
+                max_seqlen_q = metadata.max_seq_len_q
+                q_reshaped = q.contiguous().view(
+                    -1, layer.tp_q_head_num, layer.head_dim
+                )
+
+                # Default: single-token self-attention
+                result = flash_attn_with_kvcache(
+                    q=q_reshaped,
+                    k_cache=key_cache,
+                    v_cache=value_cache,
+                    page_table=page_table,
+                    cache_seqlens=cache_seqlens,
+                    cu_seqlens_q=metadata.cu_seqlens_q,
+                    cu_seqlens_k_new=cu_seqlens_k,
+                    max_seqlen_q=max_seqlen_q,
+                    softmax_scale=layer.scaling,
+                    causal=False if use_cascade_attn else causal,
+                    window_size=window_size,
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                    return_softmax_lse=use_cascade_attn,
+                )
+                if use_cascade_attn:
+                    o, softmax_lse, *rest = result
+                    o_expand, softmax_lse_expand, *rest_expand = (
+                        flash_attn_with_kvcache(
+                            q=q_reshaped,
+                            k_cache=key_cache,
+                            v_cache=value_cache,
+                            page_table=self.forward_metadata_spec_decode_expand.page_table,
+                            cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
+                            cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
+                            cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k,
+                            max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
+                            softmax_scale=layer.scaling,
+                            causal=False,
+                            window_size=window_size,
+                            softcap=layer.logit_cap,
+                            k_descale=k_descale,
+                            v_descale=v_descale,
+                            return_softmax_lse=True,
+                        )
+                    )
+                    o, _ = merge_state_v2(
+                        o,
+                        softmax_lse.T.contiguous(),
+                        o_expand,
+                        softmax_lse_expand.T.contiguous(),
+                    )
+                else:
+                    o = result
+        else:
+            # Do absorbed multi-latent attention
+            kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
+            k_rope = kv_cache[:, :, layer.v_head_dim :]
+            c_kv = kv_cache[:, :, : layer.v_head_dim]
+            k_rope_cache = k_rope.view(
+                -1,
+                self.page_size,
+                layer.tp_k_head_num,
+                layer.head_dim - layer.v_head_dim,
+            )
+            c_kv_cache = c_kv.view(
+                -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
+            )
+
+            if q_rope is not None:
+                q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+                q_rope = q_rope.view(
+                    -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+                )
+            else:
+                q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
+                q_nope = q_all[:, :, : layer.v_head_dim]
+                q_rope = q_all[:, :, layer.v_head_dim :]
+            max_seqlen_q = metadata.max_seq_len_q
+
+            result = flash_attn_with_kvcache(
+                q=q_rope,
+                k_cache=k_rope_cache,
+                v_cache=c_kv_cache,
+                qv=q_nope,
+                page_table=metadata.page_table,
+                cache_seqlens=metadata.cache_seqlens_int32,
+                cu_seqlens_q=metadata.cu_seqlens_q,
+                cu_seqlens_k_new=metadata.cu_seqlens_k,
+                max_seqlen_q=max_seqlen_q,
+                softmax_scale=layer.scaling,
+                causal=False if use_cascade_attn else causal,
+                softcap=layer.logit_cap,
+                k_descale=k_descale,
+                v_descale=v_descale,
+                return_softmax_lse=use_cascade_attn,  # softmax_lse is needed for merge states
+            )
+            if use_cascade_attn:
+                o, softmax_lse, *rest = result
+                o_expand, softmax_lse_expand, *rest_expand = flash_attn_with_kvcache(
+                    q=q_rope,
+                    k_cache=k_rope_cache,
+                    v_cache=c_kv_cache,
+                    qv=q_nope,
+                    page_table=self.forward_metadata_spec_decode_expand.page_table,
+                    cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
+                    cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
+                    cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k,
+                    max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
+                    softmax_scale=layer.scaling,
+                    causal=False,
+                    window_size=window_size,
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                    return_softmax_lse=True,
+                )
+                o, _ = merge_state_v2(
+                    o,
+                    softmax_lse.T.contiguous(),
+                    o_expand,
+                    softmax_lse_expand.T.contiguous(),
+                )
+            else:
+                o = result
+
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+
+    def init_cuda_graph_state(self, max_bs: int):
+        """Initialize CUDA graph state for the attention backend.
+
+        Args:
+            max_bs (int): Maximum batch size to support in CUDA graphs
+
+        This creates fixed-size tensors that will be reused during CUDA graph replay
+        to avoid memory allocations.
+        """
+
+        # This is being used by normal decode and draft decode when topk == 1
+        self.decode_cuda_graph_metadata = {
+            "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
+            "cu_seqlens_q": torch.arange(
+                0, max_bs + 1, dtype=torch.int32, device=self.device
+            ),
+            "cu_seqlens_k": torch.zeros(
+                max_bs + 1, dtype=torch.int32, device=self.device
+            ),
+            "page_table": torch.zeros(
+                max_bs,
+                (self.max_context_len + self.page_size - 1) // self.page_size,
+                dtype=torch.int32,
+                device=self.device,
+            ),
+            "page_table_draft_decode": torch.zeros(
+                max_bs,
+                (self.max_context_len + self.page_size - 1) // self.page_size,
+                dtype=torch.int32,
+                device=self.device,
+            ),
+            "strided_indices": torch.arange(
+                0, self.max_context_len, self.page_size, device=self.device
+            ),
+        }
+
+        # This is used by draft decode's first half of metadata when topk > 1
+        if self.topk > 1:
+            self.draft_decode_metadata_topk_normal = {
+                "cache_seqlens": torch.zeros(
+                    max_bs, dtype=torch.int32, device=self.device
+                ),
+                "cu_seqlens_q": torch.arange(
+                    0,
+                    max_bs * self.topk + 1,
+                    step=self.topk,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+                "cu_seqlens_k": torch.zeros(
+                    max_bs + 1, dtype=torch.int32, device=self.device
+                ),
+                "page_table": torch.zeros(
+                    max_bs,
+                    self.max_context_len,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+            }
+
+            # This is used by draft decode's second half of metadata when topk > 1
+            decode_length = self.speculative_step_id + 1
+            self.draft_decode_metadata_topk_expand = {
+                "cache_seqlens": torch.full(
+                    (max_bs * self.topk,),
+                    decode_length,
+                    device=self.device,
+                    dtype=torch.int32,
+                ),
+                "cu_seqlens_q": torch.arange(
+                    0,
+                    max_bs * self.topk + 1,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+                "cu_seqlens_k": torch.arange(
+                    0,
+                    max_bs * self.topk * decode_length + 1,
+                    step=decode_length,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+                "page_table": torch.zeros(
+                    max_bs * self.topk,
+                    decode_length,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+            }
+
+        if (
+            self.speculative_num_draft_tokens is not None
+            and self.speculative_num_draft_tokens > 0
+        ):
+            self.target_verify_metadata = {
+                "cache_seqlens": torch.zeros(
+                    max_bs, dtype=torch.int32, device=self.device
+                ),
+                "cu_seqlens_q": torch.arange(
+                    0,
+                    max_bs * self.speculative_num_draft_tokens + 1,
+                    step=self.speculative_num_draft_tokens,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+                "cu_seqlens_k": torch.zeros(
+                    max_bs + 1, dtype=torch.int32, device=self.device
+                ),
+                "page_table": torch.zeros(
+                    max_bs,
+                    (self.max_context_len + self.page_size - 1) // self.page_size,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+                "strided_indices": torch.arange(
+                    0, self.max_context_len, self.page_size, device=self.device
+                ),
+            }
+
+        if self.topk > 1:
+            self.target_verify_metadata_topk_normal = {
+                "cache_seqlens": torch.zeros(
+                    max_bs, dtype=torch.int32, device=self.device
+                ),
+                "cu_seqlens_q": torch.arange(
+                    0,
+                    max_bs * self.speculative_num_draft_tokens + 1,
+                    step=self.speculative_num_draft_tokens,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+                "cu_seqlens_k": torch.zeros(
+                    max_bs + 1, dtype=torch.int32, device=self.device
+                ),
+                "page_table": torch.zeros(
+                    max_bs,
+                    self.max_context_len,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+            }
+
+            self.target_verify_metadata_topk_expand = {
+                "cache_seqlens": torch.zeros(
+                    max_bs * self.speculative_num_draft_tokens,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+                "cu_seqlens_k": torch.zeros(
+                    max_bs * self.speculative_num_draft_tokens + 1,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+                "cu_seqlens_q": torch.arange(
+                    0,
+                    max_bs * self.speculative_num_draft_tokens + 1,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+                "page_table": torch.zeros(
+                    max_bs * self.speculative_num_draft_tokens,
+                    self.speculative_num_draft_tokens,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+            }
+
+        self.encoder_metadata = {
+            "encoder_page_table": torch.zeros(
+                max_bs,
+                self.max_context_len,
+                dtype=torch.int32,
+                device=self.device,
+            ),
+            "encoder_lens_int32": torch.zeros(
+                max_bs, dtype=torch.int32, device=self.device
+            ),
+            "encoder_cu_seqlens_k": torch.zeros(
+                max_bs + 1, dtype=torch.int32, device=self.device
+            ),
+        }
+
+    def init_forward_metadata_capture_cuda_graph(
+        self,
+        bs: int,
+        num_tokens: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        """Initialize forward metadata for capturing CUDA graph."""
+        metadata = FlashAttentionMetadata()
+
+        # metadata_expand is needed for Spec Decoding when top k > 1
+        metadata_expand = FlashAttentionMetadata()
+
+        device = seq_lens.device
+        if forward_mode.is_decode_or_idle():
+            if spec_info is not None:
+                # Draft Decode
+                if self.topk <= 1:
+                    # When topk = 1, we use the normal decode metadata
+                    metadata.cache_seqlens_int32 = self.decode_cuda_graph_metadata[
+                        "cache_seqlens"
+                    ][:bs]
+                    metadata.max_seq_len_k = seq_lens.max().item() + (
+                        self.speculative_step_id + 1
+                    )
+                    metadata.cu_seqlens_q = self.decode_cuda_graph_metadata[
+                        "cu_seqlens_q"
+                    ][: bs + 1]
+                    metadata.cu_seqlens_k = torch.nn.functional.pad(
+                        torch.cumsum(
+                            metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+                        ),
+                        (1, 0),
+                    )
+                    metadata.page_table = self.decode_cuda_graph_metadata[
+                        "page_table_draft_decode"
+                    ][req_pool_indices, :]
+                    self.decode_cuda_graph_metadata[bs] = metadata
+                else:
+                    # When top k > 1, we need two specific draft decode metadata, and then merge states
+                    # 1. The first half of metadata for prefix tokens
+                    metadata.cache_seqlens_int32 = (
+                        self.draft_decode_metadata_topk_normal["cache_seqlens"][:bs]
+                    )
+                    metadata.max_seq_len_q = self.topk
+                    metadata.max_seq_len_k = seq_lens.max().item()
+                    metadata.cu_seqlens_q = self.draft_decode_metadata_topk_normal[
+                        "cu_seqlens_q"
+                    ][: bs + 1]
+                    metadata.cu_seqlens_k = self.draft_decode_metadata_topk_normal[
+                        "cu_seqlens_k"
+                    ][: bs + 1]
+                    metadata.page_table = self.draft_decode_metadata_topk_normal[
+                        "page_table"
+                    ][req_pool_indices, :]
+
+                    # 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk)
+                    metadata_expand.cache_seqlens_int32 = (
+                        self.draft_decode_metadata_topk_expand["cache_seqlens"][
+                            : bs * self.topk
+                        ]
+                    )
+                    metadata_expand.max_seq_len_q = 1
+                    metadata_expand.max_seq_len_k = (
+                        self.speculative_step_id + 1
+                    )  # , do this in replay
+                    metadata_expand.cu_seqlens_q = (
+                        self.draft_decode_metadata_topk_expand["cu_seqlens_q"][
+                            : bs * self.topk + 1
+                        ]
+                    )
+                    metadata_expand.cu_seqlens_k = (
+                        self.draft_decode_metadata_topk_expand["cu_seqlens_k"][
+                            : bs * self.topk + 1
+                        ]
+                    )
+                    metadata_expand.page_table = self.draft_decode_metadata_topk_expand[
+                        "page_table"
+                    ][: bs * self.topk]
+                    self.draft_decode_metadata_topk_normal[bs] = metadata
+                    self.draft_decode_metadata_topk_expand[bs] = metadata_expand
+            else:
+                # Normal Decode
+                # Get sequence information
+                metadata.cache_seqlens_int32 = seq_lens.to(torch.int32)
+                batch_size = len(seq_lens)
+                device = seq_lens.device
+                metadata.cu_seqlens_k = torch.nn.functional.pad(
+                    torch.cumsum(seq_lens, dim=0, dtype=torch.int32), (1, 0)
+                )
+                # Precompute maximum sequence length
+                metadata.max_seq_len_k = seq_lens.max().item()
+                # Precompute page table
+                metadata.page_table = self.decode_cuda_graph_metadata["page_table"][
+                    req_pool_indices, :
+                ]
+                # Precompute cumulative sequence lengths
+                metadata.cu_seqlens_q = torch.arange(
+                    0, batch_size + 1, dtype=torch.int32, device=device
+                )
+                self.decode_cuda_graph_metadata[bs] = metadata
+
+        elif forward_mode.is_target_verify():
+            if self.topk <= 1:
+                metadata.cache_seqlens_int32 = self.target_verify_metadata[
+                    "cache_seqlens"
+                ][:bs]
+                metadata.cache_seqlens_int32.copy_(
+                    (seq_lens + self.speculative_num_draft_tokens).to(torch.int32)
+                )
+
+                metadata.max_seq_len_q = self.speculative_num_draft_tokens
+                metadata.max_seq_len_k = (
+                    seq_lens.max().item() + self.speculative_num_draft_tokens
+                )
+
+                metadata.cu_seqlens_q = torch.arange(
+                    0,
+                    bs * self.speculative_num_draft_tokens + 1,
+                    self.speculative_num_draft_tokens,
+                    dtype=torch.int32,
+                    device=device,
+                )
+
+                metadata.cu_seqlens_k = self.target_verify_metadata["cu_seqlens_k"][
+                    : (bs + 1)
+                ]
+
+                metadata.page_table = self.target_verify_metadata["page_table"][
+                    req_pool_indices, :
+                ]
+
+                self.target_verify_metadata[bs] = metadata
+            else:
+                # When topk > 1, we need two specific target verify metadata, and then merge states
+                # 1. The first half of metadata for prefix tokens
+                metadata.cache_seqlens_int32 = self.target_verify_metadata_topk_normal[
+                    "cache_seqlens"
+                ][:bs]
+                metadata.max_seq_len_q = self.speculative_num_draft_tokens
+                # metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item(), do this in replay
+                metadata.cu_seqlens_q = self.target_verify_metadata_topk_normal[
+                    "cu_seqlens_q"
+                ][: bs + 1]
+                metadata.cu_seqlens_k = self.target_verify_metadata_topk_normal[
+                    "cu_seqlens_k"
+                ][: bs + 1]
+                metadata.page_table = self.target_verify_metadata_topk_normal[
+                    "page_table"
+                ][req_pool_indices, :]
+
+                # 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk)
+                metadata_expand.cache_seqlens_int32 = (
+                    self.target_verify_metadata_topk_expand["cache_seqlens"][
+                        : bs * self.speculative_num_draft_tokens
+                    ]
+                )
+                metadata_expand.max_seq_len_q = 1
+                metadata_expand.cu_seqlens_q = self.target_verify_metadata_topk_expand[
+                    "cu_seqlens_q"
+                ][: bs * self.speculative_num_draft_tokens + 1]
+                metadata_expand.cu_seqlens_k = self.target_verify_metadata_topk_expand[
+                    "cu_seqlens_k"
+                ][: bs * self.speculative_num_draft_tokens + 1]
+
+                metadata_expand.page_table = self.target_verify_metadata_topk_expand[
+                    "page_table"
+                ][: bs * self.speculative_num_draft_tokens]
+
+                self.target_verify_metadata_topk_normal[bs] = metadata
+                self.target_verify_metadata_topk_expand[bs] = metadata_expand
+
+        if encoder_lens is not None:
+            encoder_bs = encoder_lens.numel()
+            metadata.encoder_lens_int32 = self.encoder_metadata["encoder_lens_int32"][
+                :encoder_bs
+            ]
+            metadata.encoder_cu_seqlens_k = self.encoder_metadata[
+                "encoder_cu_seqlens_k"
+            ][: (encoder_bs + 1)]
+
+            metadata.encoder_page_table = self.encoder_metadata["encoder_page_table"][
+                req_pool_indices, :
+            ]
+
+        self.forward_metadata = metadata
+        self.forward_metadata_spec_decode_expand = metadata_expand
+
+    def init_forward_metadata_replay_cuda_graph(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        seq_lens_cpu: Optional[torch.Tensor],
+        out_cache_loc: torch.Tensor = None,
+    ):
+        """Initialize forward metadata for replaying CUDA graph."""
+        seq_lens = seq_lens[:bs]
+        seq_lens_cpu = seq_lens_cpu[:bs]
+        req_pool_indices = req_pool_indices[:bs]
+        device = seq_lens.device
+        metadata = None
+        metadata_expand = None
+
+        if forward_mode.is_decode_or_idle():
+
+            if spec_info is not None:
+                # Draft Decode
+                if self.topk <= 1:
+                    metadata = self.decode_cuda_graph_metadata[bs]
+                    # When topk = 1, we use the normal decode metadata
+                    metadata.cache_seqlens_int32.copy_(
+                        (seq_lens + (self.speculative_step_id + 1)).to(torch.int32)
+                    )
+
+                    metadata.max_seq_len_k = seq_lens_cpu.max().item() + (
+                        self.speculative_step_id + 1
+                    )
+                    metadata.cu_seqlens_k.copy_(
+                        torch.nn.functional.pad(
+                            torch.cumsum(
+                                metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+                            ),
+                            (1, 0),
+                        )
+                    )
+
+                    max_seq_pages = (
+                        metadata.max_seq_len_k + self.page_size - 1
+                    ) // self.page_size
+                    page_indices = self.req_to_token[
+                        req_pool_indices[:, None],
+                        self.decode_cuda_graph_metadata["strided_indices"][
+                            :max_seq_pages
+                        ],
+                    ]
+
+                    page_indices //= self.page_size
+                    metadata.page_table[:, :max_seq_pages].copy_(page_indices)
+                else:
+                    # When top k > 1, we need two specific draft decode metadata, and then merge states
+                    # 1. The first half of metadata for prefix tokens
+                    metadata = self.draft_decode_metadata_topk_normal[bs]
+                    metadata.cache_seqlens_int32.copy_(seq_lens.to(torch.int32))
+                    # metadata.max_seq_len_q = self.topk, already set in capture
+                    metadata.max_seq_len_k = seq_lens_cpu.max().item()
+                    # metadata.cu_seqlens_q already set in capture
+                    metadata.cu_seqlens_k.copy_(
+                        torch.nn.functional.pad(
+                            torch.cumsum(
+                                metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+                            ),
+                            (1, 0),
+                        )
+                    )
+
+                    page_table = self.req_to_token[
+                        req_pool_indices, : metadata.max_seq_len_k
+                    ]
+
+                    metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table)
+
+                    # 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk)
+                    metadata_expand = self.draft_decode_metadata_topk_expand[bs]
+                    decode_length = self.speculative_step_id + 1
+                    cache_loc = out_cache_loc.view(
+                        self.speculative_num_steps, -1
+                    ).T.contiguous()
+                    metadata_expand.page_table[: cache_loc.shape[0]].copy_(
+                        cache_loc[:, :decode_length].contiguous().to(torch.int32)
+                    )
+                # TODO: we need to test this part for llama 4 eagle case
+                self._init_local_attn_metadata(metadata, device)
+            else:
+                metadata = self.decode_cuda_graph_metadata[bs]
+                # Normal Decode
+                max_len = seq_lens_cpu.max().item()
+                metadata.max_seq_len_k = max_len
+
+                metadata.cache_seqlens_int32 = seq_lens.to(torch.int32)
+                metadata.cu_seqlens_k = torch.nn.functional.pad(
+                    torch.cumsum(seq_lens, dim=0, dtype=torch.int32), (1, 0)
+                )
+
+                max_seq_pages = (
+                    metadata.max_seq_len_k + self.page_size - 1
+                ) // self.page_size
+                page_indices = self.req_to_token[
+                    req_pool_indices[:, None],
+                    self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages][
+                        None, :
+                    ],
+                ]
+                page_indices //= self.page_size
+                metadata.page_table[:, :max_seq_pages].copy_(page_indices)
+                metadata.page_table[:, max_seq_pages:].fill_(0)
+
+                self._init_local_attn_metadata(metadata, device)
+        elif forward_mode.is_target_verify():
+            if self.topk <= 1:
+                metadata = self.target_verify_metadata[bs]
+                metadata.cache_seqlens_int32.copy_(
+                    (seq_lens + self.speculative_num_draft_tokens).to(torch.int32)
+                )
+
+                metadata.max_seq_len_k = (
+                    seq_lens_cpu.max().item() + self.speculative_num_draft_tokens
+                )
+                metadata.cu_seqlens_k.copy_(
+                    torch.nn.functional.pad(
+                        torch.cumsum(
+                            metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+                        ),
+                        (1, 0),
+                    )
+                )
+                max_seq_pages = (
+                    metadata.max_seq_len_k + self.page_size - 1
+                ) // self.page_size
+                page_indices = self.req_to_token[
+                    req_pool_indices[:, None],
+                    self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages],
+                ]
+                page_indices //= self.page_size
+                metadata.page_table[:, :max_seq_pages].copy_(page_indices)
+            else:
+                # When topk > 1, we need two specific target verify metadata, and then merge states
+                # 1. The first half of metadata for prefix tokens
+                metadata = self.target_verify_metadata_topk_normal[bs]
+                metadata.cache_seqlens_int32.copy_(seq_lens.to(torch.int32))
+                # metadata.max_seq_len_q = self.speculative_num_draft_tokens, already set in capture
+                metadata.max_seq_len_k = seq_lens_cpu.max().item()
+                # metadata.cu_seqlens_q already set in capture
+                metadata.cu_seqlens_k.copy_(
+                    torch.nn.functional.pad(
+                        torch.cumsum(
+                            metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+                        ),
+                        (1, 0),
+                    )
+                )
+                page_table = self.req_to_token[
+                    req_pool_indices, : metadata.max_seq_len_k
+                ]
+                metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table)
+
+                # 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk)
+                metadata_expand = self.target_verify_metadata_topk_expand[bs]
+                # metadata_expand.max_seq_len_q = 1, already set in capture
+                # metadata_expand.cu_seqlens_q already set in capture
+
+                offsets = torch.arange(
+                    self.speculative_num_draft_tokens, device=device
+                ).unsqueeze(
+                    0
+                )  # shape: (1, self.speculative_num_draft_tokens)
+                cols = offsets.expand(seq_lens.numel(), -1) + seq_lens.unsqueeze(1)
+                cum_len = torch.nn.functional.pad(
+                    torch.cumsum(
+                        (
+                            seq_lens + self.speculative_num_draft_tokens
+                        ).repeat_interleave(self.speculative_num_draft_tokens),
+                        dim=0,
+                    ),
+                    (1, 0),
+                )[:-1]
+                mask_extraction_indices = (
+                    cols.repeat_interleave(self.speculative_num_draft_tokens, dim=0)
+                    + cum_len[:, None]
+                ).view(1, -1)
+                # avoid extracting padded seq indices which will be out of boundary
+                mask_extraction_indices[
+                    :, spec_info.positions.numel() * self.speculative_num_draft_tokens :
+                ].fill_(0)
+
+                mask = spec_info.custom_mask[mask_extraction_indices].view(
+                    -1, self.speculative_num_draft_tokens
+                )  # (bsz * draft_num, draft_num)
+                col_indices = offsets.expand(
+                    mask.shape[0], self.speculative_num_draft_tokens
+                )
+                keys = torch.where(
+                    mask, col_indices, col_indices + self.speculative_num_draft_tokens
+                )
+                _, sort_order = torch.sort(keys, dim=1)
+
+                non_masked_page_table = (
+                    self.req_to_token[req_pool_indices, :]
+                    .gather(1, cols)
+                    .repeat_interleave(self.speculative_num_draft_tokens, dim=0)
+                )  # (bsz, draft_num)
+                metadata_expand.page_table.copy_(
+                    non_masked_page_table.gather(1, sort_order)
+                )
+                metadata_expand.cache_seqlens_int32.copy_(
+                    mask.sum(dim=1).to(torch.int32)
+                )
+                metadata_expand.cu_seqlens_k.copy_(
+                    torch.nn.functional.pad(
+                        torch.cumsum(
+                            metadata_expand.cache_seqlens_int32,
+                            dim=0,
+                            dtype=torch.int32,
+                        ),
+                        (1, 0),
+                    )
+                )
+                metadata_expand.max_seq_len_k = (
+                    metadata_expand.cache_seqlens_int32.max().item()
+                )
+
+        if encoder_lens is not None:
+            # Only support encoder size 1 for now
+            metadata.encoder_max_seq_len_k = encoder_lens[0]
+            metadata.encoder_lens_int32.copy_(encoder_lens[:1])
+            metadata.encoder_cu_seqlens_k.copy_(
+                torch.nn.functional.pad(
+                    torch.cumsum(metadata.encoder_lens_int32, dim=0, dtype=torch.int32),
+                    (1, 0),
+                )
+            )
+
+            metadata.encoder_page_table[:, : metadata.encoder_max_seq_len_k].copy_(
+                self.req_to_token[req_pool_indices, : metadata.encoder_max_seq_len_k]
+            )
+
+            # Update the regular page table
+            page_table = self.req_to_token[
+                req_pool_indices,
+                metadata.encoder_max_seq_len_k : (
+                    metadata.encoder_max_seq_len_k + metadata.max_seq_len_k
+                ),
+            ]
+            metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table)
+
+        self.forward_metadata = metadata
+        self.forward_metadata_spec_decode_expand = metadata_expand
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        """Get the fill value for sequence length in CUDA graph."""
+        return 0
+
+    def _init_local_attn_metadata(self, metadata: FlashAttentionMetadata, device):
+        """Centralized utility to initialize local_attn_metadata if chunked attention is enabled."""
+        if self.attention_chunk_size is None:
+            metadata.local_attn_metadata = None
+            return
+
+        cu_seqlens_q = metadata.cu_seqlens_q
+        cache_seqlens_int32 = metadata.cache_seqlens_int32
+        page_table = metadata.page_table
+        if cu_seqlens_q is None or cache_seqlens_int32 is None or page_table is None:
+            metadata.local_attn_metadata = None
+            return
+
+        cu_seqlens_q_np = cu_seqlens_q.cpu().numpy()
+        seq_lens_np = cache_seqlens_int32.cpu().numpy()
+        (
+            seqlens_q_local_np,
+            cu_seqlens_q_local_np,
+            seqlens_k_local_np,
+            block_table_local,
+        ) = make_local_attention_virtual_batches(
+            self.attention_chunk_size,
+            cu_seqlens_q_np,
+            seq_lens_np,
+            page_table,
+            self.page_size,
+        )
+        local_metadata = FlashAttentionMetadata.LocalAttentionMetadata(
+            local_query_start_loc=torch.from_numpy(cu_seqlens_q_local_np).to(device),
+            local_seqused_k=torch.from_numpy(seqlens_k_local_np).to(device),
+            local_block_table=block_table_local.to(device),
+            local_max_query_len=int(seqlens_q_local_np.max()),
+            local_max_seq_len=int(seqlens_k_local_np.max()),
+        )
+        metadata.local_attn_metadata = local_metadata
+
+
+class FlashAttentionMultiStepBackend:
+
+    def __init__(
+        self, model_runner: ModelRunner, topk: int, speculative_num_steps: int
+    ):
+        self.model_runner = model_runner
+        self.topk = topk
+        self.speculative_num_steps = speculative_num_steps
+        self.attn_backends = []
+        for i in range(self.speculative_num_steps):
+            self.attn_backends.append(
+                FlashAttentionBackend(
+                    model_runner,
+                    speculative_step_id=i,
+                    topk=self.topk,
+                    speculative_num_steps=self.speculative_num_steps,
+                )
+            )
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        for i in range(self.speculative_num_steps - 1):
+            self.attn_backends[i].init_forward_metadata(forward_batch)
+
+    def init_cuda_graph_state(self, max_bs: int):
+        for i in range(self.speculative_num_steps):
+            self.attn_backends[i].init_cuda_graph_state(max_bs)
+
+    def init_forward_metadata_capture_cuda_graph(
+        self,
+        forward_batch: ForwardBatch,
+    ):
+        assert forward_batch.spec_info is not None
+        assert isinstance(forward_batch.spec_info, EagleDraftInput)
+
+        for i in range(self.speculative_num_steps - 1):
+            self.attn_backends[i].init_forward_metadata_capture_cuda_graph(
+                forward_batch.batch_size,
+                forward_batch.batch_size * self.topk,
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                encoder_lens=forward_batch.encoder_lens,
+                forward_mode=ForwardMode.DECODE,
+                spec_info=forward_batch.spec_info,
+            )
+
+    def init_forward_metadata_replay_cuda_graph(
+        self, forward_batch: ForwardBatch, bs: int
+    ):
+        assert forward_batch.spec_info is not None
+        assert isinstance(forward_batch.spec_info, EagleDraftInput)
+
+        for i in range(self.speculative_num_steps - 1):
+            self.attn_backends[i].init_forward_metadata_replay_cuda_graph(
+                bs,
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                encoder_lens=forward_batch.encoder_lens,
+                forward_mode=ForwardMode.DECODE,
+                spec_info=forward_batch.spec_info,
+                seq_lens_cpu=forward_batch.seq_lens_cpu,
+                out_cache_loc=forward_batch.out_cache_loc,
+            )

lib/python3.12/site-packages/sglang/srt/layers/attention/flashinfer_backend.py

@@ -0,0 +1,1266 @@
+from __future__ import annotations
+
+"""
+Support different attention backends.
+Now there are two backends: FlashInfer and Triton.
+FlashInfer is faster and Triton is easier to customize.
+Each backend supports two operators: extend (i.e. prefill with cached prefix) and decode.
+"""
+
+import os
+from dataclasses import dataclass
+from enum import Enum, auto
+from functools import partial
+from typing import TYPE_CHECKING, Callable, List, Optional, Union
+
+import torch
+
+from sglang.global_config import global_config
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
+from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+from sglang.srt.utils import is_flashinfer_available, next_power_of_2
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+if is_flashinfer_available():
+    from flashinfer import (
+        BatchDecodeWithPagedKVCacheWrapper,
+        BatchPrefillWithPagedKVCacheWrapper,
+        BatchPrefillWithRaggedKVCacheWrapper,
+    )
+    from flashinfer.cascade import merge_state
+    from flashinfer.decode import _get_range_buf, get_seq_lens
+
+
+class WrapperDispatch(Enum):
+    SLIDING_WINDOW = auto()
+    CROSS_ATTENTION = auto()
+
+
+@dataclass
+class DecodeMetadata:
+    decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper]
+
+
+@dataclass
+class PrefillMetadata:
+    prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper]
+    use_ragged: bool
+    extend_no_prefix: bool
+
+
+# Reuse this workspace buffer across all flashinfer wrappers
+global_workspace_buffer = None
+
+
+class FlashInferAttnBackend(AttentionBackend):
+    """Flashinfer attention kernels."""
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        skip_prefill: bool = False,
+        kv_indptr_buf: Optional[torch.Tensor] = None,
+        kv_last_page_len_buf: Optional[torch.Tensor] = None,
+    ):
+        super().__init__()
+
+        # Parse constants
+        self.decode_use_tensor_cores = should_use_tensor_core(
+            kv_cache_dtype=model_runner.kv_cache_dtype,
+            num_attention_heads=model_runner.model_config.num_attention_heads
+            // get_attention_tp_size(),
+            num_kv_heads=model_runner.model_config.get_num_kv_heads(
+                get_attention_tp_size()
+            ),
+        )
+        self.max_context_len = model_runner.model_config.context_len
+        self.skip_prefill = skip_prefill
+        self.is_multimodal = model_runner.model_config.is_multimodal
+        self.kv_cache_dtype = model_runner.kv_cache_dtype
+        self.kv_cache_dtype_str = model_runner.server_args.kv_cache_dtype
+
+        assert not (
+            model_runner.sliding_window_size is not None
+            and model_runner.model_config.is_encoder_decoder
+        ), "Sliding window and cross attention are not supported together"
+
+        if model_runner.sliding_window_size is not None:
+            self.num_wrappers = 2
+            self.dispatch_reason = WrapperDispatch.SLIDING_WINDOW
+        elif model_runner.model_config.is_encoder_decoder:
+            self.num_wrappers = 2
+            self.dispatch_reason = WrapperDispatch.CROSS_ATTENTION
+        else:
+            self.num_wrappers = 1
+            self.dispatch_reason = None
+
+        # Qwen2/Qwen3 models require higher flashinfer workspace size
+        if (
+            "Qwen2ForCausalLM" in model_runner.model_config.hf_config.architectures
+            or "Qwen3ForCausalLM" in model_runner.model_config.hf_config.architectures
+        ):
+            global_config.flashinfer_workspace_size = 512 * 1024 * 1024
+
+        # Allocate buffers
+        global global_workspace_buffer
+        if global_workspace_buffer is None:
+            global_workspace_buffer = torch.empty(
+                global_config.flashinfer_workspace_size,
+                dtype=torch.uint8,
+                device=model_runner.device,
+            )
+        self.workspace_buffer = global_workspace_buffer
+        max_bs = model_runner.req_to_token_pool.size
+        if kv_indptr_buf is None:
+            self.kv_indptr = [
+                torch.zeros(
+                    (max_bs + 1,), dtype=torch.int32, device=model_runner.device
+                )
+                for _ in range(self.num_wrappers)
+            ]
+        else:
+            assert self.num_wrappers == 1
+            self.kv_indptr = [kv_indptr_buf]
+
+        if kv_last_page_len_buf is None:
+            self.kv_last_page_len = torch.ones(
+                (max_bs,), dtype=torch.int32, device=model_runner.device
+            )
+        else:
+            assert self.num_wrappers == 1
+            self.kv_last_page_len = kv_last_page_len_buf
+
+        if not self.skip_prefill:
+            self.qo_indptr = [
+                torch.zeros(
+                    (max_bs + 1,), dtype=torch.int32, device=model_runner.device
+                )
+                for _ in range(self.num_wrappers)
+            ]
+
+        self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
+            self.workspace_buffer, "NHD"
+        )
+
+        # Two wrappers: one for sliding window attention and one for full attention.
+        # Using two wrappers is unnecessary in the current PR, but are prepared for future PRs
+        self.prefill_wrappers_paged = []
+        self.prefill_wrappers_verify = []
+        self.decode_wrappers = []
+        for _ in range(self.num_wrappers):
+            if not skip_prefill:
+                self.prefill_wrappers_paged.append(
+                    BatchPrefillWithPagedKVCacheWrapper(
+                        self.workspace_buffer,
+                        "NHD",
+                        backend="fa2",
+                    )
+                )
+                self.prefill_wrappers_verify.append(
+                    BatchPrefillWithPagedKVCacheWrapper(
+                        self.workspace_buffer,
+                        "NHD",
+                    )
+                )
+            self.decode_wrappers.append(
+                BatchDecodeWithPagedKVCacheWrapper(
+                    self.workspace_buffer,
+                    "NHD",
+                    use_tensor_cores=self.decode_use_tensor_cores,
+                )
+            )
+
+        # Create indices updater
+        if not skip_prefill:
+            self.indices_updater_prefill = FlashInferIndicesUpdaterPrefill(
+                model_runner, self
+            )  # for verify
+        self.indices_updater_decode = FlashInferIndicesUpdaterDecode(model_runner, self)
+
+        # Other metadata
+        self.forward_metadata: Union[PrefillMetadata, DecodeMetadata] = None
+        self.decode_cuda_graph_metadata = {}
+        self.prefill_cuda_graph_metadata = {}  # For verify
+        self.draft_extend_cuda_graph_metadata = {}  # For draft extend
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        if forward_batch.forward_mode.is_decode_or_idle():
+            self.indices_updater_decode.update(
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                decode_wrappers=self.decode_wrappers,
+                encoder_lens=forward_batch.encoder_lens,
+                spec_info=forward_batch.spec_info,
+            )
+            self.forward_metadata = DecodeMetadata(self.decode_wrappers)
+        elif forward_batch.forward_mode.is_draft_extend():
+            self.indices_updater_prefill.update(
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                prefix_lens=None,
+                prefill_wrappers=self.prefill_wrappers_paged,
+                use_ragged=False,
+                encoder_lens=forward_batch.encoder_lens,
+                spec_info=forward_batch.spec_info,
+            )
+            self.forward_metadata = PrefillMetadata(
+                self.prefill_wrappers_paged, False, False
+            )
+        elif forward_batch.forward_mode.is_target_verify():
+            self.indices_updater_prefill.update(
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                prefix_lens=None,
+                prefill_wrappers=self.prefill_wrappers_verify,
+                use_ragged=False,
+                encoder_lens=forward_batch.encoder_lens,
+                spec_info=forward_batch.spec_info,
+            )
+            self.forward_metadata = PrefillMetadata(
+                self.prefill_wrappers_verify, False, False
+            )
+        else:
+            prefix_lens = forward_batch.extend_prefix_lens
+
+            if self.is_multimodal:
+                use_ragged = False
+                extend_no_prefix = False
+            else:
+                use_ragged = True
+                extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
+
+            self.indices_updater_prefill.update(
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                prefix_lens,
+                prefill_wrappers=self.prefill_wrappers_paged,
+                use_ragged=use_ragged,
+                encoder_lens=forward_batch.encoder_lens,
+                spec_info=None,
+            )
+            self.forward_metadata = PrefillMetadata(
+                self.prefill_wrappers_paged, use_ragged, extend_no_prefix
+            )
+
+    def init_cuda_graph_state(
+        self, max_bs: int, kv_indices_buf: Optional[torch.Tensor] = None
+    ):
+        if kv_indices_buf is None:
+            cuda_graph_kv_indices = torch.zeros(
+                (max_bs * self.max_context_len,),
+                dtype=torch.int32,
+                device="cuda",
+            )
+        else:
+            cuda_graph_kv_indices = kv_indices_buf
+
+        self.cuda_graph_kv_indices = [cuda_graph_kv_indices] + [
+            cuda_graph_kv_indices.clone() for _ in range(self.num_wrappers - 1)
+        ]
+
+        if not self.skip_prefill:
+            self.cuda_graph_custom_mask = torch.zeros(
+                (max_bs * self.max_context_len),
+                dtype=torch.uint8,
+                device="cuda",
+            )
+            self.cuda_graph_qk_indptr = [x.clone() for x in self.kv_indptr]
+            self.cuda_graph_qo_indptr = [x.clone() for x in self.kv_indptr]
+
+    def init_forward_metadata_capture_cuda_graph(
+        self,
+        bs: int,
+        num_tokens: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        if forward_mode.is_decode_or_idle():
+            decode_wrappers = []
+            for i in range(self.num_wrappers):
+                decode_wrappers.append(
+                    BatchDecodeWithPagedKVCacheWrapper(
+                        self.workspace_buffer,
+                        "NHD",
+                        use_cuda_graph=True,
+                        use_tensor_cores=self.decode_use_tensor_cores,
+                        paged_kv_indptr_buffer=self.kv_indptr[i][: num_tokens + 1],
+                        paged_kv_indices_buffer=self.cuda_graph_kv_indices[i],
+                        paged_kv_last_page_len_buffer=self.kv_last_page_len[
+                            :num_tokens
+                        ],
+                    )
+                )
+            seq_lens_sum = seq_lens.sum().item()
+            self.indices_updater_decode.update(
+                req_pool_indices,
+                seq_lens,
+                seq_lens_sum,
+                decode_wrappers=decode_wrappers,
+                encoder_lens=encoder_lens,
+                spec_info=spec_info,
+            )
+            self.decode_cuda_graph_metadata[bs] = decode_wrappers
+            self.forward_metadata = DecodeMetadata(decode_wrappers)
+            for i in range(self.num_wrappers):
+                decode_wrappers[i].begin_forward = partial(
+                    fast_decode_plan, decode_wrappers[i]
+                )
+        elif forward_mode.is_target_verify():
+            prefill_wrappers = []
+            for i in range(self.num_wrappers):
+                prefill_wrappers.append(
+                    BatchPrefillWithPagedKVCacheWrapper(
+                        self.workspace_buffer,
+                        "NHD",
+                        use_cuda_graph=True,
+                        qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1],
+                        paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1],
+                        paged_kv_indices_buf=self.cuda_graph_kv_indices[i],
+                        paged_kv_last_page_len_buf=self.kv_last_page_len[:bs],
+                        custom_mask_buf=self.cuda_graph_custom_mask,
+                        mask_indptr_buf=self.cuda_graph_qk_indptr[i][: bs + 1],
+                    )
+                )
+            seq_lens_sum = seq_lens.sum().item()
+            self.indices_updater_prefill.update(
+                req_pool_indices,
+                seq_lens,
+                seq_lens_sum,
+                prefix_lens=None,
+                prefill_wrappers=prefill_wrappers,
+                use_ragged=False,
+                encoder_lens=encoder_lens,
+                spec_info=spec_info,
+            )
+            self.prefill_cuda_graph_metadata[bs] = prefill_wrappers
+            self.forward_metadata = PrefillMetadata(prefill_wrappers, False, False)
+        else:
+            raise ValueError(f"Invalid mode: {forward_mode=}")
+
+    def init_forward_metadata_replay_cuda_graph(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        seq_lens_cpu: Optional[torch.Tensor],
+    ):
+        if forward_mode.is_decode_or_idle():
+            self.indices_updater_decode.update(
+                req_pool_indices[:bs],
+                seq_lens[:bs],
+                seq_lens_sum,
+                decode_wrappers=self.decode_cuda_graph_metadata[bs],
+                encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
+                spec_info=spec_info,
+            )
+        elif forward_mode.is_target_verify():
+            self.indices_updater_prefill.update(
+                req_pool_indices[:bs],
+                seq_lens[:bs],
+                seq_lens_sum,
+                prefix_lens=None,
+                prefill_wrappers=self.prefill_cuda_graph_metadata[bs],
+                use_ragged=False,
+                encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
+                spec_info=spec_info,
+            )
+        else:
+            raise ValueError("Invalid forward mode")
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        return 0
+
+    def forward_extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+    ):
+        k_scale = layer.k_scale_float if self.kv_cache_dtype_str != "auto" else None
+        v_scale = layer.v_scale_float if self.kv_cache_dtype_str != "auto" else None
+        prefill_wrapper_paged = self.forward_metadata.prefill_wrappers[
+            self._get_wrapper_idx(layer)
+        ]
+        cache_loc = (
+            forward_batch.out_cache_loc
+            if not layer.is_cross_attention
+            else forward_batch.encoder_out_cache_loc
+        )
+
+        logits_soft_cap = layer.logit_cap
+
+        if not self.forward_metadata.use_ragged:
+            if k is not None:
+                assert v is not None
+                if save_kv_cache:
+                    forward_batch.token_to_kv_pool.set_kv_buffer(
+                        layer, cache_loc, k, v, k_scale, v_scale
+                    )
+
+            o = prefill_wrapper_paged.forward(
+                q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id),
+                causal=not layer.is_cross_attention,
+                sm_scale=layer.scaling,
+                window_left=layer.sliding_window_size,
+                logits_soft_cap=logits_soft_cap,
+                k_scale=k_scale,
+                v_scale=v_scale,
+            )
+        else:
+            o1, s1 = self.prefill_wrapper_ragged.forward_return_lse(
+                q.view(-1, layer.tp_q_head_num, layer.head_dim),
+                k.view(-1, layer.tp_k_head_num, layer.head_dim),
+                v.view(-1, layer.tp_v_head_num, layer.head_dim),
+                causal=True,
+                sm_scale=layer.scaling,
+                logits_soft_cap=logits_soft_cap,
+            )
+
+            if self.forward_metadata.extend_no_prefix:
+                o = o1
+            else:
+                o2, s2 = prefill_wrapper_paged.forward_return_lse(
+                    q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                    forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id),
+                    causal=False,
+                    sm_scale=layer.scaling,
+                    logits_soft_cap=logits_soft_cap,
+                )
+
+                o, _ = merge_state(o1, s1, o2, s2)
+
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer, cache_loc, k, v, k_scale, v_scale
+                )
+
+        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+    ):
+        k_scale = layer.k_scale_float if self.kv_cache_dtype_str != "auto" else None
+        v_scale = layer.v_scale_float if self.kv_cache_dtype_str != "auto" else None
+        decode_wrapper = self.forward_metadata.decode_wrappers[
+            self._get_wrapper_idx(layer)
+        ]
+        cache_loc = (
+            forward_batch.out_cache_loc
+            if not layer.is_cross_attention
+            else forward_batch.encoder_out_cache_loc
+        )
+
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer, cache_loc, k, v, k_scale, v_scale
+                )
+
+        o = decode_wrapper.forward(
+            q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+            forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id),
+            sm_scale=layer.scaling,
+            logits_soft_cap=layer.logit_cap,
+            k_scale=k_scale,
+            v_scale=v_scale,
+        )
+
+        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
+
+    def _get_wrapper_idx(self, layer: RadixAttention):
+        if self.num_wrappers == 1:
+            return 0
+
+        if self.dispatch_reason == WrapperDispatch.SLIDING_WINDOW:
+            return layer.sliding_window_size == -1
+        if self.dispatch_reason == WrapperDispatch.CROSS_ATTENTION:
+            return layer.is_cross_attention
+
+        raise ValueError(f"Unknown dispatch reason: {self.dispatch_reason}")
+
+
+class FlashInferIndicesUpdaterDecode:
+    def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
+        # Parse Constants
+        self.num_qo_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
+            get_attention_tp_size()
+        )
+        self.head_dim = model_runner.model_config.head_dim
+        self.data_type = model_runner.kv_cache_dtype
+        self.q_data_type = model_runner.dtype
+        self.sliding_window_size = model_runner.sliding_window_size
+        self.attn_backend = attn_backend
+
+        # Buffers and wrappers
+        self.kv_indptr = attn_backend.kv_indptr
+        self.kv_last_page_len = attn_backend.kv_last_page_len
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+
+        # Dispatch the update function
+        if self.attn_backend.dispatch_reason == WrapperDispatch.SLIDING_WINDOW:
+            self.update = self.update_sliding_window
+        elif self.attn_backend.dispatch_reason == WrapperDispatch.CROSS_ATTENTION:
+            self.update = self.update_cross_attention
+        else:
+            assert self.attn_backend.num_wrappers == 1
+            self.update = self.update_single_wrapper
+
+    def update(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
+        encoder_lens: Optional[torch.Tensor],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        # Keep the signature for type checking. It will be assigned during runtime.
+        raise NotImplementedError()
+
+    def update_single_wrapper(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
+        encoder_lens: Optional[torch.Tensor],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        decode_wrappers = decode_wrappers or self.decode_wrappers
+        self.call_begin_forward(
+            decode_wrappers[0],
+            req_pool_indices,
+            seq_lens,
+            seq_lens_sum,
+            self.kv_indptr[0],
+            None,
+            spec_info,
+        )
+
+    def update_sliding_window(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
+        encoder_lens: Optional[torch.Tensor],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        for wrapper_id in range(2):
+            if wrapper_id == 0:
+                # Sliding window attention
+                paged_kernel_lens_tmp = torch.minimum(  # TODO: replace this with clamp
+                    seq_lens,
+                    torch.tensor(self.sliding_window_size + 1),
+                )
+                paged_kernel_lens_sum_tmp = paged_kernel_lens_tmp.sum().item()
+                kv_start_idx_tmp = seq_lens - paged_kernel_lens_tmp
+            else:
+                # Full attention
+                paged_kernel_lens_tmp = seq_lens
+                paged_kernel_lens_sum_tmp = seq_lens_sum
+                kv_start_idx_tmp = None
+
+            self.call_begin_forward(
+                decode_wrappers[wrapper_id],
+                req_pool_indices,
+                paged_kernel_lens_tmp,
+                paged_kernel_lens_sum_tmp,
+                self.kv_indptr[wrapper_id],
+                kv_start_idx_tmp,
+                spec_info,
+            )
+
+    def update_cross_attention(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
+        encoder_lens: Optional[torch.Tensor],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        for wrapper_id in range(2):
+            if wrapper_id == 0:
+                # Normal attention
+                paged_kernel_lens = seq_lens
+                kv_start_idx = encoder_lens
+            else:
+                # Cross attention
+                paged_kernel_lens = encoder_lens
+                kv_start_idx = torch.zeros_like(encoder_lens)
+                seq_lens_sum = encoder_lens.sum().item()
+
+            self.call_begin_forward(
+                decode_wrappers[wrapper_id],
+                req_pool_indices,
+                paged_kernel_lens,
+                seq_lens_sum,
+                self.kv_indptr[wrapper_id],
+                kv_start_idx,
+                spec_info,
+            )
+
+    def call_begin_forward(
+        self,
+        wrapper: BatchDecodeWithPagedKVCacheWrapper,
+        req_pool_indices: torch.Tensor,
+        paged_kernel_lens: torch.Tensor,
+        paged_kernel_lens_sum: int,
+        kv_indptr: torch.Tensor,
+        kv_start_idx: torch.Tensor,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        if spec_info is None:
+            bs = len(req_pool_indices)
+            kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
+            kv_indptr = kv_indptr[: bs + 1]
+
+            if wrapper.is_cuda_graph_enabled:
+                # Directly write to the cuda graph input buffer
+                kv_indices = wrapper._paged_kv_indices_buf
+            else:
+                kv_indices = torch.empty(
+                    paged_kernel_lens_sum, dtype=torch.int32, device="cuda"
+                )
+
+            create_flashinfer_kv_indices_triton[(bs,)](
+                self.req_to_token,
+                req_pool_indices,
+                paged_kernel_lens,
+                kv_indptr,
+                kv_start_idx,
+                kv_indices,
+                self.req_to_token.shape[1],
+            )
+        else:
+            kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
+            bs = kv_indptr.shape[0] - 1
+
+        wrapper.begin_forward(
+            kv_indptr,
+            kv_indices,
+            self.kv_last_page_len[:bs],
+            self.num_qo_heads,
+            self.num_kv_heads,
+            self.head_dim,
+            1,
+            data_type=self.data_type,
+            q_data_type=self.q_data_type,
+            non_blocking=True,
+        )
+
+
+class FlashInferIndicesUpdaterPrefill:
+    def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
+        # Parse Constants
+        self.num_qo_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
+            get_attention_tp_size()
+        )
+        self.head_dim = model_runner.model_config.head_dim
+        self.data_type = model_runner.kv_cache_dtype
+        self.q_data_type = model_runner.dtype
+        self.sliding_window_size = model_runner.sliding_window_size
+        self.attn_backend = attn_backend
+
+        # Buffers and wrappers
+        self.kv_indptr = attn_backend.kv_indptr
+        self.kv_last_page_len = attn_backend.kv_last_page_len
+        self.qo_indptr = attn_backend.qo_indptr
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged
+
+        # Dispatch the update function
+        if self.attn_backend.dispatch_reason == WrapperDispatch.SLIDING_WINDOW:
+            self.update = self.update_sliding_window
+        elif self.attn_backend.dispatch_reason == WrapperDispatch.CROSS_ATTENTION:
+            self.update = self.update_cross_attention
+        else:
+            assert self.attn_backend.num_wrappers == 1
+            self.update = self.update_single_wrapper
+
+    def update(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        prefix_lens: torch.Tensor,
+        prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
+        use_ragged: bool,
+        encoder_lens: Optional[torch.Tensor],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        # Keep the signature for type checking. It will be assigned during runtime.
+        raise NotImplementedError()
+
+    def update_single_wrapper(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        prefix_lens: torch.Tensor,
+        prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
+        use_ragged: bool,
+        encoder_lens: Optional[torch.Tensor],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        if use_ragged:
+            paged_kernel_lens = prefix_lens
+            paged_kernel_lens_sum = paged_kernel_lens.sum().item()
+        else:
+            paged_kernel_lens = seq_lens
+            paged_kernel_lens_sum = seq_lens_sum
+
+        self.call_begin_forward(
+            self.prefill_wrapper_ragged,
+            prefill_wrappers[0],
+            req_pool_indices,
+            paged_kernel_lens,
+            paged_kernel_lens_sum,
+            seq_lens,
+            prefix_lens,
+            None,
+            self.kv_indptr[0],
+            self.qo_indptr[0],
+            use_ragged,
+            spec_info,
+        )
+
+    def update_sliding_window(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        prefix_lens: torch.Tensor,
+        prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
+        use_ragged: bool,
+        encoder_lens: Optional[torch.Tensor],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        for wrapper_id in range(2):
+            if wrapper_id == 0:
+                # window attention use paged only
+                paged_kernel_lens = torch.minimum(
+                    seq_lens,
+                    torch.tensor(self.sliding_window_size) + seq_lens - prefix_lens,
+                )
+                paged_kernel_lens_sum = paged_kernel_lens.sum().item()
+            else:
+                # full attention
+                paged_kernel_lens = seq_lens
+                paged_kernel_lens_sum = seq_lens_sum
+
+            kv_start_idx = seq_lens - paged_kernel_lens
+
+            self.call_begin_forward(
+                self.prefill_wrapper_ragged,
+                prefill_wrappers[wrapper_id],
+                req_pool_indices,
+                paged_kernel_lens,
+                paged_kernel_lens_sum,
+                seq_lens,
+                prefix_lens,
+                kv_start_idx,
+                self.kv_indptr[wrapper_id],
+                self.qo_indptr[wrapper_id],
+                use_ragged,
+                spec_info,
+            )
+
+    def update_cross_attention(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        prefix_lens: torch.Tensor,
+        prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
+        use_ragged: bool,
+        encoder_lens: Optional[torch.Tensor],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        for wrapper_id in range(2):
+            if wrapper_id == 0:
+                # normal attention
+                paged_kernel_lens = seq_lens
+                kv_start_idx = encoder_lens
+                paged_kernel_lens_sum = seq_lens_sum
+            else:
+                # cross attention
+                paged_kernel_lens = encoder_lens
+                kv_start_idx = torch.zeros_like(encoder_lens)
+                paged_kernel_lens_sum = paged_kernel_lens.sum().item()
+
+            self.call_begin_forward(
+                self.prefill_wrapper_ragged,
+                prefill_wrappers[wrapper_id],
+                req_pool_indices,
+                paged_kernel_lens,
+                paged_kernel_lens_sum,
+                seq_lens,
+                prefix_lens,
+                kv_start_idx,
+                self.kv_indptr[wrapper_id],
+                self.qo_indptr[wrapper_id],
+                use_ragged,
+                spec_info,
+            )
+
+    def call_begin_forward(
+        self,
+        wrapper_ragged: BatchPrefillWithRaggedKVCacheWrapper,
+        wrapper_paged: BatchPrefillWithPagedKVCacheWrapper,
+        req_pool_indices: torch.Tensor,
+        paged_kernel_lens: torch.Tensor,
+        paged_kernel_lens_sum: int,
+        seq_lens: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        kv_start_idx: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        qo_indptr: torch.Tensor,
+        use_ragged: bool,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        bs = len(seq_lens)
+        if spec_info is None:
+            assert len(seq_lens) == len(req_pool_indices)
+            # Normal extend
+            kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
+            kv_indptr = kv_indptr[: bs + 1]
+            kv_indices = torch.empty(
+                paged_kernel_lens_sum + 256,
+                dtype=torch.int32,
+                device=req_pool_indices.device,
+            )
+            create_flashinfer_kv_indices_triton[(bs,)](
+                self.req_to_token,
+                req_pool_indices,
+                paged_kernel_lens,
+                kv_indptr,
+                kv_start_idx,
+                kv_indices,
+                self.req_to_token.shape[1],
+            )
+            qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
+            qo_indptr = qo_indptr[: bs + 1]
+            custom_mask = None
+        else:
+            assert isinstance(spec_info, EagleDraftInput) or isinstance(
+                spec_info, EagleVerifyInput
+            )
+            kv_indices, kv_indptr, qo_indptr, custom_mask = (
+                spec_info.generate_attn_arg_prefill(
+                    req_pool_indices,
+                    paged_kernel_lens,
+                    paged_kernel_lens_sum,
+                    self.req_to_token,
+                )
+            )
+
+        # extend part
+        if use_ragged:
+            wrapper_ragged.begin_forward(
+                qo_indptr,
+                qo_indptr,
+                self.num_qo_heads,
+                self.num_kv_heads,
+                self.head_dim,
+                q_data_type=self.q_data_type,
+            )
+
+        # cached part
+        wrapper_paged.begin_forward(
+            qo_indptr,
+            kv_indptr,
+            kv_indices,
+            self.kv_last_page_len[:bs],
+            self.num_qo_heads,
+            self.num_kv_heads,
+            self.head_dim,
+            1,
+            q_data_type=self.q_data_type,
+            kv_data_type=self.data_type,
+            custom_mask=custom_mask,
+            non_blocking=True,
+        )
+
+
+# Use as a fast path to override the indptr in flashinfer's plan function
+# This is used to remove some host-to-device copy overhead.
+global global_override_indptr_cpu
+
+
+class FlashInferMultiStepDraftBackend:
+    """
+    Wrap multiple flashinfer attention backends as one for multiple consecutive
+    draft decoding steps.
+    """
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        topk: int,
+        speculative_num_steps: int,
+    ):
+        from sglang.srt.speculative.eagle_utils import generate_draft_decode_kv_indices
+
+        self.topk = topk
+        self.speculative_num_steps = speculative_num_steps
+        self.generate_draft_decode_kv_indices = generate_draft_decode_kv_indices
+        self.page_size = model_runner.page_size
+
+        max_bs = model_runner.req_to_token_pool.size * self.topk
+        self.kv_indptr = torch.zeros(
+            (
+                self.speculative_num_steps,
+                max_bs + 1,
+            ),
+            dtype=torch.int32,
+            device=model_runner.device,
+        )
+        self.kv_last_page_len = torch.ones(
+            (max_bs,), dtype=torch.int32, device=model_runner.device
+        )
+        self.attn_backends = []
+        for i in range(self.speculative_num_steps):
+            self.attn_backends.append(
+                FlashInferAttnBackend(
+                    model_runner,
+                    skip_prefill=True,
+                    kv_indptr_buf=self.kv_indptr[i],
+                    kv_last_page_len_buf=self.kv_last_page_len,
+                )
+            )
+
+        self.max_context_len = self.attn_backends[0].max_context_len
+
+        # Cached variables for generate_draft_decode_kv_indices
+        self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1]
+
+    def common_template(
+        self,
+        forward_batch: ForwardBatch,
+        kv_indices_buffer: torch.Tensor,
+        call_fn: Callable,
+    ):
+        num_seqs = forward_batch.batch_size
+        bs = self.topk * num_seqs
+        seq_lens_sum = forward_batch.seq_lens_sum
+
+        self.generate_draft_decode_kv_indices[
+            (self.speculative_num_steps, num_seqs, self.topk)
+        ](
+            forward_batch.req_pool_indices,
+            forward_batch.req_to_token_pool.req_to_token,
+            forward_batch.seq_lens,
+            kv_indices_buffer,
+            self.kv_indptr,
+            forward_batch.positions,
+            num_seqs,
+            self.topk,
+            self.pool_len,
+            kv_indices_buffer.shape[1],
+            self.kv_indptr.shape[1],
+            next_power_of_2(num_seqs),
+            next_power_of_2(self.speculative_num_steps),
+            next_power_of_2(bs),
+        )
+
+        assert forward_batch.spec_info is not None
+        assert isinstance(forward_batch.spec_info, EagleDraftInput)
+
+        # Copy the kv_indptr once to avoid multiple device-to-host copies in flashinfer's plan.
+        indptr_cpu_whole = self.kv_indptr[:, : bs + 1].cpu()
+        global global_override_indptr_cpu
+
+        for i in range(self.speculative_num_steps - 1):
+            forward_batch.spec_info.kv_indptr = self.kv_indptr[i, : bs + 1]
+            forward_batch.spec_info.kv_indices = kv_indices_buffer[i][
+                : seq_lens_sum * self.topk + bs * (i + 1)
+            ]
+            global_override_indptr_cpu = indptr_cpu_whole[i]
+            call_fn(i, forward_batch)
+
+        global_override_indptr_cpu = None
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        kv_indices = torch.empty(
+            (
+                self.speculative_num_steps,
+                forward_batch.batch_size * self.topk * self.max_context_len,
+            ),
+            dtype=torch.int32,
+            device="cuda",
+        )
+
+        def call_fn(i, forward_batch):
+            forward_batch.spec_info.kv_indptr = (
+                forward_batch.spec_info.kv_indptr.clone()
+            )
+            forward_batch.spec_info.kv_indices = (
+                forward_batch.spec_info.kv_indices.clone()
+            )
+            self.attn_backends[i].init_forward_metadata(forward_batch)
+
+        self.common_template(forward_batch, kv_indices, call_fn)
+
+    def init_cuda_graph_state(self, max_bs: int):
+        self.cuda_graph_kv_indices = torch.zeros(
+            (self.speculative_num_steps, max_bs * self.max_context_len),
+            dtype=torch.int32,
+            device="cuda",
+        )
+
+        for i in range(self.speculative_num_steps):
+            self.attn_backends[i].init_cuda_graph_state(
+                max_bs, kv_indices_buf=self.cuda_graph_kv_indices[i]
+            )
+
+    def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch):
+        def call_fn(i, forward_batch):
+            self.attn_backends[i].init_forward_metadata_capture_cuda_graph(
+                forward_batch.batch_size,
+                forward_batch.batch_size * self.topk,
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                encoder_lens=None,
+                forward_mode=ForwardMode.DECODE,
+                spec_info=forward_batch.spec_info,
+            )
+
+        self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn)
+
+    def init_forward_metadata_replay_cuda_graph(
+        self, forward_batch: ForwardBatch, bs: int
+    ):
+        def call_fn(i, forward_batch):
+            self.attn_backends[i].init_forward_metadata_replay_cuda_graph(
+                bs,
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                seq_lens_sum=-1,
+                encoder_lens=None,
+                forward_mode=ForwardMode.DECODE,
+                spec_info=forward_batch.spec_info,
+                seq_lens_cpu=None,
+            )
+
+        self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn)
+
+
+def should_use_tensor_core(
+    kv_cache_dtype: torch.dtype,
+    num_attention_heads: int,
+    num_kv_heads: int,
+) -> bool:
+    """
+    Determine whether to use tensor cores for attention computation.
+
+    Args:
+        kv_cache_dtype: Data type of the KV cache
+        num_attention_heads: Number of attention heads
+        num_kv_heads: Number of key/value heads
+
+    Returns:
+        bool: Whether to use tensor cores
+    """
+    # Try to use environment variable first
+    env_override = os.environ.get("SGLANG_FLASHINFER_USE_TENSOR_CORE")
+    if env_override is not None:
+        return env_override.lower() == "true"
+
+    # Try to use _grouped_size_compiled_for_decode_kernels if available
+    # This is for flashinfer <=0.1.6. Otherwise, there is an accuracy bug
+    try:
+        from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
+
+        if not _grouped_size_compiled_for_decode_kernels(
+            num_attention_heads,
+            num_kv_heads,
+        ):
+            return True
+        else:
+            return False
+    except (ImportError, AttributeError):
+        pass
+
+    # Calculate GQA group size
+    gqa_group_size = num_attention_heads // num_kv_heads
+
+    # Determine based on dtype and GQA group size
+    if kv_cache_dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+        return True
+    elif kv_cache_dtype in (torch.float16, torch.half, torch.bfloat16):
+        return gqa_group_size > 4
+    else:
+        return False
+
+
+# Use as a fast path to override the indptr in flashinfer's plan function
+# This is used to remove some host-to-device copy overhead.
+global_override_indptr_cpu = None
+
+
+def fast_decode_plan(
+    self,
+    indptr: torch.Tensor,
+    indices: torch.Tensor,
+    last_page_len: torch.Tensor,
+    num_qo_heads: int,
+    num_kv_heads: int,
+    head_dim: int,
+    page_size: int,
+    pos_encoding_mode: str = "NONE",
+    window_left: int = -1,
+    logits_soft_cap: Optional[float] = None,
+    data_type: Union[str, torch.dtype] = "float16",
+    q_data_type: Optional[Union[str, torch.dtype]] = None,
+    sm_scale: Optional[float] = None,
+    rope_scale: Optional[float] = None,
+    rope_theta: Optional[float] = None,
+    non_blocking: bool = True,
+) -> None:
+    """
+    A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for FlashInferMultiStepDraftBackend.
+    Modifications:
+    - Remove unnecessary device-to-device copy for the cuda graph buffers.
+    - Remove unnecessary host-to-device copy for the metadata buffers.
+    """
+    batch_size = len(last_page_len)
+    if logits_soft_cap is None:
+        logits_soft_cap = 0.0
+
+    if self.use_tensor_cores:
+        qo_indptr_host = _get_range_buf(batch_size + 1, "cpu")
+
+    if self.is_cuda_graph_enabled:
+        if batch_size != self._fixed_batch_size:
+            raise ValueError(
+                "The batch size should be fixed in cudagraph mode, the runtime batch size {} "
+                " mismatches the batch size set during initialization {}".format(
+                    batch_size, self._fixed_batch_size
+                )
+            )
+        if len(indices) > len(self._paged_kv_indices_buf):
+            raise ValueError(
+                "The size of indices should be less than or equal to the allocated buffer"
+            )
+        # Skip these copies because we directly write to them during prepartion
+        # self._paged_kv_indptr_buf.copy_(indptr)
+        # self._paged_kv_indices_buf[: len(indices)] = indices
+        # self._paged_kv_last_page_len_buf.copy_(last_page_len)
+    else:
+        self._paged_kv_indptr_buf = indptr
+        self._paged_kv_indices_buf = indices
+        self._paged_kv_last_page_len_buf = last_page_len
+        self._qo_indptr_buf = qo_indptr_host.to(self.device, non_blocking=non_blocking)
+
+    # NOTE(Zihao): the following tensors acts as placeholder to pass dtype info
+    if not q_data_type:
+        q_data_type = data_type
+
+    if not hasattr(self, "empty_q_data"):
+        self.empty_q_data = torch.empty(
+            0,
+            dtype=(
+                getattr(torch, q_data_type)
+                if isinstance(q_data_type, str)
+                else q_data_type
+            ),
+        )
+        self.empty_kv_cache = torch.empty(
+            0,
+            dtype=(
+                getattr(torch, data_type) if isinstance(data_type, str) else data_type
+            ),
+        )
+        self.last_page_len = torch.ones(32768, dtype=torch.int32)
+
+    indptr_host = (
+        global_override_indptr_cpu
+        if global_override_indptr_cpu is not None
+        else indptr.cpu()
+    )
+
+    if self.use_tensor_cores:
+        kv_lens_arr_host = get_seq_lens(
+            indptr_host, self.last_page_len[:batch_size], page_size
+        )
+
+        self._plan_info = self._cached_module.plan(
+            self._float_workspace_buffer,
+            self._int_workspace_buffer,
+            self._pin_memory_int_workspace_buffer,
+            qo_indptr_host,
+            indptr_host,
+            kv_lens_arr_host,
+            batch_size,  # total_num_rows
+            batch_size,
+            num_qo_heads,
+            num_kv_heads,
+            page_size,
+            self.is_cuda_graph_enabled,
+            head_dim,
+            head_dim,
+            False,  # causal
+            torch.cuda.current_stream().cuda_stream,
+        )
+    else:
+        self._plan_info = self._cached_module.plan(
+            self._float_workspace_buffer,
+            self._int_workspace_buffer,
+            self._pin_memory_int_workspace_buffer,
+            indptr_host,
+            batch_size,
+            num_qo_heads,
+            num_kv_heads,
+            page_size,
+            self.is_cuda_graph_enabled,
+            window_left,
+            logits_soft_cap,
+            head_dim,
+            head_dim,
+            self.empty_q_data,
+            self.empty_kv_cache,
+            torch.cuda.current_stream().cuda_stream,
+        )
+
+    self._pos_encoding_mode = pos_encoding_mode
+    self._window_left = window_left
+    self._logits_soft_cap = logits_soft_cap
+    self._sm_scale = sm_scale
+    self._rope_scale = rope_scale
+    self._rope_theta = rope_theta

lib/python3.12/site-packages/sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -0,0 +1,840 @@
+from __future__ import annotations
+
+"""
+Support attention backend for flashinfer MLA.
+The flashinfer_mla_disable_ragged flag controls whether to use ragged prefill wrapper and defaults to be false.
+When it's set to false, all wrappers are BatchMLAPaged wrapper.
+When it's set to true, the backend uses BatchRagged and BatchMLAPaged wrapper for prefilling,
+and uses BatchMLAPaged wrapper for decoding.
+More details can be found in https://docs.flashinfer.ai/api/mla.html
+"""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import TYPE_CHECKING, Callable, Optional, Union
+
+import torch
+import triton
+
+from sglang.global_config import global_config
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.flashinfer_backend import (
+    create_flashinfer_kv_indices_triton,
+)
+from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.managers.schedule_batch import global_server_args_dict
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+from sglang.srt.utils import is_flashinfer_available
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.spec_info import SpecInfo
+
+if is_flashinfer_available():
+    from flashinfer import (
+        BatchMLAPagedAttentionWrapper,
+        BatchPrefillWithRaggedKVCacheWrapper,
+    )
+
+
+@dataclass
+class DecodeMetadata:
+    decode_wrapper: BatchMLAPagedAttentionWrapper
+
+
+@dataclass
+class PrefillMetadata:
+    prefill_wrapper: BatchMLAPagedAttentionWrapper
+    use_ragged: bool
+
+
+# Reuse this workspace buffer across all flashinfer wrappers
+global_workspace_buffer = None
+
+
+class FlashInferMLAAttnBackend(AttentionBackend):
+    """Flashinfer attention kernels."""
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        skip_prefill: bool = False,
+        kv_indptr_buf: Optional[torch.Tensor] = None,
+        q_indptr_decode_buf: Optional[torch.Tensor] = None,
+    ):
+        super().__init__()
+
+        # Parse constants
+        self.max_context_len = model_runner.model_config.context_len
+        self.device = model_runner.device
+        self.skip_prefill = skip_prefill
+
+        # Allocate buffers
+        global global_workspace_buffer
+        if global_workspace_buffer is None:
+            global_workspace_buffer = torch.empty(
+                global_config.flashinfer_workspace_size,
+                dtype=torch.uint8,
+                device=model_runner.device,
+            )
+        self.workspace_buffer = global_workspace_buffer
+
+        max_bs = model_runner.req_to_token_pool.size
+        if kv_indptr_buf is None:
+            self.kv_indptr = torch.zeros(
+                (max_bs + 1,), dtype=torch.int32, device=model_runner.device
+            )
+        else:
+            self.kv_indptr = kv_indptr_buf
+
+        if not self.skip_prefill:
+            self.qo_indptr = torch.zeros(
+                (max_bs + 1,), dtype=torch.int32, device=model_runner.device
+            )
+
+        if q_indptr_decode_buf is None:
+            self.q_indptr_decode = torch.arange(
+                0, max_bs + 1, dtype=torch.int32, device=model_runner.device
+            )
+        else:
+            self.q_indptr_decode = q_indptr_decode_buf
+
+        self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
+            self.workspace_buffer, "NHD"
+        )
+
+        if not self.skip_prefill:
+            self.prefill_wrapper_paged = BatchMLAPagedAttentionWrapper(
+                self.workspace_buffer,
+                backend="auto",
+            )
+
+            # FlashinferMLA backend uses mla wrapper for target verify
+            self.prefill_wrapper_verify = BatchMLAPagedAttentionWrapper(
+                self.workspace_buffer,
+                backend="auto",
+            )
+
+        self.decode_wrapper = BatchMLAPagedAttentionWrapper(
+            self.workspace_buffer, backend="auto"
+        )
+
+        # Create indices updater
+        if not skip_prefill:
+            self.indices_updater_prefill = FlashInferMLAIndicesUpdaterPrefill(
+                model_runner, self
+            )
+
+        self.indices_updater_decode = FlashInferMLAIndicesUpdaterDecode(
+            model_runner, self
+        )
+
+        # Other metadata
+        self.forward_metadata: Union[PrefillMetadata, DecodeMetadata] = None
+        self.decode_cuda_graph_metadata = {}
+        self.prefill_cuda_graph_metadata = {}  # For verify
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        if forward_batch.forward_mode.is_decode_or_idle():
+            self.indices_updater_decode.update(
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                decode_wrapper=self.decode_wrapper,
+                init_metadata_replay=False,
+            )
+            self.forward_metadata = DecodeMetadata(self.decode_wrapper)
+        elif forward_batch.forward_mode.is_draft_extend():
+            self.indices_updater_prefill.update(
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                prefix_lens=None,
+                prefill_wrapper_paged=self.prefill_wrapper_paged,
+                use_ragged=False,
+                spec_info=forward_batch.spec_info,
+            )
+            self.forward_metadata = PrefillMetadata(self.prefill_wrapper_paged, False)
+        elif forward_batch.forward_mode.is_target_verify():
+            self.indices_updater_prefill.update(
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                prefix_lens=None,
+                prefill_wrapper_paged=self.prefill_wrapper_verify,
+                use_ragged=False,
+                spec_info=forward_batch.spec_info,
+            )
+            self.forward_metadata = PrefillMetadata(self.prefill_wrapper_verify, False)
+        else:
+            prefix_lens = forward_batch.extend_prefix_lens
+            extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
+            use_ragged = (
+                not global_server_args_dict["flashinfer_mla_disable_ragged"]
+                and extend_no_prefix
+            )
+
+            self.indices_updater_prefill.update(
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                prefix_lens,
+                prefill_wrapper_paged=self.prefill_wrapper_paged,
+                use_ragged=use_ragged,
+            )
+            self.forward_metadata = PrefillMetadata(
+                self.prefill_wrapper_paged, use_ragged
+            )
+
+    def init_cuda_graph_state(
+        self, max_bs: int, kv_indices_buf: Optional[torch.Tensor] = None
+    ):
+        if kv_indices_buf is None:
+            cuda_graph_kv_indices = torch.zeros(
+                (max_bs * self.max_context_len,),
+                dtype=torch.int32,
+                device="cuda",
+            )
+        else:
+            cuda_graph_kv_indices = kv_indices_buf
+
+        self.cuda_graph_kv_indices = cuda_graph_kv_indices
+        self.cuda_graph_qo_indptr = self.q_indptr_decode.clone()
+        self.cuda_graph_kv_indptr = self.kv_indptr.clone()
+        self.cuda_graph_kv_lens = torch.ones(
+            (max_bs,), dtype=torch.int32, device=self.device
+        )
+
+        # For fast decode plan in graph replaying
+        self.cuda_graph_qo_indptr_cpu = self.cuda_graph_qo_indptr.to("cpu")
+        self.cuda_graph_kv_indptr_cpu = self.cuda_graph_kv_indptr.to("cpu")
+        self.fast_decode_kwargs = {
+            "qo_indptr_cpu": self.cuda_graph_qo_indptr_cpu,
+            "kv_indptr_cpu": self.cuda_graph_kv_indptr_cpu,
+            "kv_indices": self.cuda_graph_kv_indices,
+        }
+
+    def init_forward_metadata_capture_cuda_graph(
+        self,
+        bs: int,
+        num_tokens: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[SpecInfo],
+    ):
+        if forward_mode.is_decode_or_idle():
+            decode_wrapper = BatchMLAPagedAttentionWrapper(
+                self.workspace_buffer,
+                use_cuda_graph=True,
+                qo_indptr=self.cuda_graph_qo_indptr[: num_tokens + 1],
+                kv_indptr=self.cuda_graph_kv_indptr[: num_tokens + 1],
+                kv_indices=self.cuda_graph_kv_indices,
+                kv_len_arr=self.cuda_graph_kv_lens[:num_tokens],
+                backend="auto",
+            )
+
+            seq_lens_sum = seq_lens.sum().item()
+            self.indices_updater_decode.update(
+                req_pool_indices,
+                seq_lens,
+                seq_lens_sum,
+                decode_wrapper=decode_wrapper,
+                init_metadata_replay=False,
+                spec_info=spec_info,
+            )
+            self.decode_cuda_graph_metadata[bs] = decode_wrapper
+            self.forward_metadata = DecodeMetadata(decode_wrapper)
+            decode_wrapper.plan = partial(fast_mla_decode_plan, decode_wrapper)
+        elif forward_mode.is_target_verify():
+            verify_wrapper = BatchMLAPagedAttentionWrapper(
+                self.workspace_buffer,
+                use_cuda_graph=True,
+                qo_indptr=self.cuda_graph_qo_indptr[: bs + 1],
+                kv_indptr=self.cuda_graph_kv_indptr[: bs + 1],
+                kv_indices=self.cuda_graph_kv_indices,
+                kv_len_arr=self.cuda_graph_kv_lens[:bs],
+                backend="auto",
+            )
+            seq_lens_sum = seq_lens.sum().item()
+            self.indices_updater_prefill.update(
+                req_pool_indices,
+                seq_lens,
+                seq_lens_sum,
+                prefix_lens=None,
+                prefill_wrapper_paged=verify_wrapper,
+                use_ragged=False,
+                spec_info=spec_info,
+            )
+            self.prefill_cuda_graph_metadata[bs] = verify_wrapper
+            self.forward_metadata = PrefillMetadata(verify_wrapper, False)
+        else:
+            raise ValueError(f"Invalid mode: {forward_mode=}")
+
+    def init_forward_metadata_replay_cuda_graph(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[SpecInfo],
+        seq_lens_cpu: Optional[torch.Tensor],
+    ):
+        if forward_mode.is_decode_or_idle():
+            assert seq_lens_cpu is not None
+            kv_len_arr_cpu = seq_lens_cpu[:bs]
+            self.cuda_graph_kv_indptr_cpu[1 : bs + 1] = torch.cumsum(
+                kv_len_arr_cpu, dim=0
+            )
+            self.fast_decode_kwargs.update(
+                {
+                    "qo_indptr_cpu": self.cuda_graph_qo_indptr_cpu[: bs + 1],
+                    "kv_indptr_cpu": self.cuda_graph_kv_indptr_cpu[: bs + 1],
+                    "kv_len_arr_cpu": kv_len_arr_cpu,
+                }
+            )
+
+            self.indices_updater_decode.update(
+                req_pool_indices[:bs],
+                seq_lens[:bs],
+                seq_lens_sum,
+                decode_wrapper=self.decode_cuda_graph_metadata[bs],
+                init_metadata_replay=True,
+                spec_info=spec_info,
+                **self.fast_decode_kwargs,
+            )
+        elif forward_mode.is_target_verify():
+            self.indices_updater_prefill.update(
+                req_pool_indices[:bs],
+                seq_lens[:bs],
+                seq_lens_sum,
+                prefix_lens=None,
+                prefill_wrapper_paged=self.prefill_cuda_graph_metadata[bs],
+                use_ragged=False,
+                spec_info=spec_info,
+            )
+        else:
+            raise ValueError(f"Invalid forward mode: {forward_mode=}")
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        return 0
+
+    def forward_extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+    ):
+
+        cache_loc = forward_batch.out_cache_loc
+        logits_soft_cap = layer.logit_cap
+        prefill_wrapper_paged = self.forward_metadata.prefill_wrapper
+        qall = q.view(-1, layer.tp_q_head_num, layer.head_dim)
+        k_buf = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
+
+        # Save kv cache
+        if save_kv_cache and k is not None:
+            assert v is not None
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
+
+        if self.forward_metadata.use_ragged:
+            # ragged prefill
+            o, _ = self.prefill_wrapper_ragged.forward_return_lse(
+                qall,
+                k.view(-1, layer.tp_k_head_num, layer.head_dim),
+                v.view(-1, layer.tp_k_head_num, layer.v_head_dim),
+                causal=True,
+                sm_scale=layer.scaling,
+                logits_soft_cap=logits_soft_cap,
+            )
+        else:
+            # mla paged prefill
+            o = prefill_wrapper_paged.run(
+                qall[:, :, : layer.v_head_dim],
+                qall[:, :, layer.v_head_dim :],
+                k_buf[:, :, : layer.v_head_dim],
+                k_buf[:, :, layer.v_head_dim :],
+            )
+
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+    ):
+        decode_wrapper = self.forward_metadata.decode_wrapper
+        cache_loc = forward_batch.out_cache_loc
+
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer,
+                    cache_loc,
+                    k,
+                    v,
+                )
+        reshaped_q = q.view(-1, layer.tp_q_head_num, layer.head_dim)
+        k_buffer = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
+        reshaped_k = k_buffer.view(-1, 1, layer.head_dim)
+        o = decode_wrapper.run(
+            reshaped_q[:, :, : layer.v_head_dim],
+            reshaped_q[:, :, layer.v_head_dim :],
+            reshaped_k[:, :, : layer.v_head_dim],
+            reshaped_k[:, :, layer.v_head_dim :],
+        )
+
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+
+
+class FlashInferMLAIndicesUpdaterDecode:
+    def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
+        # Parse Constants
+        self.num_local_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.kv_lora_rank = model_runner.model_config.kv_lora_rank
+        self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
+        self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
+        self.scaling = model_runner.model_config.scaling
+        self.data_type = model_runner.kv_cache_dtype
+        self.attn_backend = attn_backend
+
+        # Buffers and wrappers
+        self.kv_indptr = attn_backend.kv_indptr
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.q_indptr = attn_backend.q_indptr_decode
+
+    def update(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        decode_wrapper: BatchMLAPagedAttentionWrapper,
+        init_metadata_replay: bool = False,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]] = None,
+        **fast_decode_kwargs,
+    ):
+        decode_wrapper = decode_wrapper or self.decode_wrapper
+        self.call_begin_forward(
+            decode_wrapper,
+            req_pool_indices,
+            seq_lens,
+            seq_lens_sum,
+            self.q_indptr,
+            self.kv_indptr,
+            init_metadata_replay,
+            spec_info,
+            **fast_decode_kwargs,
+        )
+
+    def call_begin_forward(
+        self,
+        wrapper: BatchMLAPagedAttentionWrapper,
+        req_pool_indices: torch.Tensor,
+        paged_kernel_lens: torch.Tensor,
+        paged_kernel_lens_sum: int,
+        q_indptr: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        init_metadata_replay: bool = False,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]] = None,
+        **fast_decode_kwargs,
+    ):
+        bs = len(req_pool_indices)
+        q_indptr = q_indptr[: bs + 1]
+        kv_lens = paged_kernel_lens.to(torch.int32)
+        sm_scale = self.scaling
+        if spec_info is None:
+            kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
+            kv_indptr = kv_indptr[: bs + 1]
+            kv_indices = (
+                torch.empty(paged_kernel_lens_sum, dtype=torch.int32, device="cuda")
+                if not init_metadata_replay
+                else fast_decode_kwargs["kv_indices"]
+            )
+            create_flashinfer_kv_indices_triton[(bs,)](
+                self.req_to_token,
+                req_pool_indices,
+                paged_kernel_lens,
+                kv_indptr,
+                None,
+                kv_indices,
+                self.req_to_token.shape[1],
+            )
+        else:
+            kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
+
+        if not init_metadata_replay:
+            wrapper.plan(
+                q_indptr,
+                kv_indptr,
+                kv_indices,
+                kv_lens,
+                self.num_local_heads,
+                self.kv_lora_rank,
+                self.qk_rope_head_dim,
+                1,
+                False,
+                sm_scale,
+                self.data_type,
+                self.data_type,
+            )
+        else:
+            wrapper.plan(
+                fast_decode_kwargs["qo_indptr_cpu"],
+                fast_decode_kwargs["kv_indptr_cpu"],
+                kv_indices,
+                fast_decode_kwargs["kv_len_arr_cpu"],
+                self.num_local_heads,
+                self.kv_lora_rank,
+                self.qk_rope_head_dim,
+                1,
+                False,
+                sm_scale,
+                self.data_type,
+                self.data_type,
+            )
+
+
+class FlashInferMLAIndicesUpdaterPrefill:
+    def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
+        # Parse Constants
+        self.num_local_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.kv_lora_rank = model_runner.model_config.kv_lora_rank
+        self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
+        self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
+        self.v_head_dim = model_runner.model_config.v_head_dim
+        self.scaling = model_runner.model_config.scaling
+        self.data_type = model_runner.kv_cache_dtype
+        self.q_data_type = model_runner.dtype
+        self.attn_backend = attn_backend
+
+        # Buffers and wrappers
+        self.kv_indptr = attn_backend.kv_indptr
+        self.qo_indptr = attn_backend.qo_indptr
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged
+
+    def update(
+        self,
+        req_pool_indices: torch.Tnesor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        prefix_lens: torch.Tensor,
+        prefill_wrapper_paged: BatchMLAPagedAttentionWrapper,
+        use_ragged: bool,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]] = None,
+    ):
+        if use_ragged:
+            paged_kernel_lens = prefix_lens
+            paged_kernel_lens_sum = paged_kernel_lens.sum().item()
+        else:
+            paged_kernel_lens = seq_lens
+            paged_kernel_lens_sum = seq_lens_sum
+
+        self.call_begin_forward(
+            self.prefill_wrapper_ragged,
+            prefill_wrapper_paged,
+            req_pool_indices,
+            paged_kernel_lens,
+            paged_kernel_lens_sum,
+            seq_lens,
+            prefix_lens,
+            self.kv_indptr,
+            self.qo_indptr,
+            use_ragged,
+            spec_info,
+        )
+
+    def call_begin_forward(
+        self,
+        wrapper_ragged: BatchPrefillWithRaggedKVCacheWrapper,
+        wrapper_paged: BatchMLAPagedAttentionWrapper,
+        req_pool_indices: torch.Tensor,
+        paged_kernel_lens: torch.Tensor,
+        paged_kernel_lens_sum: int,
+        seq_lens: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        qo_indptr: torch.Tensor,
+        use_ragged: bool,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]] = None,
+    ):
+        bs = len(seq_lens)
+        sm_scale = self.scaling
+
+        if spec_info is None:
+            assert len(seq_lens) == len(req_pool_indices)
+            kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
+            kv_indptr = kv_indptr[: bs + 1]
+            kv_indices = torch.empty(
+                paged_kernel_lens_sum,
+                dtype=torch.int32,
+                device=req_pool_indices.device,
+            )
+            create_flashinfer_kv_indices_triton[(bs,)](
+                self.req_to_token,
+                req_pool_indices,
+                paged_kernel_lens,
+                kv_indptr,
+                None,
+                kv_indices,
+                self.req_to_token.shape[1],
+            )
+            qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
+            qo_indptr = qo_indptr[: bs + 1]
+            custom_mask = None
+        else:
+            assert isinstance(spec_info, EagleDraftInput) or isinstance(
+                spec_info, EagleVerifyInput
+            )
+            # TODO: Support topk > 1 with custom mask
+            kv_indices, kv_indptr, qo_indptr, custom_mask = (
+                spec_info.generate_attn_arg_prefill(
+                    req_pool_indices,
+                    paged_kernel_lens,
+                    paged_kernel_lens_sum,
+                    self.req_to_token,
+                )
+            )
+
+        if use_ragged:
+            # ragged prefill
+            wrapper_ragged.begin_forward(
+                qo_indptr=qo_indptr,
+                kv_indptr=qo_indptr,
+                num_qo_heads=self.num_local_heads,
+                num_kv_heads=self.num_local_heads,
+                head_dim_qk=self.qk_nope_head_dim + self.qk_rope_head_dim,
+                head_dim_vo=self.v_head_dim,
+                q_data_type=self.q_data_type,
+            )
+        else:
+            # mla paged prefill
+            kv_len_arr = kv_indptr[1:] - kv_indptr[:-1]
+            wrapper_paged.plan(
+                qo_indptr,
+                kv_indptr,
+                kv_indices,
+                kv_len_arr,
+                self.num_local_heads,
+                self.kv_lora_rank,
+                self.qk_rope_head_dim,
+                1,
+                True,
+                sm_scale,
+                self.q_data_type,
+                self.data_type,
+            )
+
+
+class FlashInferMLAMultiStepDraftBackend:
+    """
+    Wrap multiple flashinfer mla attention backends as one for multiple consecutive
+    draft decoding steps.
+    """
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        topk: int,
+        speculative_num_steps: int,
+    ):
+        from sglang.srt.speculative.eagle_utils import generate_draft_decode_kv_indices
+
+        if topk > 1:
+            raise ValueError(
+                f"Currently Flashinfer MLA only supports topk=1 for speculative decoding"
+            )
+        self.topk = topk
+        self.speculative_num_steps = speculative_num_steps
+        self.generate_draft_decode_kv_indices = generate_draft_decode_kv_indices
+
+        max_bs = model_runner.req_to_token_pool.size * self.topk
+        self.kv_indptr = torch.zeros(
+            (
+                self.speculative_num_steps,
+                max_bs + 1,
+            ),
+            dtype=torch.int32,
+            device=model_runner.device,
+        )
+        self.q_indptr_decode = torch.arange(
+            0, max_bs + 1, dtype=torch.int32, device=model_runner.device
+        )
+
+        self.attn_backends = []
+        for i in range(self.speculative_num_steps):
+            self.attn_backends.append(
+                FlashInferMLAAttnBackend(
+                    model_runner,
+                    skip_prefill=True,
+                    kv_indptr_buf=self.kv_indptr[i],
+                    q_indptr_decode_buf=self.q_indptr_decode,
+                )
+            )
+
+        self.max_context_len = self.attn_backends[0].max_context_len
+
+        # Cached variables for generate_draft_decode_kv_indices
+        self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1]
+
+    def common_template(
+        self,
+        forward_batch: ForwardBatch,
+        kv_indices_buffer: torch.Tensor,
+        call_fn: Callable,
+    ):
+        num_seqs = forward_batch.batch_size
+        bs = self.topk * num_seqs
+        seq_lens_sum = forward_batch.seq_lens_sum
+
+        self.generate_draft_decode_kv_indices[
+            (self.speculative_num_steps, num_seqs, self.topk)
+        ](
+            forward_batch.req_pool_indices,
+            forward_batch.req_to_token_pool.req_to_token,
+            forward_batch.seq_lens,
+            kv_indices_buffer,
+            self.kv_indptr,
+            forward_batch.positions,
+            num_seqs,
+            self.topk,
+            self.pool_len,
+            kv_indices_buffer.shape[1],
+            self.kv_indptr.shape[1],
+            triton.next_power_of_2(num_seqs),
+            triton.next_power_of_2(self.speculative_num_steps),
+            triton.next_power_of_2(bs),
+        )
+
+        assert forward_batch.spec_info is not None
+        assert isinstance(forward_batch.spec_info, EagleDraftInput)
+
+        for i in range(self.speculative_num_steps - 1):
+            forward_batch.spec_info.kv_indptr = self.kv_indptr[i, : bs + 1]
+            forward_batch.spec_info.kv_indices = kv_indices_buffer[i][
+                : seq_lens_sum * self.topk + bs * (i + 1)
+            ]
+            call_fn(i, forward_batch)
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        kv_indices = torch.zeros(
+            (
+                self.speculative_num_steps,
+                forward_batch.batch_size * self.topk * self.max_context_len,
+            ),
+            dtype=torch.int32,
+            device="cuda",
+        )
+
+        def call_fn(i, forward_batch):
+            assert forward_batch.spec_info is not None
+            assert isinstance(forward_batch.spec_info, EagleDraftInput)
+            forward_batch.spec_info.kv_indptr = (
+                forward_batch.spec_info.kv_indptr.clone()
+            )
+            forward_batch.spec_info.kv_indices = (
+                forward_batch.spec_info.kv_indices.clone()
+            )
+            self.attn_backends[i].init_forward_metadata(forward_batch)
+
+        self.common_template(forward_batch, kv_indices, call_fn)
+
+    def init_cuda_graph_state(self, max_bs: int):
+        self.cuda_graph_kv_indices = torch.zeros(
+            (self.speculative_num_steps, max_bs * self.max_context_len),
+            dtype=torch.int32,
+            device="cuda",
+        )
+
+        for i in range(self.speculative_num_steps):
+            self.attn_backends[i].init_cuda_graph_state(
+                max_bs, kv_indices_buf=self.cuda_graph_kv_indices[i]
+            )
+
+    def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch):
+        def call_fn(i, forward_batch):
+            self.attn_backends[i].init_forward_metadata_capture_cuda_graph(
+                forward_batch.batch_size,
+                forward_batch.batch_size * self.topk,
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                encoder_lens=None,
+                forward_mode=ForwardMode.DECODE,
+                spec_info=forward_batch.spec_info,
+            )
+
+        self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn)
+
+    def init_forward_metadata_replay_cuda_graph(
+        self, forward_batch: ForwardBatch, bs: int
+    ):
+        def call_fn(i, forward_batch):
+            self.attn_backends[i].init_forward_metadata_replay_cuda_graph(
+                bs,
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                seq_lens_sum=-1,
+                encoder_lens=None,
+                forward_mode=ForwardMode.DECODE,
+                spec_info=forward_batch.spec_info,
+                seq_lens_cpu=forward_batch.seq_lens_cpu,
+            )
+
+        self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn)
+
+
+def fast_mla_decode_plan(
+    self,
+    qo_indptr_cpu: torch.Tensor,
+    kv_indptr_cpu: torch.Tensor,
+    kv_indices: torch.Tensor,
+    kv_len_arr_cpu: torch.Tensor,
+    num_heads: int,
+    head_dim_ckv: int,
+    head_dim_kpe: int,
+    page_size: int,
+    causal: bool,
+    sm_scale: float,
+    q_data_type: torch.dtype,
+    kv_data_type: torch.dtype,
+) -> None:
+    """A faster version of BatchMLAPagedAttentionWrapper::plan,
+    for skipping the stream synchronization in original plan function during
+    cuda graph replaying.
+    """
+    self._causal = causal
+    self._page_size = page_size
+    self._sm_scale = sm_scale
+
+    with self.device as device:
+        stream = torch.cuda.current_stream(device).cuda_stream
+        self._cached_module.plan(
+            self._float_workspace_buffer,
+            self._int_workspace_buffer,
+            self._pin_memory_int_workspace_buffer,
+            qo_indptr_cpu,
+            kv_indptr_cpu,
+            kv_len_arr_cpu,
+            num_heads,
+            head_dim_ckv,
+            causal,
+            stream,
+        )

lib/python3.12/site-packages/sglang/srt/layers/attention/flashmla_backend.py

@@ -0,0 +1,281 @@
+from __future__ import annotations
+
+"""
+Support attention backend for FlashMLA.
+
+#TODO
+Enable speculative sampling in FlashMLA
+"""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+import triton
+from flash_mla import flash_mla_with_kvcache, get_mla_metadata
+
+from sglang.global_config import global_config
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend
+from sglang.srt.layers.attention.utils import create_flashmla_kv_indices_triton
+from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+    from sglang.srt.speculative.spec_info import SpecInfo
+
+
+# FlashMLA only supports pagesize=64
+PAGE_SIZE = 64
+# TODO The current setup is hard-coded and will be changed after integrating with MTP.
+Q_LEN = 1
+
+
+@dataclass
+class FlashMLADecodeMetadata:
+    flashmla_metadata: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+    num_splits: Optional[torch.Tensor] = None
+    block_kv_indices: Optional[torch.Tensor] = None
+
+    def __init__(
+        self,
+        flashmla_metadata: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        num_splits: Optional[torch.Tensor] = None,
+        block_kv_indices: Optional[torch.Tensor] = None,
+    ):
+        self.flashmla_metadata = flashmla_metadata
+        self.num_splits = num_splits
+        self.block_kv_indices = block_kv_indices
+
+
+class FlashMLABackend(FlashInferMLAAttnBackend):
+    """Flashinfer attention kernels."""
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        skip_prefill: bool = False,
+        kv_indptr_buf: Optional[torch.Tensor] = None,
+        kv_last_page_len_buf: Optional[torch.Tensor] = None,
+    ):
+        super().__init__(
+            model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf
+        )
+
+        self.num_q_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.num_local_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.forward_metadata: Union[FlashMLADecodeMetadata] = None
+        self.kv_lora_rank = model_runner.model_config.kv_lora_rank
+        self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
+        self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
+        self.v_head_dim = model_runner.model_config.v_head_dim
+        self.scaling = model_runner.model_config.scaling
+        self.data_type = model_runner.kv_cache_dtype
+        self.q_data_type = model_runner.dtype
+        self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+
+        bs = forward_batch.batch_size
+        spec_info = forward_batch.spec_info
+        if forward_batch.forward_mode.is_decode_or_idle():
+            if spec_info is None:
+                max_seqlen_pad = triton.cdiv(
+                    forward_batch.seq_lens_cpu.max().item(), PAGE_SIZE
+                )
+                block_kv_indices = torch.full(
+                    (bs, max_seqlen_pad),
+                    -1,
+                    dtype=torch.int32,
+                    device=forward_batch.seq_lens.device,
+                )
+                create_flashmla_kv_indices_triton[(bs,)](
+                    self.req_to_token,
+                    forward_batch.req_pool_indices,
+                    forward_batch.seq_lens,
+                    None,
+                    block_kv_indices,
+                    self.req_to_token.stride(0),
+                    max_seqlen_pad,
+                )
+                mla_metadata, num_splits = get_mla_metadata(
+                    forward_batch.seq_lens.to(torch.int32),
+                    Q_LEN * self.num_q_heads,
+                    1,
+                )
+                self.forward_metadata = FlashMLADecodeMetadata(
+                    mla_metadata,
+                    num_splits,
+                    block_kv_indices,
+                )
+            else:
+                super().init_forward_metadata(forward_batch)
+        else:
+            super().init_forward_metadata(forward_batch)
+
+    def init_cuda_graph_state(
+        self,
+        max_bs: int,
+        block_kv_indices: Optional[torch.Tensor] = None,
+    ):
+        if block_kv_indices is None:
+            cuda_graph_kv_indices = torch.full(
+                (max_bs, (self.max_context_len + PAGE_SIZE) // PAGE_SIZE),
+                1,
+                dtype=torch.int32,
+                device="cuda",
+            )
+        else:
+            cuda_graph_kv_indices = block_kv_indices
+
+        self.cuda_graph_mla_metadata, self.cuda_graph_num_splits = get_mla_metadata(
+            torch.ones(max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device),
+            Q_LEN * self.num_q_heads,
+            1,
+        )
+        self.cuda_graph_kv_indices = cuda_graph_kv_indices
+
+    def init_forward_metadata_capture_cuda_graph(
+        self,
+        bs: int,
+        num_tokens: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[SpecInfo],
+    ):
+        if forward_mode.is_decode_or_idle():
+            if spec_info is None:
+                max_seqlen_pad = triton.cdiv(seq_lens.max().item(), PAGE_SIZE)
+
+                create_flashmla_kv_indices_triton[(bs,)](
+                    self.req_to_token,
+                    req_pool_indices,
+                    seq_lens,
+                    None,
+                    self.cuda_graph_kv_indices,
+                    self.req_to_token.stride(0),
+                    self.cuda_graph_kv_indices.stride(0),
+                )
+                mla_metadata, num_splits = get_mla_metadata(
+                    seq_lens.to(torch.int32),
+                    Q_LEN * self.num_q_heads,
+                    1,
+                )
+                self.cuda_graph_mla_metadata.copy_(mla_metadata)
+                self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)
+                self.forward_metadata = FlashMLADecodeMetadata(
+                    self.cuda_graph_mla_metadata,
+                    self.cuda_graph_num_splits[: bs + 1],
+                    self.cuda_graph_kv_indices[:bs, :max_seqlen_pad],
+                )
+
+        else:
+            super().init_forward_metadata_capture_cuda_graph(
+                bs,
+                num_tokens,
+                req_pool_indices,
+                seq_lens,
+                encoder_lens,
+                forward_mode,
+                spec_info,
+            )
+
+    def init_forward_metadata_replay_cuda_graph(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[SpecInfo],
+        seq_lens_cpu: Optional[torch.Tensor],
+    ):
+
+        if forward_mode.is_decode_or_idle():
+            assert seq_lens_cpu is not None
+            seq_lens = seq_lens[:bs]
+            seq_lens_cpu = seq_lens_cpu[:bs]
+            max_seqlen_pad = triton.cdiv(seq_lens_cpu.max().item(), PAGE_SIZE)
+            create_flashmla_kv_indices_triton[(bs,)](
+                self.req_to_token,
+                req_pool_indices[:bs],
+                seq_lens,
+                None,
+                self.cuda_graph_kv_indices,
+                self.req_to_token.stride(0),
+                self.cuda_graph_kv_indices.stride(0),
+            )
+            mla_metadata, num_splits = get_mla_metadata(
+                seq_lens.to(torch.int32),
+                Q_LEN * self.num_q_heads,
+                1,
+            )
+            self.cuda_graph_mla_metadata.copy_(mla_metadata)
+            self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)
+            self.forward_metadata.mla_metadata = self.cuda_graph_mla_metadata
+            self.forward_metadata.num_splits = self.cuda_graph_num_splits[: bs + 1]
+            self.forward_metadata.block_kv_indices = self.cuda_graph_kv_indices[
+                :bs, :max_seqlen_pad
+            ]
+
+        else:
+            super().init_forward_metadata_replay_cuda_graph(
+                bs,
+                req_pool_indices,
+                seq_lens,
+                seq_lens_sum,
+                encoder_lens,
+                forward_mode,
+                spec_info,
+                seq_lens_cpu,
+            )
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+    ):
+        cache_loc = forward_batch.out_cache_loc
+
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer,
+                    cache_loc,
+                    k,
+                    v,
+                )
+        bs = forward_batch.batch_size
+        k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
+
+        reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim)
+
+        o, _ = flash_mla_with_kvcache(
+            q=reshape_q,
+            k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim),
+            block_table=self.forward_metadata.block_kv_indices,
+            cache_seqlens=forward_batch.seq_lens.to(torch.int32),
+            head_dim_v=self.kv_lora_rank,  # TODO Retrieve from config.
+            tile_scheduler_metadata=self.forward_metadata.flashmla_metadata,
+            num_splits=self.forward_metadata.num_splits,
+            softmax_scale=layer.scaling,
+            causal=False,
+        )
+
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)

lib/python3.12/site-packages/sglang/srt/layers/attention/torch_native_backend.py

@@ -0,0 +1,267 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from torch.nn.functional import scaled_dot_product_attention
+
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.radix_attention import AttentionType
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+
+class TorchNativeAttnBackend(AttentionBackend):
+    def __init__(self, model_runner: ModelRunner):
+        super().__init__()
+        self.forward_metadata = None
+        self.device = model_runner.device
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        """Init the metadata for a forward pass."""
+        pass
+
+    def _run_sdpa_forward_extend(
+        self,
+        query: torch.Tensor,
+        output: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        extend_prefix_lens: torch.Tensor,
+        extend_seq_lens: torch.Tensor,
+        scaling=None,
+        enable_gqa=False,
+        causal=False,
+    ):
+        """Run the extend forward by using torch native sdpa op.
+
+        Args:
+            query: [num_tokens, num_heads, head_size]
+            output: [num_tokens, num_heads, head_size]
+            k_cache: [max_total_num_tokens, num_heads, head_size]
+            v_cache: [max_total_num_tokens, num_heads, head_size]
+            req_to_token: [max_num_reqs, max_context_len]
+            req_pool_indices: [num_seqs]
+            seq_lens: [num_seqs]
+            extend_prefix_lens: [num_seqs]
+            extend_seq_lens: [num_seqs]
+            scaling: float or None
+            enable_gqa: bool
+            causal: bool
+
+        Returns:
+            output: [num_tokens, num_heads, head_size]
+        """
+
+        assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
+        assert seq_lens.shape[0] == extend_seq_lens.shape[0]
+
+        # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
+        query = query.movedim(0, query.dim() - 2)
+
+        start_q, start_kv = 0, 0
+        for seq_idx in range(seq_lens.shape[0]):
+            # TODO: this loop process a sequence per iter, this is inefficient.
+            # Need optimize the performance later.
+
+            extend_seq_len_q = extend_seq_lens[seq_idx]
+            prefill_seq_len_q = extend_prefix_lens[seq_idx]
+
+            seq_len_kv = seq_lens[seq_idx]
+            end_q = start_q + extend_seq_len_q
+            end_kv = start_kv + seq_len_kv
+
+            per_req_query = query[:, start_q:end_q, :]
+            per_req_query_redudant = torch.empty(
+                (per_req_query.shape[0], seq_len_kv, per_req_query.shape[2]),
+                dtype=per_req_query.dtype,
+                device=per_req_query.device,
+            )
+
+            per_req_query_redudant[:, prefill_seq_len_q:, :] = per_req_query
+
+            # get key and value from cache. per_req_tokens contains the kv cache
+            # index for each token in the sequence.
+            req_pool_idx = req_pool_indices[seq_idx]
+            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
+
+            per_req_out_redudant = (
+                scaled_dot_product_attention(
+                    per_req_query_redudant.unsqueeze(0),
+                    per_req_key.unsqueeze(0),
+                    per_req_value.unsqueeze(0),
+                    enable_gqa=enable_gqa,
+                    scale=scaling,
+                    is_causal=causal,
+                )
+                .squeeze(0)
+                .movedim(query.dim() - 2, 0)
+            )
+            output[start_q:end_q, :, :] = per_req_out_redudant[prefill_seq_len_q:, :, :]
+            start_q, start_kv = end_q, end_kv
+        return output
+
+    def _run_sdpa_forward_decode(
+        self,
+        query: torch.Tensor,
+        output: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        scaling=None,
+        enable_gqa=False,
+        causal=False,
+    ):
+        """Run the decode forward by using torch native sdpa op.
+
+        Args:
+            query: [num_tokens, num_heads, head_size]
+            output: [num_tokens, num_heads, head_size]
+            k_cache: [max_total_num_tokens, num_heads, head_size]
+            v_cache: [max_total_num_tokens, num_heads, head_size]
+            req_to_token: [max_num_reqs, max_context_len]
+            req_pool_indices: [num_seqs]
+            seq_lens: [num_seqs]
+            scaling: float or None
+            enable_gqa: bool
+            causal: bool
+
+        Returns:
+            output: [num_tokens, num_heads, head_size]
+        """
+
+        # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
+        query = query.movedim(0, query.dim() - 2)
+
+        start_q, start_kv = 0, 0
+        for seq_idx in range(seq_lens.shape[0]):
+            # TODO: this loop process a sequence per iter, this is inefficient.
+            # Need optimize the performance later.
+
+            seq_len_q = 1
+            seq_len_kv = seq_lens[seq_idx]
+            end_q = start_q + seq_len_q
+            end_kv = start_kv + seq_len_kv
+
+            per_req_query = query[:, start_q:end_q, :]
+
+            # get key and value from cache. per_req_tokens contains the kv cache
+            # index for each token in the sequence.
+            req_pool_idx = req_pool_indices[seq_idx]
+            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
+
+            per_req_out = (
+                scaled_dot_product_attention(
+                    per_req_query.unsqueeze(0),
+                    per_req_key.unsqueeze(0),
+                    per_req_value.unsqueeze(0),
+                    enable_gqa=enable_gqa,
+                    scale=scaling,
+                    is_causal=causal,
+                )
+                .squeeze(0)
+                .movedim(query.dim() - 2, 0)
+            )
+            output[start_q:end_q, :, :] = per_req_out
+            start_q, start_kv = end_q, end_kv
+
+        return output
+
+    def forward_extend(
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+    ):
+        if layer.qk_head_dim != layer.v_head_dim:
+            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
+        else:
+            o = torch.empty_like(q)
+
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v
+            )
+
+        use_gqa = layer.tp_q_head_num != layer.tp_k_head_num
+
+        q_ = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim)
+        o_ = o.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+
+        causal = True
+        if layer.is_cross_attention or layer.attn_type == AttentionType.ENCODER_ONLY:
+            causal = False
+
+        self._run_sdpa_forward_extend(
+            q_,
+            o_,
+            forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
+            forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
+            forward_batch.req_to_token_pool.req_to_token,
+            forward_batch.req_pool_indices,
+            forward_batch.seq_lens,
+            forward_batch.extend_prefix_lens,
+            forward_batch.extend_seq_lens,
+            scaling=layer.scaling,
+            enable_gqa=use_gqa,
+            causal=causal,
+        )
+        return o
+
+    def forward_decode(
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+    ):
+        # During torch.compile, there is a bug in rotary_emb that causes the
+        # output value to have a 3D tensor shape. This reshapes the output correctly.
+        q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
+
+        if layer.qk_head_dim != layer.v_head_dim:
+            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
+        else:
+            o = torch.empty_like(q)
+
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v
+            )
+
+        use_gqa = layer.tp_q_head_num != layer.tp_k_head_num
+
+        q_ = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim)
+        o_ = o.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+
+        self._run_sdpa_forward_decode(
+            q_,
+            o_,
+            forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
+            forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
+            forward_batch.req_to_token_pool.req_to_token,
+            forward_batch.req_pool_indices,
+            forward_batch.seq_lens,
+            scaling=layer.scaling,
+            enable_gqa=use_gqa,
+            causal=False,
+        )
+
+        return o

lib/python3.12/site-packages/sglang/srt/layers/attention/triton_backend.py

@@ -0,0 +1,733 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
+from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.layers.radix_attention import AttentionType
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.utils import get_bool_env_var, get_device_core_count
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+
+
+@triton.jit
+def get_num_kv_splits_triton(
+    num_kv_splits_ptr,
+    seq_lens_ptr,
+    num_seq,
+    num_group,
+    num_head,
+    num_kv_head,
+    max_kv_splits,
+    device_core_count,
+    MAX_NUM_SEQ: tl.constexpr,
+):
+    # TODO: this method is tunable, we need more online serving data to tune it
+    offs_seq = tl.arange(0, MAX_NUM_SEQ)
+    mask_seq = offs_seq < num_seq
+
+    seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=0)
+    max_seq_len = tl.max(seq_lens)
+    seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=max_seq_len)
+    min_seq_len = tl.min(seq_lens)
+    if max_seq_len * 8 < min_seq_len * 10:
+        min_seq_len = max_seq_len
+    max_kv_splits_1 = tl.minimum(tl.cdiv(max_seq_len, min_seq_len), max_kv_splits)
+    kv_chunk_size_1 = tl.cdiv(max_seq_len, max_kv_splits_1)
+
+    # NOTE: this is a hack to let num_kv_split grows up with seqlen gradually
+    ext_seq_len = tl.cast(max_seq_len, tl.float32) / 64.0
+    ext_device_core_count = tl.cast(
+        device_core_count * tl.maximum(tl.log2(ext_seq_len), 1.0), tl.int32
+    )
+    block_h, num_kv_group = 16, num_head // num_kv_head
+    if num_kv_group == 1:
+        token_grid = num_seq * num_group * num_head
+    else:
+        # from triton_ops/decode_attention.py:_decode_grouped_att_m_fwd
+        block_h = tl.minimum(block_h, num_kv_group)
+        token_grid = num_seq * num_group * tl.cdiv(num_head, block_h)
+    max_kv_splits_2 = tl.minimum(
+        tl.cdiv(ext_device_core_count, token_grid), max_kv_splits
+    )
+    kv_chunk_size_2 = tl.cdiv(max_seq_len, max_kv_splits_2)
+
+    num_kv_splits = tl.maximum(
+        tl.cdiv(seq_lens, kv_chunk_size_1), tl.cdiv(seq_lens, kv_chunk_size_2)
+    )
+
+    offs_token = offs_seq * num_group
+    mask_token = offs_token < num_seq * num_group
+    for i in range(0, num_group):
+        tl.store(num_kv_splits_ptr + i + offs_token, num_kv_splits, mask=mask_token)
+
+
+@dataclass
+class ForwardMetadata:
+    attn_logits: torch.Tensor
+    attn_lse: torch.Tensor
+    max_extend_len: int
+    num_kv_splits: torch.Tensor
+    kv_indptr: torch.Tensor
+    kv_indices: torch.Tensor
+    qo_indptr: torch.Tensor
+    custom_mask: torch.Tensor
+    mask_indptr: torch.Tensor
+
+
+class TritonAttnBackend(AttentionBackend):
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        skip_prefill: bool = False,
+        kv_indptr_buf: Optional[torch.Tensor] = None,
+    ):
+        # Lazy import to avoid the initialization of cuda context
+        from sglang.srt.layers.attention.triton_ops.decode_attention import (
+            decode_attention_fwd,
+        )
+        from sglang.srt.layers.attention.triton_ops.extend_attention import (
+            extend_attention_fwd,
+        )
+
+        super().__init__()
+
+        self.decode_attention_fwd = decode_attention_fwd
+        self.extend_attention_fwd = extend_attention_fwd
+
+        self.skip_prefill = skip_prefill
+
+        max_bs = model_runner.req_to_token_pool.size
+
+        if kv_indptr_buf is None:
+            self.kv_indptr = torch.zeros(
+                (max_bs + 1,), dtype=torch.int32, device=model_runner.device
+            )
+        else:
+            self.kv_indptr = kv_indptr_buf
+
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+
+        if not self.skip_prefill:
+            self.qo_indptr = torch.zeros(
+                (max_bs + 1,), dtype=torch.int32, device=model_runner.device
+            )
+
+            self.mask_indptr = torch.zeros(
+                (max_bs + 1,), dtype=torch.int64, device=model_runner.device
+            )
+
+        self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens
+
+        self.num_head = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.num_kv_head = model_runner.model_config.get_num_kv_heads(
+            get_attention_tp_size()
+        )
+
+        self.static_kv_splits = get_bool_env_var(
+            "SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS", "false"
+        )
+        self.max_kv_splits = model_runner.server_args.triton_attention_num_kv_splits
+        self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[-1]
+
+        self.forward_metadata: ForwardMetadata = None
+
+        self.max_context_len = model_runner.model_config.context_len
+
+        self.device = model_runner.device
+        self.device_core_count = get_device_core_count(model_runner.gpu_id)
+
+    def get_num_kv_splits(
+        self,
+        num_kv_splits: torch.Tensor,
+        seq_lens: torch.Tensor,
+    ):
+        num_token, num_seq = num_kv_splits.shape[0], seq_lens.shape[0]
+        num_group = num_token // num_seq
+
+        assert (
+            num_group * num_seq == num_token
+        ), f"num_seq({num_seq}), num_token({num_token}), something goes wrong!"
+
+        if self.static_kv_splits or self.device_core_count <= 0:
+            num_kv_splits.fill_(self.max_kv_splits)
+            return
+
+        if num_seq < 256:
+            SCHEDULE_SEQ = 256
+        else:
+            SCHEDULE_SEQ = triton.next_power_of_2(num_seq)
+
+        get_num_kv_splits_triton[(1,)](
+            num_kv_splits,
+            seq_lens,
+            num_seq,
+            num_group,
+            self.num_head,
+            self.num_kv_head,
+            self.max_kv_splits,
+            self.device_core_count,
+            MAX_NUM_SEQ=SCHEDULE_SEQ,
+        )
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        """Init auxiliary variables for triton attention backend."""
+
+        bs = forward_batch.batch_size
+        kv_indptr = self.kv_indptr
+        spec_info = forward_batch.spec_info
+
+        if forward_batch.forward_mode.is_decode_or_idle():
+            if spec_info is None:
+                kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0)
+                kv_indptr = kv_indptr[: bs + 1]
+                kv_indices = torch.empty(
+                    forward_batch.seq_lens_sum, dtype=torch.int32, device=self.device
+                )
+                create_flashinfer_kv_indices_triton[(bs,)](
+                    self.req_to_token,
+                    forward_batch.req_pool_indices,
+                    forward_batch.seq_lens,
+                    kv_indptr,
+                    None,
+                    kv_indices,
+                    self.req_to_token.stride(0),
+                )
+            else:
+                kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
+                bs = kv_indptr.shape[0] - 1
+
+            attn_logits = torch.empty(
+                (bs, self.num_head, self.max_kv_splits, self.v_head_dim),
+                dtype=torch.float32,
+                device=self.device,
+            )
+            attn_lse = torch.empty(
+                (bs, self.num_head, self.max_kv_splits),
+                dtype=torch.float32,
+                device=self.device,
+            )
+            num_kv_splits = torch.empty((bs,), dtype=torch.int32, device=self.device)
+
+            self.get_num_kv_splits(num_kv_splits, forward_batch.seq_lens)
+
+            qo_indptr = None
+            custom_mask = None
+            mask_indptr = None
+            max_extend_len = None
+        elif forward_batch.forward_mode.is_target_verify():
+            bs = len(forward_batch.req_pool_indices)
+            qo_indptr = torch.arange(
+                0,
+                (1 + bs) * self.num_draft_tokens,
+                step=self.num_draft_tokens,
+                dtype=torch.int32,
+                device=self.device,
+            )
+            # Different with flashinfer kv_indptr and kv_indices construction
+            kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0)
+            kv_indptr = kv_indptr[: bs + 1]
+            kv_indices = torch.empty(
+                kv_indptr[-1], dtype=torch.int32, device=self.device
+            )
+            create_flashinfer_kv_indices_triton[(bs,)](
+                self.req_to_token,
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                kv_indptr,
+                None,
+                kv_indices,
+                self.req_to_token.stride(0),
+            )
+
+            custom_mask = spec_info.custom_mask
+            seq_mask_len = self.num_draft_tokens * (
+                forward_batch.seq_lens + self.num_draft_tokens
+            )
+            mask_indptr = self.mask_indptr
+            mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len[:bs], dim=0)
+            mask_indptr = mask_indptr[: bs + 1]
+            max_extend_len = self.num_draft_tokens
+            num_kv_splits = None
+            attn_logits = None
+            attn_lse = None
+        elif forward_batch.forward_mode.is_draft_extend():
+            kv_indices, kv_indptr, qo_indptr, custom_mask = (
+                spec_info.generate_attn_arg_prefill(
+                    forward_batch.req_pool_indices,
+                    forward_batch.seq_lens,
+                    None,
+                    self.req_to_token,
+                )
+            )
+            mask_indptr = None
+            # TODO(FIXME): This will trigger an invalid Eagle tree when using
+            # `max(spec_info.accept_length_cpu)`.
+            # It might have been forgotten to update somewhere.
+            max_extend_len = torch.max(spec_info.accept_length).item()
+            num_kv_splits = None
+            attn_logits = None
+            attn_lse = None
+        else:
+            kv_indptr[1 : bs + 1] = torch.cumsum(
+                forward_batch.extend_prefix_lens, dim=0
+            )
+            kv_indptr = kv_indptr[: bs + 1]
+            kv_indices = torch.empty(
+                forward_batch.extend_prefix_lens.sum().item(),
+                dtype=torch.int32,
+                device=self.device,
+            )
+            create_flashinfer_kv_indices_triton[(bs,)](
+                self.req_to_token,
+                forward_batch.req_pool_indices,
+                forward_batch.extend_prefix_lens,
+                kv_indptr,
+                None,
+                kv_indices,
+                self.req_to_token.stride(0),
+            )
+
+            qo_indptr = self.qo_indptr
+            qo_indptr[1 : bs + 1] = torch.cumsum(forward_batch.extend_seq_lens, dim=0)
+            qo_indptr = qo_indptr[: bs + 1]
+            custom_mask = None
+            mask_indptr = None
+            attn_logits = None
+            attn_lse = None
+            max_extend_len = torch.max(forward_batch.extend_seq_lens).item()
+            num_kv_splits = None
+
+        self.forward_metadata = ForwardMetadata(
+            attn_logits,
+            attn_lse,
+            max_extend_len,
+            num_kv_splits,
+            kv_indptr,
+            kv_indices,
+            qo_indptr,
+            custom_mask,
+            mask_indptr,
+        )
+
+    def init_cuda_graph_state(
+        self, max_bs: int, kv_indices_buf: Optional[torch.Tensor] = None
+    ):
+        self.cuda_graph_attn_logits = torch.zeros(
+            (max_bs, self.num_head, self.max_kv_splits, self.v_head_dim),
+            dtype=torch.float32,
+            device=self.device,
+        )
+        self.cuda_graph_attn_lse = torch.zeros(
+            (max_bs, self.num_head, self.max_kv_splits),
+            dtype=torch.float32,
+            device=self.device,
+        )
+        self.cuda_graph_num_kv_splits = torch.full(
+            (max_bs,), self.max_kv_splits, dtype=torch.int32, device=self.device
+        )
+        if kv_indices_buf is None:
+            self.cuda_graph_kv_indices = torch.zeros(
+                (max_bs * self.max_context_len),
+                dtype=torch.int32,
+                device=self.device,
+            )
+        else:
+            self.cuda_graph_kv_indices = kv_indices_buf
+
+        if not self.skip_prefill:
+            self.cuda_graph_custom_mask = torch.zeros(
+                (max_bs * self.max_context_len),
+                dtype=torch.uint8,
+                device=self.device,
+            )
+
+    def init_forward_metadata_capture_cuda_graph(
+        self,
+        bs: int,
+        num_tokens: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        assert encoder_lens is None, "Not supported"
+
+        if forward_mode.is_decode_or_idle():
+            if spec_info is None:
+                kv_indptr = self.kv_indptr
+                kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
+                kv_indptr = kv_indptr[: bs + 1]
+                kv_indices = self.cuda_graph_kv_indices
+                create_flashinfer_kv_indices_triton[(bs,)](
+                    self.req_to_token,
+                    req_pool_indices,
+                    seq_lens,
+                    kv_indptr,
+                    None,
+                    kv_indices,
+                    self.req_to_token.stride(0),
+                )
+            else:
+                kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
+
+            attn_logits = self.cuda_graph_attn_logits
+            attn_lse = self.cuda_graph_attn_lse
+            max_extend_len = None
+            num_kv_splits = self.cuda_graph_num_kv_splits
+            qo_indptr = None
+            custom_mask = None
+            mask_indptr = None
+        elif forward_mode.is_target_verify():
+            qo_indptr = self.qo_indptr[: bs + 1]
+            qo_indptr[: bs + 1] = torch.arange(
+                0,
+                (1 + bs) * self.num_draft_tokens,
+                step=self.num_draft_tokens,
+                dtype=torch.int32,
+                device=self.device,
+            )
+            kv_indptr = self.kv_indptr[: bs + 1]
+            kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
+            kv_indices = self.cuda_graph_kv_indices
+            create_flashinfer_kv_indices_triton[(bs,)](
+                self.req_to_token,
+                req_pool_indices,
+                seq_lens,
+                kv_indptr,
+                None,
+                kv_indices,
+                self.req_to_token.stride(0),
+            )
+
+            custom_mask = self.cuda_graph_custom_mask
+            seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens)
+            mask_indptr = self.mask_indptr[: bs + 1]
+            mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0)
+            max_extend_len = self.num_draft_tokens
+            num_kv_splits = None
+            attn_logits = None
+            attn_lse = None
+        else:
+            raise ValueError(
+                f"Invalid forward mode: {forward_mode=} for CUDA Graph capture."
+            )
+
+        self.forward_metadata = ForwardMetadata(
+            attn_logits,
+            attn_lse,
+            max_extend_len,
+            num_kv_splits,
+            kv_indptr,
+            kv_indices,
+            qo_indptr,
+            custom_mask,
+            mask_indptr,
+        )
+
+    def init_forward_metadata_replay_cuda_graph(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        seq_lens_cpu: Optional[torch.Tensor],
+    ):
+        # NOTE: encoder_lens expected to be zeros or None
+        if forward_mode.is_decode_or_idle():
+            # Update kv_indptr, kv_indices
+            kv_indptr = self.kv_indptr
+            kv_indices = self.cuda_graph_kv_indices
+            num_kv_splits = self.cuda_graph_num_kv_splits
+            if spec_info is None:
+                kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens[:bs], dim=0)
+                kv_indptr = kv_indptr[: bs + 1]
+                create_flashinfer_kv_indices_triton[(bs,)](
+                    self.req_to_token,
+                    req_pool_indices[:bs],
+                    seq_lens[:bs],
+                    kv_indptr,
+                    None,
+                    kv_indices,
+                    self.req_to_token.stride(0),
+                )
+                num_token = bs
+            else:
+                kv_indptr[: spec_info.kv_indptr.shape[0]] = spec_info.kv_indptr
+                kv_indices[: spec_info.kv_indices.shape[0]] = spec_info.kv_indices
+                num_token = spec_info.kv_indptr.shape[0] - 1
+            self.get_num_kv_splits(num_kv_splits[:num_token], seq_lens[:bs])
+        elif forward_mode.is_target_verify():
+            # Update qo_indptr, kv_indptr, kv_indices, custom_mask, mask_indptr
+            bs = len(req_pool_indices)
+            qo_indptr = self.qo_indptr[: bs + 1]
+            qo_indptr[: bs + 1] = torch.arange(
+                0,
+                (1 + bs) * self.num_draft_tokens,
+                step=self.num_draft_tokens,
+                dtype=torch.int32,
+                device=self.device,
+            )
+            kv_indptr = self.kv_indptr[: bs + 1]
+            kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
+            kv_indices = self.cuda_graph_kv_indices
+            create_flashinfer_kv_indices_triton[(bs,)](
+                self.req_to_token,
+                req_pool_indices,
+                seq_lens,
+                kv_indptr,
+                None,
+                kv_indices,
+                self.req_to_token.stride(0),
+            )
+            custom_mask = self.cuda_graph_custom_mask
+            custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask
+            seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens)
+            mask_indptr = self.mask_indptr[: bs + 1]
+            mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0)
+        else:
+            raise ValueError(
+                f"Invalid forward mode: {forward_mode=} for CUDA Graph replay."
+            )
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        return 1
+
+    def forward_extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+    ):
+        # TODO: reuse the buffer across layers
+        if layer.qk_head_dim != layer.v_head_dim:
+            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
+        else:
+            o = torch.empty_like(q)
+
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v
+            )
+
+        causal = True
+        if layer.attn_type == AttentionType.ENCODER_ONLY:
+            causal = False
+
+        self.extend_attention_fwd(
+            q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+            k.contiguous(),
+            v.contiguous(),
+            o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
+            forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
+            forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
+            self.forward_metadata.qo_indptr,
+            self.forward_metadata.kv_indptr,
+            self.forward_metadata.kv_indices,
+            self.forward_metadata.custom_mask,
+            causal,
+            self.forward_metadata.mask_indptr,
+            self.forward_metadata.max_extend_len,
+            layer.scaling,
+            layer.logit_cap,
+        )
+        return o
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+    ):
+        # During torch.compile, there is a bug in rotary_emb that causes the
+        # output value to have a 3D tensor shape. This reshapes the output correctly.
+        q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
+
+        # TODO: reuse the buffer across layers
+        if layer.qk_head_dim != layer.v_head_dim:
+            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
+        else:
+            o = torch.empty_like(q)
+
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v
+            )
+
+        self.decode_attention_fwd(
+            q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+            forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
+            forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
+            o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
+            self.forward_metadata.kv_indptr,
+            self.forward_metadata.kv_indices,
+            self.forward_metadata.attn_logits,
+            self.forward_metadata.attn_lse,
+            self.forward_metadata.num_kv_splits,
+            self.max_kv_splits,
+            layer.scaling,
+            layer.logit_cap,
+        )
+        return o
+
+
+class TritonMultiStepDraftBackend:
+    """
+    Wrap multiple triton attention backends as one for multiple consecutive
+    draft decoding steps.
+    """
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        topk: int,
+        speculative_num_steps: int,
+    ):
+        from sglang.srt.speculative.eagle_utils import generate_draft_decode_kv_indices
+
+        self.topk = topk
+        self.speculative_num_steps = speculative_num_steps
+        self.generate_draft_decode_kv_indices = generate_draft_decode_kv_indices
+        max_bs = model_runner.req_to_token_pool.size * self.topk
+        self.kv_indptr = torch.zeros(
+            (
+                self.speculative_num_steps,
+                max_bs + 1,
+            ),
+            dtype=torch.int32,
+            device=model_runner.device,
+        )
+        self.attn_backends = []
+        for i in range(self.speculative_num_steps):
+            self.attn_backends.append(
+                TritonAttnBackend(
+                    model_runner,
+                    skip_prefill=True,
+                    kv_indptr_buf=self.kv_indptr[i],
+                )
+            )
+        self.max_context_len = self.attn_backends[0].max_context_len
+        self.num_head = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.device = model_runner.device
+        # Cached variables for generate_draft_decode_kv_indices
+        self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1]
+
+    def common_template(
+        self, forward_batch: ForwardBatch, kv_indices_buffer: torch.Tensor, call_fn: int
+    ):
+        num_seqs = forward_batch.batch_size
+        bs = self.topk * num_seqs
+        seq_lens_sum = forward_batch.seq_lens_sum
+
+        self.generate_draft_decode_kv_indices[
+            (self.speculative_num_steps, num_seqs, self.topk)
+        ](
+            forward_batch.req_pool_indices,
+            forward_batch.req_to_token_pool.req_to_token,
+            forward_batch.seq_lens,
+            kv_indices_buffer,
+            self.kv_indptr,
+            forward_batch.positions,
+            num_seqs,
+            self.topk,
+            self.pool_len,
+            kv_indices_buffer.shape[1],
+            self.kv_indptr.shape[1],
+            triton.next_power_of_2(num_seqs),
+            triton.next_power_of_2(self.speculative_num_steps),
+            triton.next_power_of_2(bs),
+        )
+
+        for i in range(self.speculative_num_steps):
+            forward_batch.spec_info.kv_indptr = self.kv_indptr[i, : bs + 1]
+            forward_batch.spec_info.kv_indices = kv_indices_buffer[i][
+                : seq_lens_sum * self.topk + bs * (i + 1)
+            ]
+            call_fn(i, forward_batch)
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        kv_indices = torch.empty(
+            (
+                self.speculative_num_steps,
+                forward_batch.batch_size * self.topk * self.max_context_len,
+            ),
+            dtype=torch.int32,
+            device=self.device,
+        )
+
+        def call_fn(i, forward_batch):
+            forward_batch.spec_info.kv_indptr = (
+                forward_batch.spec_info.kv_indptr.clone()
+            )
+            forward_batch.spec_info.kv_indices = (
+                forward_batch.spec_info.kv_indices.clone()
+            )
+            self.attn_backends[i].init_forward_metadata(forward_batch)
+
+        self.common_template(forward_batch, kv_indices, call_fn)
+
+    def init_cuda_graph_state(self, max_bs: int):
+        self.cuda_graph_kv_indices = torch.zeros(
+            (self.speculative_num_steps, max_bs * self.max_context_len),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        for i in range(self.speculative_num_steps):
+            self.attn_backends[i].init_cuda_graph_state(
+                max_bs, kv_indices_buf=self.cuda_graph_kv_indices[i]
+            )
+
+    def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch):
+        def call_fn(i, forward_batch):
+            self.attn_backends[i].init_forward_metadata_capture_cuda_graph(
+                forward_batch.batch_size,
+                forward_batch.batch_size * self.topk,
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                encoder_lens=None,
+                forward_mode=ForwardMode.DECODE,
+                spec_info=forward_batch.spec_info,
+            )
+
+        self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn)
+
+    def init_forward_metadata_replay_cuda_graph(
+        self, forward_batch: ForwardBatch, bs: int
+    ):
+        def call_fn(i, forward_batch):
+            self.attn_backends[i].init_forward_metadata_replay_cuda_graph(
+                bs,
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                seq_lens_sum=-1,
+                encoder_lens=None,
+                forward_mode=ForwardMode.DECODE,
+                spec_info=forward_batch.spec_info,
+                seq_lens_cpu=None,
+            )
+
+        self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn)

lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/decode_attention.py

@@ -0,0 +1,733 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""
+Memory-efficient attention for decoding.
+It supports page size = 1.
+"""
+
+# Adapted from
+# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py
+# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py
+
+import logging
+
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
+
+logger = logging.getLogger(__name__)
+
+# TODO: Remove this when triton>=3.2.0. This issue will not affect performance and accuracy.
+logger.warning(
+    "The following error message 'operation scheduled before its operands' can be ignored."
+)
+
+
+_MIN_BLOCK_KV = 32
+
+
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+
+@triton.jit
+def _fwd_kernel_stage1(
+    Q,
+    K_Buffer,
+    V_Buffer,
+    sm_scale,
+    kv_indptr,
+    kv_indices,
+    Att_Out,
+    Att_Lse,
+    num_kv_splits,
+    stride_qbs,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    stride_buf_vbs,
+    stride_buf_vh,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    kv_group_num: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
+    logit_cap: tl.constexpr,
+    Lk: tl.constexpr,
+    Lv: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    split_kv_id = tl.program_id(2)
+
+    cur_kv_head = cur_head // kv_group_num
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    mask_d = offs_d < Lk
+    mask_dv = offs_dv < Lv
+
+    cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+    kv_splits = tl.load(num_kv_splits + cur_batch)
+
+    off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d
+
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
+    split_kv_start = kv_len_per_split * split_kv_id
+    split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+    e_max = -float("inf")
+    e_sum = 0.0
+    acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
+
+    if split_kv_end > split_kv_start:
+        q = tl.load(Q + off_q, mask=mask_d, other=0.0)
+        for start_n in range(split_kv_start, split_kv_end, BLOCK_N):
+            offs_n = start_n + tl.arange(0, BLOCK_N)
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + offs_n,
+                mask=offs_n < split_kv_end,
+                other=0,
+            )
+            offs_buf_k = (
+                kv_loc[:, None] * stride_buf_kbs
+                + cur_kv_head * stride_buf_kh
+                + offs_d[None, :]
+            )
+            k = tl.load(
+                K_Buffer + offs_buf_k,
+                mask=(offs_n[:, None] < split_kv_end) & (mask_d[None, :]),
+                other=0.0,
+            )
+            qk = tl.sum(q[None, :] * k, 1)
+            qk *= sm_scale
+
+            if logit_cap > 0:
+                qk = logit_cap * tanh(qk / logit_cap)
+
+            qk = tl.where(offs_n < split_kv_end, qk, float("-inf"))
+
+            offs_buf_v = (
+                kv_loc[:, None] * stride_buf_vbs
+                + cur_kv_head * stride_buf_vh
+                + offs_dv[None, :]
+            )
+            v = tl.load(
+                V_Buffer + offs_buf_v,
+                mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),
+                other=0.0,
+            )
+
+            n_e_max = tl.maximum(tl.max(qk, 0), e_max)
+            re_scale = tl.exp(e_max - n_e_max)
+            p = tl.exp(qk - n_e_max)
+            acc *= re_scale
+            acc += tl.sum(p[:, None] * v, 0)
+
+            e_sum = e_sum * re_scale + tl.sum(p, 0)
+            e_max = n_e_max
+
+        offs_mid_o = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + split_kv_id * stride_mid_os
+            + offs_dv
+        )
+
+        tl.store(
+            Att_Out + offs_mid_o,
+            acc / e_sum,
+            mask=(mask_dv),
+        )
+
+        offs_mid_o_1 = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + split_kv_id * stride_mid_os
+        ) // Lv
+
+        tl.store(
+            Att_Lse + offs_mid_o_1,
+            e_max + tl.log(e_sum),
+        )
+
+
+def _decode_att_m_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    att_out,
+    att_lse,
+    kv_indptr,
+    kv_indices,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale,
+    logit_cap,
+):
+    BLOCK = 64
+    # [TODO] work around SGPR limit on MI3xx
+    if _is_hip:
+        BLOCK = 8
+    MAX_KV_SPLITS = max_kv_splits
+    Lk = k_buffer.shape[-1]
+    Lv = v_buffer.shape[-1]
+
+    batch, head_num = kv_indptr.shape[0] - 1, q.shape[1]
+
+    grid = (batch, head_num, MAX_KV_SPLITS)
+    kv_group_num = q.shape[1] // k_buffer.shape[1]
+
+    if kv_group_num == 1:
+        num_warps = 4
+    else:
+        num_warps = 2
+        if _is_hip:
+            num_warps = 1
+
+    BLOCK_DMODEL = triton.next_power_of_2(Lk)
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    _fwd_kernel_stage1[grid](
+        q,
+        k_buffer,
+        v_buffer,
+        sm_scale,
+        kv_indptr,
+        kv_indices,
+        att_out,
+        att_lse,
+        num_kv_splits,
+        q.stride(0),
+        q.stride(1),
+        k_buffer.stride(0),
+        k_buffer.stride(1),
+        v_buffer.stride(0),
+        v_buffer.stride(1),
+        att_out.stride(0),
+        att_out.stride(1),
+        att_out.stride(2),
+        kv_group_num=kv_group_num,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DV=BLOCK_DV,
+        BLOCK_N=BLOCK,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
+        logit_cap=logit_cap,
+        num_warps=num_warps,
+        num_stages=2,
+        Lk=Lk,
+        Lv=Lv,
+    )
+
+
+@triton.jit
+def _fwd_grouped_kernel_stage1(
+    Q,
+    K_Buffer,
+    V_Buffer,
+    sm_scale,
+    kv_indptr,
+    kv_indices,
+    Att_Out,
+    Att_Lse,
+    num_kv_splits,
+    stride_qbs,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    stride_buf_vbs,
+    stride_buf_vh,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    kv_group_num: tl.constexpr,
+    q_head_num: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DPE: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
+    logit_cap: tl.constexpr,
+    Lk: tl.constexpr,
+    Lv: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head_id = tl.program_id(1)
+    cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H)
+    split_kv_id = tl.program_id(2)
+
+    if BLOCK_H < kv_group_num:
+        VALID_BLOCK_H: tl.constexpr = BLOCK_H
+    else:
+        VALID_BLOCK_H: tl.constexpr = kv_group_num
+    cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H
+    mask_h = mask_h & (cur_head < q_head_num)
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    mask_d = offs_d < Lk
+    mask_dv = offs_dv < Lv
+
+    cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+    kv_splits = tl.load(num_kv_splits + cur_batch)
+
+    offs_q = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_d[None, :]
+
+    if BLOCK_DPE > 0:
+        offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
+        mask_dpe = offs_dpe < Lk
+        off_qpe = (
+            cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :]
+        )
+
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
+    split_kv_start = kv_len_per_split * split_kv_id
+    split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+    e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf")
+    e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32)
+
+    if split_kv_end > split_kv_start:
+        q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_d[None, :]), other=0.0)
+        if BLOCK_DPE > 0:
+            qpe = tl.load(
+                Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0
+            )
+        for start_n in range(split_kv_start, split_kv_end, BLOCK_N):
+            offs_n = start_n + tl.arange(0, BLOCK_N)
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + offs_n,
+                mask=offs_n < split_kv_end,
+                other=0,
+            )
+            offs_buf_k = (
+                kv_loc[None, :] * stride_buf_kbs
+                + cur_kv_head * stride_buf_kh
+                + offs_d[:, None]
+            )
+            k = tl.load(
+                K_Buffer + offs_buf_k,
+                mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]),
+                other=0.0,
+            )
+            qk = tl.dot(q, k.to(q.dtype))
+            if BLOCK_DPE > 0:
+                offs_buf_kpe = (
+                    kv_loc[None, :] * stride_buf_kbs
+                    + cur_kv_head * stride_buf_kh
+                    + offs_dpe[:, None]
+                )
+                kpe = tl.load(
+                    K_Buffer + offs_buf_kpe,
+                    mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]),
+                    other=0.0,
+                )
+                qk += tl.dot(qpe, kpe.to(qpe.dtype))
+            qk *= sm_scale
+
+            if logit_cap > 0:
+                qk = logit_cap * tanh(qk / logit_cap)
+
+            qk = tl.where(
+                mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf")
+            )
+
+            offs_buf_v = (
+                kv_loc[:, None] * stride_buf_vbs
+                + cur_kv_head * stride_buf_vh
+                + offs_dv[None, :]
+            )
+            v = tl.load(
+                V_Buffer + offs_buf_v,
+                mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),
+                other=0.0,
+            )
+
+            n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+            re_scale = tl.exp(e_max - n_e_max)
+            p = tl.exp(qk - n_e_max[:, None])
+            acc *= re_scale[:, None]
+            acc += tl.dot(p.to(v.dtype), v)
+
+            e_sum = e_sum * re_scale + tl.sum(p, 1)
+            e_max = n_e_max
+
+        offs_mid_o = (
+            cur_batch * stride_mid_ob
+            + cur_head[:, None] * stride_mid_oh
+            + split_kv_id * stride_mid_os
+            + offs_dv[None, :]
+        )
+
+        tl.store(
+            Att_Out + offs_mid_o,
+            acc / e_sum[:, None],
+            mask=(mask_h[:, None]) & (mask_dv[None, :]),
+        )
+
+        offs_mid_o_1 = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + split_kv_id * stride_mid_os
+        ) // Lv
+
+        tl.store(
+            Att_Lse + offs_mid_o_1,
+            e_max + tl.log(e_sum),
+            mask=mask_h,
+        )
+
+
+def _decode_grouped_att_m_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    att_out,
+    att_lse,
+    kv_indptr,
+    kv_indices,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale,
+    logit_cap,
+):
+    BLOCK = 32
+    Lk = k_buffer.shape[-1]
+    Lv = v_buffer.shape[-1]
+
+    # [TODO] work around shmem limit on MI3xx
+    if _is_hip and Lk >= 576:
+        BLOCK = 16
+
+    if Lk == 576:
+        BLOCK_DMODEL = 512
+        BLOCK_DPE = 64
+    elif Lk == 288:
+        BLOCK_DMODEL = 256
+        BLOCK_DPE = 32
+    else:
+        BLOCK_DMODEL = triton.next_power_of_2(Lk)
+        BLOCK_DPE = 0
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    batch, head_num = kv_indptr.shape[0] - 1, q.shape[1]
+    kv_group_num = q.shape[1] // k_buffer.shape[1]
+
+    BLOCK_H = 16
+    MAX_KV_SPLITS = max_kv_splits
+    grid = (
+        batch,
+        triton.cdiv(head_num, min(BLOCK_H, kv_group_num)),
+        MAX_KV_SPLITS,
+    )
+
+    extra_kargs = {}
+    num_stages = 2
+    if _is_hip:
+        # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
+        # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
+        extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}
+        num_stages = 1
+
+    _fwd_grouped_kernel_stage1[grid](
+        q,
+        k_buffer,
+        v_buffer,
+        sm_scale,
+        kv_indptr,
+        kv_indices,
+        att_out,
+        att_lse,
+        num_kv_splits,
+        q.stride(0),
+        q.stride(1),
+        k_buffer.stride(0),
+        k_buffer.stride(1),
+        v_buffer.stride(0),
+        v_buffer.stride(1),
+        att_out.stride(0),
+        att_out.stride(1),
+        att_out.stride(2),
+        kv_group_num=kv_group_num,
+        q_head_num=head_num,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DPE=BLOCK_DPE,
+        BLOCK_DV=BLOCK_DV,
+        BLOCK_N=BLOCK,
+        BLOCK_H=BLOCK_H,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
+        logit_cap=logit_cap,
+        num_warps=4,
+        num_stages=num_stages,
+        Lk=Lk,
+        Lv=Lv,
+        **extra_kargs,
+    )
+
+
+@triton.jit
+def _fwd_kernel_stage2(
+    Mid_O,
+    Mid_O_1,
+    O,
+    kv_indptr,
+    num_kv_splits,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_obs,
+    stride_oh,
+    MAX_KV_SPLITS: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    Lv: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - tl.load(
+        kv_indptr + cur_batch
+    )
+    kv_splits = tl.load(num_kv_splits + cur_batch)
+
+    offs_d = tl.arange(0, BLOCK_DV)
+    mask_d = offs_d < Lv
+
+    e_sum = 0.0
+    e_max = -float("inf")
+    acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
+
+    offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d
+    offs_logic = (cur_batch * stride_mid_ob + cur_head * stride_mid_oh) // Lv
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
+
+    for split_kv_id in range(0, MAX_KV_SPLITS):
+        split_kv_start = kv_len_per_split * split_kv_id
+        split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+        if split_kv_end > split_kv_start:
+            tv = tl.load(
+                Mid_O + offs_v + split_kv_id * stride_mid_os, mask=mask_d, other=0.0
+            )
+            tlogic = tl.load(Mid_O_1 + offs_logic + split_kv_id * stride_mid_os // Lv)
+            n_e_max = tl.maximum(tlogic, e_max)
+
+            old_scale = tl.exp(e_max - n_e_max)
+            acc *= old_scale
+            exp_logic = tl.exp(tlogic - n_e_max)
+            acc += exp_logic * tv
+
+            e_sum = e_sum * old_scale + exp_logic
+            e_max = n_e_max
+
+    tl.store(
+        O + cur_batch * stride_obs + cur_head * stride_oh + offs_d,
+        acc / e_sum,
+        mask=mask_d,
+    )
+
+
+def _decode_softmax_reducev_fwd(
+    logits,
+    lse,
+    q,
+    o,
+    v_buffer,
+    kv_indptr,
+    num_kv_splits,
+    max_kv_splits,
+):
+    batch, head_num = q.shape[0], q.shape[1]
+    Lv = v_buffer.shape[-1]
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    MAX_KV_SPLITS = max_kv_splits
+
+    extra_kargs = {}
+    if _is_hip:
+        # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
+        # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
+        extra_kargs = {"waves_per_eu": 4, "matrix_instr_nonkdim": 16, "kpack": 2}
+
+    grid = (batch, head_num)
+    _fwd_kernel_stage2[grid](
+        logits,
+        lse,
+        o,
+        kv_indptr,
+        num_kv_splits,
+        logits.stride(0),
+        logits.stride(1),
+        logits.stride(2),
+        o.stride(0),
+        o.stride(1),
+        MAX_KV_SPLITS=MAX_KV_SPLITS,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
+        BLOCK_DV=BLOCK_DV,
+        Lv=Lv,
+        num_warps=4,
+        num_stages=2,
+        **extra_kargs,
+    )
+
+
+def decode_attention_fwd_normal(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    kv_indptr,
+    kv_indices,
+    attn_logits,
+    attn_lse,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale,
+    logit_cap=0.0,
+):
+    _decode_att_m_fwd(
+        q,
+        k_buffer,
+        v_buffer,
+        attn_logits,
+        attn_lse,
+        kv_indptr,
+        kv_indices,
+        num_kv_splits,
+        max_kv_splits,
+        sm_scale,
+        logit_cap,
+    )
+    _decode_softmax_reducev_fwd(
+        attn_logits,
+        attn_lse,
+        q,
+        o,
+        v_buffer,
+        kv_indptr,
+        num_kv_splits,
+        max_kv_splits,
+    )
+
+
+def decode_attention_fwd_grouped(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    kv_indptr,
+    kv_indices,
+    attn_logits,
+    attn_lse,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale,
+    logit_cap=0.0,
+):
+    _decode_grouped_att_m_fwd(
+        q,
+        k_buffer,
+        v_buffer,
+        attn_logits,
+        attn_lse,
+        kv_indptr,
+        kv_indices,
+        num_kv_splits,
+        max_kv_splits,
+        sm_scale,
+        logit_cap,
+    )
+    _decode_softmax_reducev_fwd(
+        attn_logits,
+        attn_lse,
+        q,
+        o,
+        v_buffer,
+        kv_indptr,
+        num_kv_splits,
+        max_kv_splits,
+    )
+
+
+def decode_attention_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    kv_indptr,
+    kv_indices,
+    attn_logits,
+    attn_lse,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale,
+    logit_cap=0.0,
+):
+    assert max_kv_splits == attn_logits.shape[2]
+    assert q.shape[0] <= kv_indptr.shape[0] - 1
+    assert q.shape[0] <= attn_logits.shape[0]
+
+    kv_group_num = q.shape[1] // v_buffer.shape[1]
+
+    if kv_group_num == 1:
+        # MHA
+        decode_attention_fwd_normal(
+            q,
+            k_buffer,
+            v_buffer,
+            o,
+            kv_indptr,
+            kv_indices,
+            attn_logits,
+            attn_lse,
+            num_kv_splits,
+            max_kv_splits,
+            sm_scale,
+            logit_cap,
+        )
+    else:
+        # GQA/MQA/MLA
+        decode_attention_fwd_grouped(
+            q,
+            k_buffer,
+            v_buffer,
+            o,
+            kv_indptr,
+            kv_indices,
+            attn_logits,
+            attn_lse,
+            num_kv_splits,
+            max_kv_splits,
+            sm_scale,
+            logit_cap,
+        )

lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/double_sparsity_attention.py

@@ -0,0 +1,1106 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.managers.schedule_batch import global_server_args_dict
+from sglang.srt.utils import is_cuda, is_hip
+
+_is_cuda = is_cuda()
+if _is_cuda:
+    CUDA_CAPABILITY = torch.cuda.get_device_capability()
+
+_is_hip = is_hip()
+
+if global_server_args_dict.get("attention_reduce_in_fp32", False):
+    REDUCE_TRITON_TYPE = tl.float32
+    REDUCE_TORCH_TYPE = torch.float32
+else:
+    REDUCE_TRITON_TYPE = tl.float16
+    REDUCE_TORCH_TYPE = torch.float16
+
+
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+
+@triton.jit
+def _fwd_kernel_flash_decode_stage1(
+    Q,
+    K,
+    V,
+    sm_scale,
+    Req_to_tokens,
+    B_req_idx,
+    B_Seqlen,
+    Mid_O,  # [batch, head, seq_block_num, head_dim]
+    Mid_O_LogExpSum,  # [batch, head, seq_block_num]
+    stride_req_to_tokens_b,
+    stride_req_to_tokens_s,
+    stride_qbs,
+    stride_qh,
+    stride_qd,
+    stride_kbs,
+    stride_kh,
+    stride_kd,
+    stride_vbs,
+    stride_vh,
+    stride_vd,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_mid_od,
+    stride_mid_o_eb,
+    stride_mid_o_eh,
+    stride_mid_o_es,
+    gqa_group_size,
+    BLOCK_SEQ: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    seq_start_block = tl.program_id(2)
+    cur_kv_head = cur_head // gqa_group_size
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
+    cur_batch_start_index = seq_start_block * BLOCK_SEQ
+    cur_batch_end_index = tl.minimum(
+        cur_batch_seq_len, cur_batch_start_index + BLOCK_SEQ
+    )
+
+    off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d
+
+    block_n_size = (
+        tl.where(
+            cur_batch_end_index - cur_batch_start_index <= 0,
+            0,
+            cur_batch_end_index - cur_batch_start_index + BLOCK_N - 1,
+        )
+        // BLOCK_N
+    )
+
+    offs_n = cur_batch_start_index + tl.arange(0, BLOCK_N)
+
+    q = tl.load(Q + off_q)
+
+    sum_exp = 0.0
+    max_logic = -float("inf")
+    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
+
+    for start_n in range(0, block_n_size, 1):
+        offs_n_new = start_n * BLOCK_N + offs_n
+        k_loc = tl.load(
+            Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_n_new,
+            mask=offs_n_new < cur_batch_end_index,
+            other=0,
+        )
+        off_k = k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :]
+        k = tl.load(
+            K + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0
+        )
+        att_value = tl.sum(q[None, :] * k, 1)
+        att_value *= sm_scale
+        att_value = tl.where(offs_n_new < cur_batch_end_index, att_value, float("-inf"))
+        v = tl.load(
+            V + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0
+        )
+
+        cur_max_logic = tl.max(att_value, axis=0)
+        new_max_logic = tl.maximum(cur_max_logic, max_logic)
+
+        exp_logic = tl.exp(att_value - new_max_logic)
+        logic_scale = tl.exp(max_logic - new_max_logic)
+        acc *= logic_scale
+        acc += tl.sum(exp_logic[:, None] * v, axis=0)
+
+        sum_exp = sum_exp * logic_scale + tl.sum(exp_logic, axis=0)
+        max_logic = new_max_logic
+
+    need_store = tl.where(block_n_size == 0, 0, 1)
+    for _ in range(0, need_store, 1):
+        off_mid_o = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + seq_start_block * stride_mid_os
+            + offs_d
+        )
+        off_mid_o_logexpsum = (
+            cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh + seq_start_block
+        )
+        tl.store(Mid_O + off_mid_o, acc / sum_exp)
+        tl.store(Mid_O_LogExpSum + off_mid_o_logexpsum, max_logic + tl.log(sum_exp))
+    return
+
+
+@triton.jit
+def _fwd_kernel_flash_decode_stage2(
+    B_Seqlen,
+    Mid_O,  # [batch, head, seq_block_num, head_dim]
+    Mid_O_LogExpSum,  # [batch, head, seq_block_num]
+    O,  # [batch, head, head_dim]
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_mid_od,
+    stride_mid_o_eb,
+    stride_mid_o_eh,
+    stride_mid_o_es,
+    stride_obs,
+    stride_oh,
+    stride_od,
+    BLOCK_SEQ: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+
+    block_n_size = (
+        tl.where(cur_batch_seq_len <= 0, 0, cur_batch_seq_len + BLOCK_SEQ - 1)
+        // BLOCK_SEQ
+    )
+
+    sum_exp = 0.0
+    max_logic = -float("inf")
+    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
+
+    offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d
+    offs_logic = cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh
+    for block_seq_n in range(0, block_n_size, 1):
+        tv = tl.load(Mid_O + offs_v + block_seq_n * stride_mid_os)
+        tlogic = tl.load(Mid_O_LogExpSum + offs_logic + block_seq_n)
+        new_max_logic = tl.maximum(tlogic, max_logic)
+
+        old_scale = tl.exp(max_logic - new_max_logic)
+        acc *= old_scale
+        exp_logic = tl.exp(tlogic - new_max_logic)
+        acc += exp_logic * tv
+        sum_exp = sum_exp * old_scale + exp_logic
+        max_logic = new_max_logic
+
+    tl.store(O + cur_batch * stride_obs + cur_head * stride_oh + offs_d, acc / sum_exp)
+    return
+
+
+@torch.no_grad()
+def flash_decode_stage1(
+    q,
+    k,
+    v,
+    Req_to_tokens,
+    B_req_idx,
+    B_Seqlen,
+    max_len_in_batch,
+    mid_out,
+    mid_out_logsumexp,
+    block_seq,
+):
+    BLOCK_SEQ = block_seq
+    BLOCK_N = 16
+    assert BLOCK_SEQ % BLOCK_N == 0
+    # shape constraints
+    Lq, Lk = q.shape[-1], k.shape[-1]
+    assert Lq == Lk
+    assert Lk in {16, 32, 64, 128}
+    sm_scale = 1.0 / (Lk**0.5)
+    batch, head_num = B_req_idx.shape[0], q.shape[1]
+    grid = (batch, head_num, triton.cdiv(max_len_in_batch, BLOCK_SEQ))
+    gqa_group_size = q.shape[1] // k.shape[1]
+
+    _fwd_kernel_flash_decode_stage1[grid](
+        q,
+        k,
+        v,
+        sm_scale,
+        Req_to_tokens,
+        B_req_idx,
+        B_Seqlen,
+        mid_out,
+        mid_out_logsumexp,
+        Req_to_tokens.stride(0),
+        Req_to_tokens.stride(1),
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        v.stride(0),
+        v.stride(1),
+        v.stride(2),
+        mid_out.stride(0),
+        mid_out.stride(1),
+        mid_out.stride(2),
+        mid_out.stride(3),
+        mid_out_logsumexp.stride(0),
+        mid_out_logsumexp.stride(1),
+        mid_out_logsumexp.stride(2),
+        gqa_group_size,
+        BLOCK_SEQ=BLOCK_SEQ,
+        BLOCK_DMODEL=Lk,
+        BLOCK_N=BLOCK_N,
+        num_warps=1,
+        num_stages=2,
+    )
+    return
+
+
+@torch.no_grad()
+def flash_decode_stage2(mid_out, mid_out_logexpsum, B_Seqlen, O, block_seq):
+    Lk = mid_out.shape[-1]
+    assert Lk in {16, 32, 64, 128}
+    batch, head_num = mid_out.shape[0], mid_out.shape[1]
+    grid = (batch, head_num)
+
+    _fwd_kernel_flash_decode_stage2[grid](
+        B_Seqlen,
+        mid_out,
+        mid_out_logexpsum,
+        O,
+        mid_out.stride(0),
+        mid_out.stride(1),
+        mid_out.stride(2),
+        mid_out.stride(3),
+        mid_out_logexpsum.stride(0),
+        mid_out_logexpsum.stride(1),
+        mid_out_logexpsum.stride(2),
+        O.stride(0),
+        O.stride(1),
+        O.stride(2),
+        BLOCK_SEQ=block_seq,
+        BLOCK_DMODEL=Lk,
+        num_warps=4,
+        num_stages=2,
+    )
+    return
+
+
+def flash_decode_attention_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    req_to_token,
+    b_req_idx,
+    b_start_loc,
+    b_seq_len,
+    attn_logits,
+    max_len_in_batch,
+    sm_scale,
+    logit_cap=0.0,
+):
+    BLOCK_SEQ = 256
+    kv_group_num = q.shape[1] // v_buffer.shape[1]
+    # batch_size = q.shape[0]
+
+    block_seq_num = (max_len_in_batch + BLOCK_SEQ - 1) // BLOCK_SEQ
+
+    mid_o = torch.empty(
+        [q.shape[0], q.shape[1], block_seq_num, q.shape[-1]],
+        dtype=torch.float32,
+        device="cuda",
+    )
+    mid_o_logexpsum = torch.empty(
+        [q.shape[0], q.shape[1], block_seq_num], dtype=torch.float32, device="cuda"
+    )
+
+    flash_decode_stage1(
+        q,
+        k_buffer,
+        v_buffer,
+        req_to_token,
+        b_req_idx,
+        b_seq_len,
+        max_len_in_batch,
+        mid_o,
+        mid_o_logexpsum,
+        BLOCK_SEQ,
+    )
+    flash_decode_stage2(mid_o, mid_o_logexpsum, b_seq_len, o, BLOCK_SEQ)
+
+
+@triton.jit
+def _sparse_fwd_kernel_flash_decode_stage1(  # Double Sparsity's approximate attention
+    Q_Label,
+    K_Label_Buffer,
+    sm_scale,
+    Req_to_tokens,  # shape: [B, S]
+    B_Seqlen,
+    Att_Out,  # shape: [H, B, S] easier for topk
+    stride_req_to_tokens_b,
+    stride_qbs,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    att_stride_h,
+    att_stride_b,
+    kv_group_num: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    logit_cap: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    start_n = tl.program_id(2)
+
+    cur_kv_head = cur_head // kv_group_num
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+
+    cur_batch_start_index = 0
+    cur_batch_end_index = cur_batch_seq_len
+
+    min_val = -float("inf")
+    att_value = tl.full([BLOCK_N], min_val, dtype=tl.float32)
+
+    off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d
+
+    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    block_index = start_n * BLOCK_N
+    block_mask = tl.where(block_index < cur_batch_seq_len, 1, 0)
+
+    for start_mark in range(0, block_mask, 1):
+        q = tl.load(Q_Label + off_q + start_mark).to(REDUCE_TRITON_TYPE)
+        offs_n_new = cur_batch_start_index + offs_n
+        k_loc = tl.load(
+            Req_to_tokens + stride_req_to_tokens_b * cur_batch + offs_n_new,
+            mask=offs_n_new < cur_batch_end_index,
+            other=0,
+        )
+        offs_buf_k = (
+            k_loc[:, None] * stride_buf_kbs
+            + cur_kv_head * stride_buf_kh
+            + offs_d[None, :]
+        )
+        k = tl.load(
+            K_Label_Buffer + offs_buf_k,
+            mask=offs_n_new[:, None] < cur_batch_end_index,
+            other=0.0,
+        ).to(REDUCE_TRITON_TYPE)
+
+        att_value = tl.sum(q[None, :] * k, 1)
+        att_value *= sm_scale
+
+        if logit_cap > 0:
+            att_value = logit_cap * tanh(att_value / logit_cap)
+
+    att_value = tl.where(offs_n < cur_batch_end_index, att_value, min_val)
+    off_o = cur_head * att_stride_h + (cur_batch * att_stride_b + offs_n)
+    tl.store(Att_Out + off_o, att_value)
+
+
+@triton.jit
+def _sparse_fwd_kernel_flash_decode_stage2(
+    Q,
+    K,
+    V,
+    sm_scale,
+    Req_to_tokens,  # shape: [B, S]
+    Topk_token_indices,  # shape: [H, B, k]
+    Mid_O,  # [batch, head, seq_block_num, head_dim]
+    Mid_O_LogExpSum,  # [batch, head, seq_block_num]
+    Heavy_token_num,  # NOTE: This can be used as constexpr but we may support dynamic heavy token number in the future
+    stride_req_to_tokens_b,
+    stride_topk_token_indices_h,
+    stride_topk_token_indices_b,
+    stride_qbs,
+    stride_qh,
+    stride_kbs,
+    stride_kh,
+    stride_vbs,
+    stride_vh,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_mid_o_eb,
+    stride_mid_o_eh,
+    gqa_group_size,
+    BLOCK_SEQ: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    seq_start_block = tl.program_id(2)
+    cur_kv_head = cur_head // gqa_group_size
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    cur_batch_start_index = seq_start_block * BLOCK_SEQ
+    cur_batch_end_index = tl.minimum(Heavy_token_num, cur_batch_start_index + BLOCK_SEQ)
+
+    off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d
+
+    block_n_size = (
+        tl.where(
+            cur_batch_end_index - cur_batch_start_index <= 0,
+            0,
+            cur_batch_end_index - cur_batch_start_index + BLOCK_N - 1,
+        )
+        // BLOCK_N
+    )
+
+    # offs_n = cur_batch_start_index + tl.arange(0, BLOCK_N)
+    offs_n = tl.arange(0, BLOCK_N)
+
+    q = tl.load(Q + off_q)
+
+    sum_exp = 0.0
+    max_logic = -float("inf")
+    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
+
+    for start_n in range(cur_batch_start_index, cur_batch_end_index, BLOCK_N):
+        # for start_n in range(0, block_n_size, 1):
+        # offs_n_new = start_n * BLOCK_N + offs_n
+        offs_n_new = start_n + offs_n
+        # offs_n_new = cur_batch_start_index + start_n * BLOCK_N + offs_n
+        topk_token_indices = tl.load(
+            Topk_token_indices
+            + stride_topk_token_indices_h * cur_head
+            + stride_topk_token_indices_b * cur_batch
+            + offs_n_new,
+            mask=offs_n_new < cur_batch_end_index,
+            other=0,
+        )
+        k_loc = tl.load(
+            Req_to_tokens + stride_req_to_tokens_b * cur_batch + topk_token_indices,
+            mask=offs_n_new < cur_batch_end_index,
+            other=0,
+        )
+        off_k = k_loc[:, None] * stride_kbs + cur_kv_head * stride_kh + offs_d[None, :]
+        k = tl.load(
+            K + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0
+        )
+        att_value = tl.sum(q[None, :] * k, 1)
+        att_value *= sm_scale
+        att_value = tl.where(offs_n_new < cur_batch_end_index, att_value, float("-inf"))
+        v = tl.load(
+            V + off_k, mask=offs_n_new[:, None] < cur_batch_end_index, other=0.0
+        )
+
+        cur_max_logic = tl.max(att_value, axis=0)
+        new_max_logic = tl.maximum(cur_max_logic, max_logic)
+
+        exp_logic = tl.exp(att_value - new_max_logic)
+        logic_scale = tl.exp(max_logic - new_max_logic)
+        acc *= logic_scale
+        acc += tl.sum(exp_logic[:, None] * v, axis=0)
+
+        sum_exp = sum_exp * logic_scale + tl.sum(exp_logic, axis=0)
+        max_logic = new_max_logic
+
+    # need_store = tl.where(block_n_size == 0, 0, 1)
+    need_store = 1
+    for _ in range(0, need_store, 1):
+        off_mid_o = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + seq_start_block * stride_mid_os
+            + offs_d
+        )
+        off_mid_o_logexpsum = (
+            cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh + seq_start_block
+        )
+        tl.store(Mid_O + off_mid_o, acc / sum_exp)
+        tl.store(Mid_O_LogExpSum + off_mid_o_logexpsum, max_logic + tl.log(sum_exp))
+    return
+
+
+@triton.jit
+def _sparse_fwd_kernel_flash_decode_stage3(
+    Mid_O,  # [batch, head, seq_block_num, head_dim]
+    Mid_O_LogExpSum,  # [batch, head, seq_block_num]
+    O,  # [batch, head, head_dim]
+    seq_len,  # NOTE: This can be used as constexpr but we may support dynamic heavy token number in the future
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_mid_o_eb,
+    stride_mid_o_eh,
+    stride_obs,
+    stride_oh,
+    BLOCK_SEQ: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+
+    block_n_size = tl.where(seq_len <= 0, 0, seq_len + BLOCK_SEQ - 1) // BLOCK_SEQ
+
+    sum_exp = 0.0
+    max_logic = -float("inf")
+    acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32)
+
+    offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d
+    offs_logic = cur_batch * stride_mid_o_eb + cur_head * stride_mid_o_eh
+    for block_seq_n in range(0, block_n_size, 1):
+        tv = tl.load(Mid_O + offs_v + block_seq_n * stride_mid_os)
+        tlogic = tl.load(Mid_O_LogExpSum + offs_logic + block_seq_n)
+        new_max_logic = tl.maximum(tlogic, max_logic)
+
+        old_scale = tl.exp(max_logic - new_max_logic)
+        acc *= old_scale
+        exp_logic = tl.exp(tlogic - new_max_logic)
+        acc += exp_logic * tv
+        sum_exp = sum_exp * old_scale + exp_logic
+        max_logic = new_max_logic
+
+    tl.store(O + cur_batch * stride_obs + cur_head * stride_oh + offs_d, acc / sum_exp)
+    return
+
+
+def sparse_flash_decode_stage1(
+    q_label,
+    k_label_buffer,
+    att_out,
+    Req_to_tokens,
+    B_Seqlen,
+    max_len_in_batch,
+    sm_scale,
+    logit_cap,
+):
+    BLOCK = 32
+    # shape constraints
+    Lq, Lk = q_label.shape[-1], k_label_buffer.shape[-1]
+    assert Lq == Lk
+    assert Lk in {16, 32, 64, 128, 256, 576}
+
+    BLOCK_DMODEL = Lk
+
+    batch, head_num = q_label.shape[0], q_label.shape[1]
+
+    grid = (batch, head_num, triton.cdiv(max_len_in_batch, BLOCK))
+    kv_group_num = q_label.shape[1] // k_label_buffer.shape[1]
+
+    if kv_group_num == 1:
+        num_warps = 4
+    else:
+        num_warps = 2
+
+    _sparse_fwd_kernel_flash_decode_stage1[grid](
+        q_label,
+        k_label_buffer,
+        sm_scale,
+        Req_to_tokens,
+        B_Seqlen,
+        att_out,
+        Req_to_tokens.stride(0),
+        q_label.stride(0),
+        q_label.stride(1),
+        k_label_buffer.stride(0),
+        k_label_buffer.stride(1),
+        att_out.stride(0),
+        att_out.stride(1),
+        kv_group_num,
+        BLOCK_DMODEL,
+        BLOCK,
+        logit_cap,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+
+
+@torch.no_grad()
+def sparse_flash_decode_stage2(
+    q,
+    k,
+    v,
+    Req_to_tokens,
+    Topk_token_indices,
+    heavy_token_num,
+    mid_out,
+    mid_out_logsumexp,
+    block_seq,
+    sm_scale,
+):
+    BLOCK_SEQ = block_seq
+    BLOCK_N = 16
+    assert BLOCK_SEQ % BLOCK_N == 0
+    # shape constraints
+    Lq, Lk = q.shape[-1], k.shape[-1]
+    assert Lq == Lk
+    assert Lk in {16, 32, 64, 128}
+    assert heavy_token_num == Topk_token_indices.shape[-1]
+    # sm_scale = 1.0 / (Lk ** 0.5)
+    batch, head_num = q.shape[0], q.shape[1]
+    grid = (batch, head_num, triton.cdiv(heavy_token_num, BLOCK_SEQ))
+
+    gqa_group_size = q.shape[1] // k.shape[1]
+
+    _sparse_fwd_kernel_flash_decode_stage2[grid](
+        q,
+        k,
+        v,
+        sm_scale,
+        Req_to_tokens,
+        Topk_token_indices,
+        mid_out,
+        mid_out_logsumexp,
+        heavy_token_num,
+        Req_to_tokens.stride(0),
+        Topk_token_indices.stride(0),
+        Topk_token_indices.stride(1),
+        q.stride(0),
+        q.stride(1),
+        k.stride(0),
+        k.stride(1),
+        v.stride(0),
+        v.stride(1),
+        mid_out.stride(0),
+        mid_out.stride(1),
+        mid_out.stride(2),
+        mid_out_logsumexp.stride(0),
+        mid_out_logsumexp.stride(1),
+        gqa_group_size,
+        BLOCK_SEQ=BLOCK_SEQ,
+        BLOCK_DMODEL=Lk,
+        BLOCK_N=BLOCK_N,
+        num_warps=1,
+        num_stages=2,
+    )
+    return
+
+
+@torch.no_grad()
+def sparse_flash_decode_stage3(Seqlen, mid_out, mid_out_logexpsum, O, block_seq):
+    Lk = mid_out.shape[-1]
+    assert Lk in {16, 32, 64, 128}
+    batch, head_num = mid_out.shape[0], mid_out.shape[1]
+    grid = (batch, head_num)
+
+    _sparse_fwd_kernel_flash_decode_stage3[grid](
+        mid_out,
+        mid_out_logexpsum,
+        O,
+        Seqlen,
+        mid_out.stride(0),
+        mid_out.stride(1),
+        mid_out.stride(2),
+        mid_out_logexpsum.stride(0),
+        mid_out_logexpsum.stride(1),
+        O.stride(0),
+        O.stride(1),
+        BLOCK_SEQ=block_seq,
+        BLOCK_DMODEL=Lk,
+        num_warps=4,
+        num_stages=2,
+    )
+    return
+
+
+def flash_decode_sparse_attention_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    q_label,
+    k_label_buffer,
+    req_to_token,
+    b_seq_len,
+    max_len_in_batch,
+    sm_scale,
+    logit_cap,
+    heavy_token_num=32,
+    att_out_approx=None,
+    mid_out=None,
+    mid_o_logexpsum=None,
+    BLOCK_SEQ=256,
+):
+    # TODO(Andy): Tune BLOCK_SEQ & BLOCK_D
+    kv_group_num = q.shape[1] // v_buffer.shape[1]
+    # batch_size = q.shape[0]
+
+    # Step 1: BGEMV approximate attention (page implementation)
+
+    if att_out_approx is None:
+        att_out_approx = torch.empty(
+            [q.shape[1], q.shape[0], max_len_in_batch],
+            dtype=REDUCE_TORCH_TYPE,
+            device=q.device,
+        )
+
+    if mid_out is None:
+        block_seq_num = (heavy_token_num + BLOCK_SEQ - 1) // BLOCK_SEQ
+
+        mid_out = torch.empty(
+            [q.shape[0], q.shape[1], block_seq_num, q.shape[-1]],
+            dtype=torch.float32,
+            device=q.device,
+        )
+        mid_o_logexpsum = torch.empty(
+            [q.shape[0], q.shape[1], block_seq_num],
+            dtype=torch.float32,
+            device=q.device,
+        )
+
+    sparse_flash_decode_stage1(
+        q_label,
+        k_label_buffer,
+        att_out_approx,
+        req_to_token,
+        b_seq_len,
+        max_len_in_batch,
+        sm_scale,
+        logit_cap,
+    )
+
+    # Step 2: TopK token selection
+    # NOTE(Andy): Apply sparse decoding when min > heavy_token_num and max > sparse decoding threshold
+    # TODO(Andy): Change a faster topk implementation
+    topk_token_indices = torch.topk(att_out_approx, heavy_token_num, dim=-1).indices
+    # topk_token_indices: [H, B, k], Req_to_tokens: [B, S]
+    # topk_token_indices = torch.arange(0, heavy_token_num, device=q.device).unsqueeze(0).unsqueeze(0).expand(q.shape[1], q.shape[0], -1)
+
+    sparse_flash_decode_stage2(
+        q,
+        k_buffer,
+        v_buffer,
+        req_to_token,
+        topk_token_indices,
+        heavy_token_num,
+        mid_out,
+        mid_o_logexpsum,
+        BLOCK_SEQ,
+        sm_scale,
+    )
+
+    sparse_flash_decode_stage3(heavy_token_num, mid_out, mid_o_logexpsum, o, BLOCK_SEQ)
+
+
+# Extend attention kernel for Double Sparsity
+# Moved from https://github.com/sgl-project/sglang/blob/v0.4.2.post1/python/sglang/srt/layers/attention/triton_ops/extend_attention.py
+@triton.jit
+def _fwd_kernel(
+    Q_Extend,
+    K_Extend,
+    V_Extend,
+    O_Extend,
+    K_Buffer,
+    V_Buffer,
+    Req_to_tokens,
+    B_req_idx,
+    B_Seq_Len,
+    B_Start_Loc_Extend,
+    B_Seq_Len_Extend,
+    sm_scale,
+    kv_group_num,
+    stride_qbs,
+    stride_qh,
+    stride_kbs,
+    stride_kh,
+    stride_vbs,
+    stride_vh,
+    stride_obs,
+    stride_oh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    stride_buf_vbs,
+    stride_buf_vh,
+    stride_req_to_tokens_b,
+    logit_cap: tl.constexpr,
+    Lq: tl.constexpr,
+    Lv: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DPE: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    cur_seq = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    cur_block_m = tl.program_id(2)
+    cur_kv_head = cur_head // kv_group_num
+
+    cur_seq_len = tl.load(B_Seq_Len + cur_seq)
+    cur_seq_len_extend = tl.load(B_Seq_Len_Extend + cur_seq)
+    cur_seq_len_prefix = cur_seq_len - cur_seq_len_extend
+
+    cur_seq_prefix_start_in_loc = 0
+    cur_seq_extend_start_contiguous = tl.load(B_Start_Loc_Extend + cur_seq)
+    cur_batch_req_idx = tl.load(B_req_idx + cur_seq)
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    offs_m = tl.arange(0, BLOCK_M)
+    mask_m = (cur_block_m * BLOCK_M + offs_m) < cur_seq_len_extend
+
+    mask_d = offs_d < Lq
+    mask_dv = offs_dv < Lv
+
+    offs_q = (
+        (cur_seq_extend_start_contiguous + cur_block_m * BLOCK_M + offs_m[:, None])
+        * stride_qbs
+        + cur_head * stride_qh
+        + offs_d[None, :]
+    )
+    q = tl.load(
+        Q_Extend + offs_q, mask=(mask_m[:, None]) & (mask_d[None, :]), other=0.0
+    )
+
+    if BLOCK_DPE > 0:
+        offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
+        offs_qpe = (
+            (cur_seq_extend_start_contiguous + cur_block_m * BLOCK_M + offs_m[:, None])
+            * stride_qbs
+            + cur_head * stride_qh
+            + offs_dpe[None, :]
+        )
+        qpe = tl.load(Q_Extend + offs_qpe, mask=mask_m[:, None], other=0.0)
+
+    # stage 1: compute scores with prefix
+    offs_n = tl.arange(0, BLOCK_N)
+
+    acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32)
+    deno = tl.zeros([BLOCK_M], dtype=tl.float32)
+    e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+
+    for start_n in range(0, cur_seq_len_prefix, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        mask_n = (start_n + offs_n) < cur_seq_len_prefix
+        offs_b_loc_prefix = cur_batch_req_idx * stride_req_to_tokens_b + (
+            cur_seq_prefix_start_in_loc + start_n + offs_n
+        )
+        offs_kv_loc = tl.load(Req_to_tokens + offs_b_loc_prefix, mask=mask_n, other=0)
+
+        # load k in transposed way
+        offs_buf_k = (
+            offs_kv_loc[None, :] * stride_buf_kbs
+            + cur_kv_head * stride_buf_kh
+            + offs_d[:, None]
+        )
+        k = tl.load(
+            K_Buffer + offs_buf_k, mask=(mask_n[None, :]) & (mask_d[:, None]), other=0.0
+        )
+
+        qk = tl.dot(q.to(k.dtype), k)
+        if BLOCK_DPE > 0:
+            offs_kpe = (
+                offs_kv_loc[None, :] * stride_buf_kbs
+                + cur_kv_head * stride_buf_kh
+                + offs_dpe[:, None]
+            )
+            kpe = tl.load(
+                K_Buffer + offs_kpe,
+                mask=mask_n[None, :],
+                other=0.0,
+            )
+            qk += tl.dot(qpe.to(kpe.dtype), kpe)
+        qk *= sm_scale
+
+        if logit_cap > 0:
+            qk = logit_cap * tanh(qk / logit_cap)
+
+        qk = tl.where(mask_m[:, None] & mask_n[None, :], qk, float("-inf"))
+
+        n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+        re_scale = tl.exp(e_max - n_e_max)
+        p = tl.exp(qk - n_e_max[:, None])
+        deno = deno * re_scale + tl.sum(p, 1)
+
+        offs_buf_v = (
+            offs_kv_loc[:, None] * stride_buf_vbs
+            + cur_kv_head * stride_buf_vh
+            + offs_dv[None, :]
+        )
+        v = tl.load(
+            V_Buffer + offs_buf_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0
+        )
+        p = p.to(v.dtype)
+        acc = acc * re_scale[:, None] + tl.dot(p, v)
+
+        e_max = n_e_max
+
+    # stage 2: compute the trianlge part
+
+    cur_block_m_end = tl.minimum(cur_seq_len_extend, (cur_block_m + 1) * BLOCK_M)
+    for start_n in range(0, cur_block_m_end, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        mask_n = (start_n + offs_n) < cur_block_m_end
+
+        # load k in transposed way
+        offs_k = (
+            (cur_seq_extend_start_contiguous + start_n + offs_n[None, :]) * stride_kbs
+            + cur_kv_head * stride_kh
+            + offs_d[:, None]
+        )
+        k = tl.load(
+            K_Extend + offs_k, mask=(mask_n[None, :]) & (mask_d[:, None]), other=0.0
+        )
+
+        qk = tl.dot(q, k, out_dtype=tl.float32)
+        if BLOCK_DPE > 0:
+            offs_kpe = (
+                (cur_seq_extend_start_contiguous + start_n + offs_n[None, :])
+                * stride_kbs
+                + cur_kv_head * stride_kh
+                + offs_dpe[:, None]
+            )
+            kpe = tl.load(
+                K_Extend + offs_kpe,
+                mask=mask_n[None, :],
+                other=0.0,
+            )
+            qk += tl.dot(qpe, kpe)
+
+        qk *= sm_scale
+
+        if logit_cap > 0:
+            qk = logit_cap * tanh(qk / logit_cap)
+
+        mask_causual = (cur_block_m * BLOCK_M + offs_m[:, None]) >= (
+            start_n + offs_n[None, :]
+        )
+        mask_causual &= mask_m[:, None] & mask_n[None, :]
+        qk = tl.where(mask_causual, qk, float("-inf"))
+
+        n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+        re_scale = tl.exp(e_max - n_e_max)
+        p = tl.exp(qk - n_e_max[:, None])
+        deno = deno * re_scale + tl.sum(p, 1)
+
+        offs_v = (
+            (cur_seq_extend_start_contiguous + start_n + offs_n[:, None]) * stride_vbs
+            + cur_kv_head * stride_vh
+            + offs_dv[None, :]
+        )
+        v = tl.load(
+            V_Extend + offs_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0
+        )
+        p = p.to(v.dtype)
+        acc = acc * re_scale[:, None] + tl.dot(p, v)
+
+        e_max = n_e_max
+
+    offs_o = (
+        (cur_seq_extend_start_contiguous + cur_block_m * BLOCK_M + offs_m[:, None])
+        * stride_obs
+        + cur_head * stride_oh
+        + offs_dv[None, :]
+    )
+    tl.store(
+        O_Extend + offs_o, acc / deno[:, None], mask=mask_m[:, None] & mask_dv[None, :]
+    )
+
+
+def extend_attention_fwd(
+    q_extend,
+    k_extend,
+    v_extend,
+    o_extend,
+    k_buffer,
+    v_buffer,
+    req_to_tokens,
+    b_req_idx,
+    b_seq_len,
+    b_seq_len_extend,
+    b_start_loc_extend,
+    max_len_extend,
+    sm_scale=None,
+    logit_cap=0.0,
+):
+    """
+    q_extend, k_extend, v_extend, o_extend: contiguous tensors
+
+    k_buffer, v_buffer: (prefix + extend) tensors in mem_manager
+    """
+    Lq, Lk, Lv = (
+        q_extend.shape[-1],
+        k_extend.shape[-1],
+        v_extend.shape[-1],
+    )
+
+    if Lq == 576:
+        BLOCK_DMODEL = 512
+        BLOCK_DPE = 64
+    elif Lq == 288:
+        BLOCK_DMODEL = 256
+        BLOCK_DPE = 32
+    elif Lq == 192:
+        BLOCK_DMODEL = 128
+        BLOCK_DPE = 64
+    else:
+        BLOCK_DMODEL = triton.next_power_of_2(Lq)
+        BLOCK_DPE = 0
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    if _is_hip:
+        BLOCK_M, BLOCK_N = (64, 64)
+        num_warps = 4
+
+    else:
+        if _is_cuda and CUDA_CAPABILITY[0] >= 9:
+            if Lq <= 256:
+                BLOCK_M, BLOCK_N = (128, 64)
+            else:
+                BLOCK_M, BLOCK_N = (32, 64)
+        elif _is_cuda and CUDA_CAPABILITY[0] >= 8:
+            if Lq <= 128:
+                BLOCK_M, BLOCK_N = (128, 128)
+            elif Lq <= 256:
+                BLOCK_M, BLOCK_N = (64, 64)
+            else:
+                BLOCK_M, BLOCK_N = (32, 64)
+        else:
+            BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32)
+
+        num_warps = 4 if Lk <= 64 else 8
+
+    sm_scale = sm_scale or 1.0 / (Lq**0.5)
+    batch_size, head_num = b_seq_len.shape[0], q_extend.shape[1]
+    kv_group_num = q_extend.shape[1] // k_extend.shape[1]
+
+    grid = (batch_size, head_num, triton.cdiv(max_len_extend, BLOCK_M))
+    num_stages = 1
+
+    extra_kargs = {}
+    if _is_hip:
+        extra_kargs = {"waves_per_eu": 4, "matrix_instr_nonkdim": 16, "kpack": 2}
+
+    _fwd_kernel[grid](
+        q_extend,
+        k_extend,
+        v_extend,
+        o_extend,
+        k_buffer,
+        v_buffer,
+        req_to_tokens,
+        b_req_idx,
+        b_seq_len,
+        b_start_loc_extend,
+        b_seq_len_extend,
+        sm_scale,
+        kv_group_num,
+        q_extend.stride(0),
+        q_extend.stride(1),
+        k_extend.stride(0),
+        k_extend.stride(1),
+        v_extend.stride(0),
+        v_extend.stride(1),
+        o_extend.stride(0),
+        o_extend.stride(1),
+        k_buffer.stride(0),
+        k_buffer.stride(1),
+        v_buffer.stride(0),
+        v_buffer.stride(1),
+        req_to_tokens.stride(0),
+        logit_cap=logit_cap,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DPE=BLOCK_DPE,
+        BLOCK_DV=BLOCK_DV,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        Lq=Lq,
+        Lv=Lv,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        **extra_kargs,
+    )

lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/extend_attention.py

@@ -0,0 +1,467 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""
+Memory-efficient attention for prefill.
+It supports page size = 1 and prefill with KV cache (i.e. extend).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.attention.triton_ops.prefill_attention import (
+    context_attention_fwd,
+)
+from sglang.srt.utils import is_cuda, is_hip
+
+_is_cuda = is_cuda()
+if _is_cuda:
+    CUDA_CAPABILITY = torch.cuda.get_device_capability()
+
+_is_hip = is_hip()
+
+
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+
+@triton.jit
+def _fwd_kernel(
+    Q_Extend,
+    K_Extend,
+    V_Extend,
+    O_Extend,
+    K_Buffer,
+    V_Buffer,
+    qo_indptr,
+    kv_indptr,
+    kv_indices,
+    mask_ptr,
+    mask_indptr,
+    sm_scale,
+    kv_group_num,
+    stride_qbs,
+    stride_qh,
+    stride_kbs,
+    stride_kh,
+    stride_vbs,
+    stride_vh,
+    stride_obs,
+    stride_oh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    stride_buf_vbs,
+    stride_buf_vh,
+    logit_cap: tl.constexpr,
+    Lq: tl.constexpr,
+    Lv: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DPE: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    USE_CUSTOM_MASK: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    SKIP_PREFIX_CUSTOM_MASK: tl.constexpr,
+    STORE_TRANSPOSE: tl.constexpr,
+):
+    cur_seq = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    cur_block_m = tl.program_id(2)
+    cur_kv_head = cur_head // kv_group_num
+
+    cur_seq_extend_start_idx = tl.load(qo_indptr + cur_seq)
+    cur_seq_len_extend = tl.load(qo_indptr + cur_seq + 1) - cur_seq_extend_start_idx
+    cur_seq_kv_start_idx = tl.load(kv_indptr + cur_seq)
+    cur_seq_len_prefix = tl.load(kv_indptr + cur_seq + 1) - cur_seq_kv_start_idx
+    cur_seq_len = cur_seq_len_prefix + cur_seq_len_extend
+
+    if USE_CUSTOM_MASK:
+        cur_seq_mask_start_idx = tl.load(mask_indptr + cur_seq)
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    offs_m = tl.arange(0, BLOCK_M)
+    mask_m = (cur_block_m * BLOCK_M + offs_m) < cur_seq_len_extend
+
+    mask_d = offs_d < Lq
+    mask_dv = offs_dv < Lv
+
+    offs_q = (
+        (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None])
+        * stride_qbs
+        + cur_head * stride_qh
+        + offs_d[None, :]
+    )
+    q = tl.load(
+        Q_Extend + offs_q, mask=(mask_m[:, None]) & (mask_d[None, :]), other=0.0
+    )
+
+    if BLOCK_DPE > 0:
+        offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
+        offs_qpe = (
+            (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None])
+            * stride_qbs
+            + cur_head * stride_qh
+            + offs_dpe[None, :]
+        )
+        qpe = tl.load(Q_Extend + offs_qpe, mask=mask_m[:, None], other=0.0)
+
+    # stage 1: compute scores with prefix
+    offs_n = tl.arange(0, BLOCK_N)
+
+    acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32)
+    deno = tl.zeros([BLOCK_M], dtype=tl.float32)
+    e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+
+    for start_n in range(0, cur_seq_len_prefix, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        mask_n = (start_n + offs_n) < cur_seq_len_prefix
+
+        offs_kv_loc = tl.load(
+            kv_indices + cur_seq_kv_start_idx + start_n + offs_n, mask=mask_n, other=0
+        )
+
+        # load k in transposed way
+        offs_buf_k = (
+            offs_kv_loc[None, :] * stride_buf_kbs
+            + cur_kv_head * stride_buf_kh
+            + offs_d[:, None]
+        )
+        k = tl.load(
+            K_Buffer + offs_buf_k, mask=(mask_n[None, :]) & (mask_d[:, None]), other=0.0
+        )
+
+        qk = tl.dot(q.to(k.dtype), k)
+        if BLOCK_DPE > 0:
+            offs_kpe = (
+                offs_kv_loc[None, :] * stride_buf_kbs
+                + cur_kv_head * stride_buf_kh
+                + offs_dpe[:, None]
+            )
+            kpe = tl.load(
+                K_Buffer + offs_kpe,
+                mask=mask_n[None, :],
+                other=0.0,
+            )
+            qk += tl.dot(qpe.to(kpe.dtype), kpe)
+        qk *= sm_scale
+
+        if logit_cap > 0:
+            qk = logit_cap * tanh(qk / logit_cap)
+
+        if USE_CUSTOM_MASK and not SKIP_PREFIX_CUSTOM_MASK:
+            custom_mask = tl.load(
+                mask_ptr
+                + cur_seq_mask_start_idx
+                + (cur_block_m * BLOCK_M + offs_m[:, None]) * cur_seq_len
+                + start_n
+                + offs_n[None, :],
+                mask=(mask_m[:, None] & mask_n[None, :]),
+                other=0,
+            )
+            custom_mask &= mask_m[:, None] & mask_n[None, :]
+            qk = tl.where(custom_mask, qk, float("-inf"))
+        else:
+            qk = tl.where(mask_m[:, None] & mask_n[None, :], qk, float("-inf"))
+
+        n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+        re_scale = tl.exp(e_max - n_e_max)
+        p = tl.exp(qk - n_e_max[:, None])
+        deno = deno * re_scale + tl.sum(p, 1)
+
+        offs_buf_v = (
+            offs_kv_loc[:, None] * stride_buf_vbs
+            + cur_kv_head * stride_buf_vh
+            + offs_dv[None, :]
+        )
+        v = tl.load(
+            V_Buffer + offs_buf_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0
+        )
+        p = p.to(v.dtype)
+        acc = acc * re_scale[:, None] + tl.dot(p, v)
+
+        e_max = n_e_max
+
+    # stage 2: compute the triangle part
+
+    cur_block_m_end = (
+        cur_seq_len_extend
+        if not IS_CAUSAL
+        else tl.minimum(cur_seq_len_extend, (cur_block_m + 1) * BLOCK_M)
+    )
+    for start_n in range(0, cur_block_m_end, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        mask_n = (start_n + offs_n) < cur_block_m_end
+
+        # load k in transposed way
+        offs_k = (
+            (cur_seq_extend_start_idx + start_n + offs_n[None, :]) * stride_kbs
+            + cur_kv_head * stride_kh
+            + offs_d[:, None]
+        )
+        k = tl.load(
+            K_Extend + offs_k, mask=(mask_n[None, :]) & (mask_d[:, None]), other=0.0
+        )
+
+        qk = tl.dot(q, k, out_dtype=tl.float32)
+        if BLOCK_DPE > 0:
+            offs_kpe = (
+                (cur_seq_extend_start_idx + start_n + offs_n[None, :]) * stride_kbs
+                + cur_kv_head * stride_kh
+                + offs_dpe[:, None]
+            )
+            kpe = tl.load(
+                K_Extend + offs_kpe,
+                mask=mask_n[None, :],
+                other=0.0,
+            )
+            qk += tl.dot(qpe, kpe)
+
+        qk *= sm_scale
+
+        if logit_cap > 0:
+            qk = logit_cap * tanh(qk / logit_cap)
+
+        if USE_CUSTOM_MASK:
+            custom_mask = tl.load(
+                mask_ptr
+                + cur_seq_mask_start_idx
+                + (cur_block_m * BLOCK_M + offs_m[:, None]) * cur_seq_len
+                + cur_seq_len_prefix
+                + start_n
+                + offs_n[None, :],
+                mask=(mask_m[:, None] & mask_n[None, :]),
+                other=0,
+            )
+            custom_mask &= mask_m[:, None] & mask_n[None, :]
+            qk = tl.where(custom_mask, qk, float("-inf"))
+        elif IS_CAUSAL:
+            mask_causual = (cur_block_m * BLOCK_M + offs_m[:, None]) >= (
+                start_n + offs_n[None, :]
+            )
+            mask_causual &= mask_m[:, None] & mask_n[None, :]
+            qk = tl.where(mask_causual, qk, float("-inf"))
+        else:
+            mask_non_causal = mask_m[:, None] & mask_n[None, :]
+            qk = tl.where(mask_non_causal, qk, float("-inf"))
+
+        n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+        re_scale = tl.exp(e_max - n_e_max)
+        p = tl.exp(qk - n_e_max[:, None])
+        deno = deno * re_scale + tl.sum(p, 1)
+
+        offs_v = (
+            (cur_seq_extend_start_idx + start_n + offs_n[:, None]) * stride_vbs
+            + cur_kv_head * stride_vh
+            + offs_dv[None, :]
+        )
+        v = tl.load(
+            V_Extend + offs_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0
+        )
+        p = p.to(v.dtype)
+        acc = acc * re_scale[:, None] + tl.dot(p, v)
+
+        e_max = n_e_max
+
+    offs_o = (
+        (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None])
+        * stride_obs
+        + cur_head * stride_oh
+        + offs_dv[None, :]
+    )
+    if STORE_TRANSPOSE:
+        tl.store(
+            O_Extend + offs_o.T,
+            (acc / deno[:, None]).T,
+            mask=(mask_m[:, None] & mask_dv[None, :]).T,
+        )
+    else:
+        tl.store(
+            O_Extend + offs_o,
+            acc / deno[:, None],
+            mask=mask_m[:, None] & mask_dv[None, :],
+        )
+
+
+def extend_attention_fwd(
+    q_extend,
+    k_extend,
+    v_extend,
+    o_extend,
+    k_buffer,
+    v_buffer,
+    qo_indptr,
+    kv_indptr,
+    kv_indices,
+    custom_mask,
+    is_causal,
+    mask_indptr,
+    max_len_extend,
+    sm_scale=None,
+    logit_cap=0.0,
+    skip_prefix_custom_mask=True,
+):
+    """
+    q_extend, k_extend, v_extend, o_extend: contiguous tensors
+
+    k_buffer, v_buffer: (prefix + extend) tensors in mem_manager
+    """
+    Lq, Lk, Lv = (
+        q_extend.shape[-1],
+        k_extend.shape[-1],
+        v_extend.shape[-1],
+    )
+
+    if Lq == 576:
+        BLOCK_DMODEL = 512
+        BLOCK_DPE = 64
+    elif Lq == 288:
+        BLOCK_DMODEL = 256
+        BLOCK_DPE = 32
+    elif Lq == 192:
+        BLOCK_DMODEL = 128
+        BLOCK_DPE = 64
+    else:
+        BLOCK_DMODEL = triton.next_power_of_2(Lq)
+        BLOCK_DPE = 0
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    if _is_hip:
+        BLOCK_M, BLOCK_N = (64, 64)
+        num_warps = 4
+
+    else:
+        if _is_cuda and CUDA_CAPABILITY[0] >= 9:
+            if Lq <= 256:
+                BLOCK_M, BLOCK_N = (128, 64)
+            else:
+                BLOCK_M, BLOCK_N = (32, 64)
+        elif _is_cuda and CUDA_CAPABILITY[0] >= 8:
+            # sm86/sm89 has a much smaller shared memory size (100K) than sm80 (160K)
+            if CUDA_CAPABILITY[1] == 9 or CUDA_CAPABILITY[1] == 6:
+                if Lq <= 128:
+                    BLOCK_M, BLOCK_N = (64, 128)
+                elif Lq <= 256:
+                    BLOCK_M, BLOCK_N = (64, 64)
+                else:
+                    BLOCK_M, BLOCK_N = (32, 32)
+            else:
+                if Lq <= 128:
+                    BLOCK_M, BLOCK_N = (128, 128)
+                elif Lq <= 256:
+                    BLOCK_M, BLOCK_N = (64, 64)
+                else:
+                    BLOCK_M, BLOCK_N = (32, 64)
+        else:
+            BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32)
+
+        num_warps = 4 if Lk <= 64 else 8
+
+    sm_scale = sm_scale or 1.0 / (Lq**0.5)
+    batch_size, head_num = qo_indptr.shape[0] - 1, q_extend.shape[1]
+    kv_group_num = q_extend.shape[1] // k_extend.shape[1]
+
+    USE_CUSTOM_MASK = custom_mask is not None
+    # Skip custom mask for prefix part
+    SKIP_PREFIX_CUSTOM_MASK = skip_prefix_custom_mask
+
+    grid = (batch_size, head_num, triton.cdiv(max_len_extend, BLOCK_M))
+    num_stages = 1
+
+    extra_kargs = {}
+    if _is_hip:
+        extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}
+
+    _fwd_kernel[grid](
+        q_extend,
+        k_extend,
+        v_extend,
+        o_extend,
+        k_buffer,
+        v_buffer,
+        qo_indptr,
+        kv_indptr,
+        kv_indices,
+        custom_mask,
+        mask_indptr,
+        sm_scale,
+        kv_group_num,
+        q_extend.stride(0),
+        q_extend.stride(1),
+        k_extend.stride(0),
+        k_extend.stride(1),
+        v_extend.stride(0),
+        v_extend.stride(1),
+        o_extend.stride(0),
+        o_extend.stride(1),
+        k_buffer.stride(0),
+        k_buffer.stride(1),
+        v_buffer.stride(0),
+        v_buffer.stride(1),
+        logit_cap=logit_cap,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DPE=BLOCK_DPE,
+        BLOCK_DV=BLOCK_DV,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        Lq=Lq,
+        Lv=Lv,
+        USE_CUSTOM_MASK=USE_CUSTOM_MASK,
+        IS_CAUSAL=is_causal,
+        SKIP_PREFIX_CUSTOM_MASK=SKIP_PREFIX_CUSTOM_MASK,
+        STORE_TRANSPOSE=_is_hip,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        **extra_kargs,
+    )
+
+
+def redundant_attention(
+    q_extend,
+    o_extend,
+    k_buffer,
+    v_buffer,
+    b_req_idx,
+    b_start_loc,
+    b_seq_len,
+    b_seq_len_prefix,
+    max_len_in_batch,
+):
+    total_token_num = k_buffer.shape[0]
+    B, H_Q, D = b_req_idx.shape[0], q_extend.shape[-2], q_extend.shape[-1]
+    q_buffer = torch.empty(
+        (total_token_num, H_Q, D), dtype=q_extend.dtype, device=q_extend.device
+    )
+
+    pt = 0
+    for i in range(B):
+        cur_seq_len_extend = b_seq_len[i] - b_seq_len_prefix[i]
+        pl, pr = b_start_loc[i] + b_seq_len_prefix[i], b_start_loc[i] + b_seq_len[i]
+        q_buffer[pl:pr] = q_extend[pt : pt + cur_seq_len_extend]
+        pt += cur_seq_len_extend
+
+    o_buffer = torch.empty_like(q_buffer)
+    context_attention_fwd(
+        q_buffer, k_buffer, v_buffer, o_buffer, b_start_loc, b_seq_len, max_len_in_batch
+    )
+
+    pt = 0
+    for i in range(B):
+        cur_seq_len_extend = b_seq_len[i] - b_seq_len_prefix[i]
+        pl, pr = b_start_loc[i] + b_seq_len_prefix[i], b_start_loc[i] + b_seq_len[i]
+        o_extend[pt : pt + cur_seq_len_extend] = o_buffer[pl:pr]
+        pt += cur_seq_len_extend