diff --git a/lib/python3.12/site-packages/matplotlib/_docstring.pyi b/lib/python3.12/site-packages/matplotlib/_docstring.pyi new file mode 100644 index 0000000000000000000000000000000000000000..fb52d084612399f399e6bcd3f07bca85bb010ba0 --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/matplotlib/_enums.py b/lib/python3.12/site-packages/matplotlib/_enums.py new file mode 100644 index 0000000000000000000000000000000000000000..75a09b7b5d8c7f9097e965d76fbae7bab2b41a01 --- /dev/null +++ b/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 + `, 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 + `_ + + .. 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 + ` 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, +) diff --git a/lib/python3.12/site-packages/matplotlib/_type1font.py b/lib/python3.12/site-packages/matplotlib/_type1font.py new file mode 100644 index 0000000000000000000000000000000000000000..b3e08f52c035e65837e489bf79da60d28e3ac0a8 --- /dev/null +++ b/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('> 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', +} diff --git a/lib/python3.12/site-packages/matplotlib/backend_managers.pyi b/lib/python3.12/site-packages/matplotlib/backend_managers.pyi new file mode 100644 index 0000000000000000000000000000000000000000..9e59acb14eda98b25e2dcb3f6e2004ca99d0ae5e --- /dev/null +++ b/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: ... diff --git a/lib/python3.12/site-packages/matplotlib/category.py b/lib/python3.12/site-packages/matplotlib/category.py new file mode 100644 index 0000000000000000000000000000000000000000..225c837006f726aa97b990c82a3265ec06e55df3 --- /dev/null +++ b/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() diff --git a/lib/python3.12/site-packages/matplotlib/cbook.py b/lib/python3.12/site-packages/matplotlib/cbook.py new file mode 100644 index 0000000000000000000000000000000000000000..b40f6549c9457d20f0263e0cf7c0ce1b7241dbec --- /dev/null +++ b/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 :: + + [, + , + ] + + one will get :: + + + + 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"" + else: + return "" + + +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 diff --git a/lib/python3.12/site-packages/matplotlib/dviread.py b/lib/python3.12/site-packages/matplotlib/dviread.py new file mode 100644 index 0000000000000000000000000000000000000000..bd21367ce73dc39658ab3c671c2ec2ebba5341f3 --- /dev/null +++ b/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 + # # if using chemformula + # down + # push + # down + # # 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 + 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") diff --git a/lib/python3.12/site-packages/matplotlib/dviread.pyi b/lib/python3.12/site-packages/matplotlib/dviread.pyi new file mode 100644 index 0000000000000000000000000000000000000000..8073ee9fbff852bc3583cb79ade6be5aab7e58d5 --- /dev/null +++ b/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: ... diff --git a/lib/python3.12/site-packages/matplotlib/inset.py b/lib/python3.12/site-packages/matplotlib/inset.py new file mode 100644 index 0000000000000000000000000000000000000000..bab69491303e698ff6ec05a7d95ae88d1d213ac3 --- /dev/null +++ b/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] diff --git a/lib/python3.12/site-packages/matplotlib/layout_engine.py b/lib/python3.12/site-packages/matplotlib/layout_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..8a3276b53371b5a99a3226399e281e4f184258a6 --- /dev/null +++ b/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] diff --git a/lib/python3.12/site-packages/matplotlib/legend.py b/lib/python3.12/site-packages/matplotlib/legend.py new file mode 100644 index 0000000000000000000000000000000000000000..0c2dfc19705ca03fa7887413f3c1c9d643e3b1a0 --- /dev/null +++ b/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 + `. + +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`` 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 diff --git a/lib/python3.12/site-packages/matplotlib/lines.py b/lib/python3.12/site-packages/matplotlib/lines.py new file mode 100644 index 0000000000000000000000000000000000000000..21dd91b89f49348c0165ef351e6f1494fed549e3 --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/matplotlib/mathtext.py b/lib/python3.12/site-packages/matplotlib/mathtext.py new file mode 100644 index 0000000000000000000000000000000000000000..a88c35c15676b027c773eead1e9d0a8b13270768 --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/matplotlib/offsetbox.pyi b/lib/python3.12/site-packages/matplotlib/offsetbox.pyi new file mode 100644 index 0000000000000000000000000000000000000000..3b1520e171385a334e0b4893de6bf152c379c52a --- /dev/null +++ b/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: ... diff --git a/lib/python3.12/site-packages/matplotlib/pyplot.py b/lib/python3.12/site-packages/matplotlib/pyplot.py new file mode 100644 index 0000000000000000000000000000000000000000..2dd8e76d98b4842a16b101cf1dee2cbce702a5e2 --- /dev/null +++ b/lib/python3.12/site-packages/matplotlib/pyplot.py @@ -0,0 +1,4645 @@ +# Note: The first part of this file can be modified in place, but the latter +# part is autogenerated by the boilerplate.py script. + +""" +`matplotlib.pyplot` is a state-based interface to matplotlib. It provides +an implicit, MATLAB-like, way of plotting. It also opens figures on your +screen, and acts as the figure GUI manager. + +pyplot is mainly intended for interactive plots and simple cases of +programmatic plot generation:: + + import numpy as np + import matplotlib.pyplot as plt + + x = np.arange(0, 5, 0.1) + y = np.sin(x) + plt.plot(x, y) + plt.show() + +The explicit object-oriented API is recommended for complex plots, though +pyplot is still usually used to create the figure and often the Axes in the +figure. See `.pyplot.figure`, `.pyplot.subplots`, and +`.pyplot.subplot_mosaic` to create figures, and +:doc:`Axes API ` for the plotting methods on an Axes:: + + import numpy as np + import matplotlib.pyplot as plt + + x = np.arange(0, 5, 0.1) + y = np.sin(x) + fig, ax = plt.subplots() + ax.plot(x, y) + plt.show() + + +See :ref:`api_interfaces` for an explanation of the tradeoffs between the +implicit and explicit interfaces. +""" + +# fmt: off + +from __future__ import annotations + +from contextlib import AbstractContextManager, ExitStack +from enum import Enum +import functools +import importlib +import inspect +import logging +import sys +import threading +import time +from typing import TYPE_CHECKING, cast, overload + +from cycler import cycler # noqa: F401 +import matplotlib +import matplotlib.colorbar +import matplotlib.image +from matplotlib import _api +# Re-exported (import x as x) for typing. +from matplotlib import get_backend as get_backend, rcParams as rcParams +from matplotlib import cm as cm # noqa: F401 +from matplotlib import style as style # noqa: F401 +from matplotlib import _pylab_helpers +from matplotlib import interactive # noqa: F401 +from matplotlib import cbook +from matplotlib import _docstring +from matplotlib.backend_bases import ( + FigureCanvasBase, FigureManagerBase, MouseButton) +from matplotlib.figure import Figure, FigureBase, figaspect +from matplotlib.gridspec import GridSpec, SubplotSpec +from matplotlib import rcsetup, rcParamsDefault, rcParamsOrig +from matplotlib.artist import Artist +from matplotlib.axes import Axes +from matplotlib.axes import Subplot # noqa: F401 +from matplotlib.backends import BackendFilter, backend_registry +from matplotlib.projections import PolarAxes +from matplotlib.colorizer import _ColorizerInterface, ColorizingArtist, Colorizer +from matplotlib import mlab # for detrend_none, window_hanning +from matplotlib.scale import get_scale_names # noqa: F401 + +from matplotlib.cm import _colormaps +from matplotlib.colors import _color_sequences, Colormap + +import numpy as np + +if TYPE_CHECKING: + from collections.abc import Callable, Hashable, Iterable, Sequence + import datetime + import pathlib + import os + from typing import Any, BinaryIO, Literal, TypeVar + from typing_extensions import ParamSpec + + import PIL.Image + from numpy.typing import ArrayLike + + import matplotlib.axes + import matplotlib.artist + import matplotlib.backend_bases + from matplotlib.axis import Tick + from matplotlib.axes._base import _AxesBase + from matplotlib.backend_bases import Event + from matplotlib.cm import ScalarMappable + from matplotlib.contour import ContourSet, QuadContourSet + from matplotlib.collections import ( + Collection, + FillBetweenPolyCollection, + LineCollection, + PolyCollection, + PathCollection, + EventCollection, + QuadMesh, + ) + from matplotlib.colorbar import Colorbar + from matplotlib.container import ( + BarContainer, + ErrorbarContainer, + StemContainer, + ) + from matplotlib.figure import SubFigure + from matplotlib.legend import Legend + from matplotlib.mlab import GaussianKDE + from matplotlib.image import AxesImage, FigureImage + from matplotlib.patches import FancyArrow, StepPatch, Wedge + from matplotlib.quiver import Barbs, Quiver, QuiverKey + from matplotlib.scale import ScaleBase + from matplotlib.typing import ( + ColorType, + CoordsType, + HashableList, + LineStyleType, + MarkerType, + ) + from matplotlib.widgets import SubplotTool + + _P = ParamSpec('_P') + _R = TypeVar('_R') + _T = TypeVar('_T') + + +# We may not need the following imports here: +from matplotlib.colors import Normalize +from matplotlib.lines import Line2D, AxLine +from matplotlib.text import Text, Annotation +from matplotlib.patches import Arrow, Circle, Rectangle # noqa: F401 +from matplotlib.patches import Polygon +from matplotlib.widgets import Button, Slider, Widget # noqa: F401 + +from .ticker import ( # noqa: F401 + TickHelper, Formatter, FixedFormatter, NullFormatter, FuncFormatter, + FormatStrFormatter, ScalarFormatter, LogFormatter, LogFormatterExponent, + LogFormatterMathtext, Locator, IndexLocator, FixedLocator, NullLocator, + LinearLocator, LogLocator, AutoLocator, MultipleLocator, MaxNLocator) + +_log = logging.getLogger(__name__) + + +# Explicit rename instead of import-as for typing's sake. +colormaps = _colormaps +color_sequences = _color_sequences + + +@overload +def _copy_docstring_and_deprecators( + method: Any, + func: Literal[None] = None +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: ... + + +@overload +def _copy_docstring_and_deprecators( + method: Any, func: Callable[_P, _R]) -> Callable[_P, _R]: ... + + +def _copy_docstring_and_deprecators( + method: Any, + func: Callable[_P, _R] | None = None +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]] | Callable[_P, _R]: + if func is None: + return cast('Callable[[Callable[_P, _R]], Callable[_P, _R]]', + functools.partial(_copy_docstring_and_deprecators, method)) + decorators: list[Callable[[Callable[_P, _R]], Callable[_P, _R]]] = [ + _docstring.copy(method) + ] + # Check whether the definition of *method* includes @_api.rename_parameter + # or @_api.make_keyword_only decorators; if so, propagate them to the + # pyplot wrapper as well. + while hasattr(method, "__wrapped__"): + potential_decorator = _api.deprecation.DECORATORS.get(method) + if potential_decorator: + decorators.append(potential_decorator) + method = method.__wrapped__ + for decorator in decorators[::-1]: + func = decorator(func) + _add_pyplot_note(func, method) + return func + + +_NO_PYPLOT_NOTE = [ + 'FigureBase._gci', # wrapped_func is private + '_AxesBase._sci', # wrapped_func is private + 'Artist.findobj', # not a standard pyplot wrapper because it does not operate + # on the current Figure / Axes. Explanation of relation would + # be more complex and is not too important. +] + + +def _add_pyplot_note(func, wrapped_func): + """ + Add a note to the docstring of *func* that it is a pyplot wrapper. + + The note is added to the "Notes" section of the docstring. If that does + not exist, a "Notes" section is created. In numpydoc, the "Notes" + section is the third last possible section, only potentially followed by + "References" and "Examples". + """ + if not func.__doc__: + return # nothing to do + + qualname = wrapped_func.__qualname__ + if qualname in _NO_PYPLOT_NOTE: + return + + wrapped_func_is_method = True + if "." not in qualname: + # method qualnames are prefixed by the class and ".", e.g. "Axes.plot" + wrapped_func_is_method = False + link = f"{wrapped_func.__module__}.{qualname}" + elif qualname.startswith("Axes."): # e.g. "Axes.plot" + link = ".axes." + qualname + elif qualname.startswith("_AxesBase."): # e.g. "_AxesBase.set_xlabel" + link = ".axes.Axes" + qualname[9:] + elif qualname.startswith("Figure."): # e.g. "Figure.figimage" + link = "." + qualname + elif qualname.startswith("FigureBase."): # e.g. "FigureBase.gca" + link = ".Figure" + qualname[10:] + elif qualname.startswith("FigureCanvasBase."): # "FigureBaseCanvas.mpl_connect" + link = "." + qualname + else: + raise RuntimeError(f"Wrapped method from unexpected class: {qualname}") + + if wrapped_func_is_method: + message = f"This is the :ref:`pyplot wrapper ` for `{link}`." + else: + message = f"This is equivalent to `{link}`." + + # Find the correct insert position: + # - either we already have a "Notes" section into which we can insert + # - or we create one before the next present section. Note that in numpydoc, the + # "Notes" section is the third last possible section, only potentially followed + # by "References" and "Examples". + # - or we append a new "Notes" section at the end. + doc = inspect.cleandoc(func.__doc__) + if "\nNotes\n-----" in doc: + before, after = doc.split("\nNotes\n-----", 1) + elif (index := doc.find("\nReferences\n----------")) != -1: + before, after = doc[:index], doc[index:] + elif (index := doc.find("\nExamples\n--------")) != -1: + before, after = doc[:index], doc[index:] + else: + # No "Notes", "References", or "Examples" --> append to the end. + before = doc + "\n" + after = "" + + func.__doc__ = f"{before}\nNotes\n-----\n\n.. note::\n\n {message}\n{after}" + + +## Global ## + + +# The state controlled by {,un}install_repl_displayhook(). +_ReplDisplayHook = Enum("_ReplDisplayHook", ["NONE", "PLAIN", "IPYTHON"]) +_REPL_DISPLAYHOOK = _ReplDisplayHook.NONE + + +def _draw_all_if_interactive() -> None: + if matplotlib.is_interactive(): + draw_all() + + +def install_repl_displayhook() -> None: + """ + Connect to the display hook of the current shell. + + The display hook gets called when the read-evaluate-print-loop (REPL) of + the shell has finished the execution of a command. We use this callback + to be able to automatically update a figure in interactive mode. + + This works both with IPython and with vanilla python shells. + """ + global _REPL_DISPLAYHOOK + + if _REPL_DISPLAYHOOK is _ReplDisplayHook.IPYTHON: + return + + # See if we have IPython hooks around, if so use them. + # Use ``sys.modules.get(name)`` rather than ``name in sys.modules`` as + # entries can also have been explicitly set to None. + mod_ipython = sys.modules.get("IPython") + if not mod_ipython: + _REPL_DISPLAYHOOK = _ReplDisplayHook.PLAIN + return + ip = mod_ipython.get_ipython() + if not ip: + _REPL_DISPLAYHOOK = _ReplDisplayHook.PLAIN + return + + ip.events.register("post_execute", _draw_all_if_interactive) + _REPL_DISPLAYHOOK = _ReplDisplayHook.IPYTHON + + if mod_ipython.version_info[:2] < (8, 24): + # Use of backend2gui is not needed for IPython >= 8.24 as that functionality + # has been moved to Matplotlib. + # This code can be removed when Python 3.12, the latest version supported by + # IPython < 8.24, reaches end-of-life in late 2028. + from IPython.core.pylabtools import backend2gui + ipython_gui_name = backend2gui.get(get_backend()) + else: + _, ipython_gui_name = backend_registry.resolve_backend(get_backend()) + # trigger IPython's eventloop integration, if available + if ipython_gui_name: + ip.enable_gui(ipython_gui_name) + + +def uninstall_repl_displayhook() -> None: + """Disconnect from the display hook of the current shell.""" + global _REPL_DISPLAYHOOK + if _REPL_DISPLAYHOOK is _ReplDisplayHook.IPYTHON: + from IPython import get_ipython + ip = get_ipython() + ip.events.unregister("post_execute", _draw_all_if_interactive) + _REPL_DISPLAYHOOK = _ReplDisplayHook.NONE + + +draw_all = _pylab_helpers.Gcf.draw_all + + +# Ensure this appears in the pyplot docs. +@_copy_docstring_and_deprecators(matplotlib.set_loglevel) +def set_loglevel(*args, **kwargs) -> None: + return matplotlib.set_loglevel(*args, **kwargs) + + +@_copy_docstring_and_deprecators(Artist.findobj) +def findobj( + o: Artist | None = None, + match: Callable[[Artist], bool] | type[Artist] | None = None, + include_self: bool = True +) -> list[Artist]: + if o is None: + o = gcf() + return o.findobj(match, include_self=include_self) + + +_backend_mod: type[matplotlib.backend_bases._Backend] | None = None + + +def _get_backend_mod() -> type[matplotlib.backend_bases._Backend]: + """ + Ensure that a backend is selected and return it. + + This is currently private, but may be made public in the future. + """ + if _backend_mod is None: + # Use rcParams._get("backend") to avoid going through the fallback + # logic (which will (re)import pyplot and then call switch_backend if + # we need to resolve the auto sentinel) + switch_backend(rcParams._get("backend")) + return cast(type[matplotlib.backend_bases._Backend], _backend_mod) + + +def switch_backend(newbackend: str) -> None: + """ + Set the pyplot backend. + + Switching to an interactive backend is possible only if no event loop for + another interactive backend has started. Switching to and from + non-interactive backends is always possible. + + If the new backend is different than the current backend then all open + Figures will be closed via ``plt.close('all')``. + + Parameters + ---------- + newbackend : str + The case-insensitive name of the backend to use. + + """ + global _backend_mod + # make sure the init is pulled up so we can assign to it later + import matplotlib.backends + + if newbackend is rcsetup._auto_backend_sentinel: + current_framework = cbook._get_running_interactive_framework() + + if (current_framework and + (backend := backend_registry.backend_for_gui_framework( + current_framework))): + candidates = [backend] + else: + candidates = [] + candidates += [ + "macosx", "qtagg", "gtk4agg", "gtk3agg", "tkagg", "wxagg"] + + # Don't try to fallback on the cairo-based backends as they each have + # an additional dependency (pycairo) over the agg-based backend, and + # are of worse quality. + for candidate in candidates: + try: + switch_backend(candidate) + except ImportError: + continue + else: + rcParamsOrig['backend'] = candidate + return + else: + # Switching to Agg should always succeed; if it doesn't, let the + # exception propagate out. + switch_backend("agg") + rcParamsOrig["backend"] = "agg" + return + old_backend = rcParams._get('backend') # get without triggering backend resolution + + module = backend_registry.load_backend_module(newbackend) + canvas_class = module.FigureCanvas + + required_framework = canvas_class.required_interactive_framework + if required_framework is not None: + current_framework = cbook._get_running_interactive_framework() + if (current_framework and required_framework + and current_framework != required_framework): + raise ImportError( + "Cannot load backend {!r} which requires the {!r} interactive " + "framework, as {!r} is currently running".format( + newbackend, required_framework, current_framework)) + + # Load the new_figure_manager() and show() functions from the backend. + + # Classically, backends can directly export these functions. This should + # keep working for backcompat. + new_figure_manager = getattr(module, "new_figure_manager", None) + show = getattr(module, "show", None) + + # In that classical approach, backends are implemented as modules, but + # "inherit" default method implementations from backend_bases._Backend. + # This is achieved by creating a "class" that inherits from + # backend_bases._Backend and whose body is filled with the module globals. + class backend_mod(matplotlib.backend_bases._Backend): + locals().update(vars(module)) + + # However, the newer approach for defining new_figure_manager and + # show is to derive them from canvas methods. In that case, also + # update backend_mod accordingly; also, per-backend customization of + # draw_if_interactive is disabled. + if new_figure_manager is None: + + def new_figure_manager_given_figure(num, figure): + return canvas_class.new_manager(figure, num) + + def new_figure_manager(num, *args, FigureClass=Figure, **kwargs): + fig = FigureClass(*args, **kwargs) + return new_figure_manager_given_figure(num, fig) + + def draw_if_interactive() -> None: + if matplotlib.is_interactive(): + manager = _pylab_helpers.Gcf.get_active() + if manager: + manager.canvas.draw_idle() + + backend_mod.new_figure_manager_given_figure = ( # type: ignore[method-assign] + new_figure_manager_given_figure) + backend_mod.new_figure_manager = ( # type: ignore[method-assign] + new_figure_manager) + backend_mod.draw_if_interactive = ( # type: ignore[method-assign] + draw_if_interactive) + + # If the manager explicitly overrides pyplot_show, use it even if a global + # show is already present, as the latter may be here for backcompat. + manager_class = getattr(canvas_class, "manager_class", None) + # We can't compare directly manager_class.pyplot_show and FMB.pyplot_show because + # pyplot_show is a classmethod so the above constructs are bound classmethods, and + # thus always different (being bound to different classes). We also have to use + # getattr_static instead of vars as manager_class could have no __dict__. + manager_pyplot_show = inspect.getattr_static(manager_class, "pyplot_show", None) + base_pyplot_show = inspect.getattr_static(FigureManagerBase, "pyplot_show", None) + if (show is None + or (manager_pyplot_show is not None + and manager_pyplot_show != base_pyplot_show)): + if not manager_pyplot_show: + raise ValueError( + f"Backend {newbackend} defines neither FigureCanvas.manager_class nor " + f"a toplevel show function") + _pyplot_show = cast('Any', manager_class).pyplot_show + backend_mod.show = _pyplot_show # type: ignore[method-assign] + + _log.debug("Loaded backend %s version %s.", + newbackend, backend_mod.backend_version) + + if newbackend in ("ipympl", "widget"): + # ipympl < 0.9.4 expects rcParams["backend"] to be the fully-qualified backend + # name "module://ipympl.backend_nbagg" not short names "ipympl" or "widget". + import importlib.metadata as im + from matplotlib import _parse_to_version_info # type: ignore[attr-defined] + try: + module_version = im.version("ipympl") + if _parse_to_version_info(module_version) < (0, 9, 4): + newbackend = "module://ipympl.backend_nbagg" + except im.PackageNotFoundError: + pass + + rcParams['backend'] = rcParamsDefault['backend'] = newbackend + _backend_mod = backend_mod + for func_name in ["new_figure_manager", "draw_if_interactive", "show"]: + globals()[func_name].__signature__ = inspect.signature( + getattr(backend_mod, func_name)) + + # Need to keep a global reference to the backend for compatibility reasons. + # See https://github.com/matplotlib/matplotlib/issues/6092 + matplotlib.backends.backend = newbackend # type: ignore[attr-defined] + + # Make sure the repl display hook is installed in case we become interactive. + install_repl_displayhook() + + +def _warn_if_gui_out_of_main_thread() -> None: + warn = False + canvas_class = cast(type[FigureCanvasBase], _get_backend_mod().FigureCanvas) + if canvas_class.required_interactive_framework: + if hasattr(threading, 'get_native_id'): + # This compares native thread ids because even if Python-level + # Thread objects match, the underlying OS thread (which is what + # really matters) may be different on Python implementations with + # green threads. + if threading.get_native_id() != threading.main_thread().native_id: + warn = True + else: + # Fall back to Python-level Thread if native IDs are unavailable, + # mainly for PyPy. + if threading.current_thread() is not threading.main_thread(): + warn = True + if warn: + _api.warn_external( + "Starting a Matplotlib GUI outside of the main thread will likely " + "fail.") + + +# This function's signature is rewritten upon backend-load by switch_backend. +def new_figure_manager(*args, **kwargs): + """Create a new figure manager instance.""" + _warn_if_gui_out_of_main_thread() + return _get_backend_mod().new_figure_manager(*args, **kwargs) + + +# This function's signature is rewritten upon backend-load by switch_backend. +def draw_if_interactive(*args, **kwargs): + """ + Redraw the current figure if in interactive mode. + + .. warning:: + + End users will typically not have to call this function because the + the interactive mode takes care of this. + """ + return _get_backend_mod().draw_if_interactive(*args, **kwargs) + + +# This function's signature is rewritten upon backend-load by switch_backend. +def show(*args, **kwargs) -> None: + """ + Display all open figures. + + Parameters + ---------- + block : bool, optional + Whether to wait for all figures to be closed before returning. + + If `True` block and run the GUI main loop until all figure windows + are closed. + + If `False` ensure that all figure windows are displayed and return + immediately. In this case, you are responsible for ensuring + that the event loop is running to have responsive figures. + + Defaults to True in non-interactive mode and to False in interactive + mode (see `.pyplot.isinteractive`). + + See Also + -------- + ion : Enable interactive mode, which shows / updates the figure after + every plotting command, so that calling ``show()`` is not necessary. + ioff : Disable interactive mode. + savefig : Save the figure to an image file instead of showing it on screen. + + Notes + ----- + **Saving figures to file and showing a window at the same time** + + If you want an image file as well as a user interface window, use + `.pyplot.savefig` before `.pyplot.show`. At the end of (a blocking) + ``show()`` the figure is closed and thus unregistered from pyplot. Calling + `.pyplot.savefig` afterwards would save a new and thus empty figure. This + limitation of command order does not apply if the show is non-blocking or + if you keep a reference to the figure and use `.Figure.savefig`. + + **Auto-show in jupyter notebooks** + + The jupyter backends (activated via ``%matplotlib inline``, + ``%matplotlib notebook``, or ``%matplotlib widget``), call ``show()`` at + the end of every cell by default. Thus, you usually don't have to call it + explicitly there. + """ + _warn_if_gui_out_of_main_thread() + return _get_backend_mod().show(*args, **kwargs) + + +def isinteractive() -> bool: + """ + Return whether plots are updated after every plotting command. + + The interactive mode is mainly useful if you build plots from the command + line and want to see the effect of each command while you are building the + figure. + + In interactive mode: + + - newly created figures will be shown immediately; + - figures will automatically redraw on change; + - `.pyplot.show` will not block by default. + + In non-interactive mode: + + - newly created figures and changes to figures will not be reflected until + explicitly asked to be; + - `.pyplot.show` will block by default. + + See Also + -------- + ion : Enable interactive mode. + ioff : Disable interactive mode. + show : Show all figures (and maybe block). + pause : Show all figures, and block for a time. + """ + return matplotlib.is_interactive() + + +# Note: The return type of ioff being AbstractContextManager +# instead of ExitStack is deliberate. +# See https://github.com/matplotlib/matplotlib/issues/27659 +# and https://github.com/matplotlib/matplotlib/pull/27667 for more info. +def ioff() -> AbstractContextManager: + """ + Disable interactive mode. + + See `.pyplot.isinteractive` for more details. + + See Also + -------- + ion : Enable interactive mode. + isinteractive : Whether interactive mode is enabled. + show : Show all figures (and maybe block). + pause : Show all figures, and block for a time. + + Notes + ----- + For a temporary change, this can be used as a context manager:: + + # if interactive mode is on + # then figures will be shown on creation + plt.ion() + # This figure will be shown immediately + fig = plt.figure() + + with plt.ioff(): + # interactive mode will be off + # figures will not automatically be shown + fig2 = plt.figure() + # ... + + To enable optional usage as a context manager, this function returns a + context manager object, which is not intended to be stored or + accessed by the user. + """ + stack = ExitStack() + stack.callback(ion if isinteractive() else ioff) + matplotlib.interactive(False) + uninstall_repl_displayhook() + return stack + + +# Note: The return type of ion being AbstractContextManager +# instead of ExitStack is deliberate. +# See https://github.com/matplotlib/matplotlib/issues/27659 +# and https://github.com/matplotlib/matplotlib/pull/27667 for more info. +def ion() -> AbstractContextManager: + """ + Enable interactive mode. + + See `.pyplot.isinteractive` for more details. + + See Also + -------- + ioff : Disable interactive mode. + isinteractive : Whether interactive mode is enabled. + show : Show all figures (and maybe block). + pause : Show all figures, and block for a time. + + Notes + ----- + For a temporary change, this can be used as a context manager:: + + # if interactive mode is off + # then figures will not be shown on creation + plt.ioff() + # This figure will not be shown immediately + fig = plt.figure() + + with plt.ion(): + # interactive mode will be on + # figures will automatically be shown + fig2 = plt.figure() + # ... + + To enable optional usage as a context manager, this function returns a + context manager object, which is not intended to be stored or + accessed by the user. + """ + stack = ExitStack() + stack.callback(ion if isinteractive() else ioff) + matplotlib.interactive(True) + install_repl_displayhook() + return stack + + +def pause(interval: float) -> None: + """ + Run the GUI event loop for *interval* seconds. + + If there is an active figure, it will be updated and displayed before the + pause, and the GUI event loop (if any) will run during the pause. + + This can be used for crude animation. For more complex animation use + :mod:`matplotlib.animation`. + + If there is no active figure, sleep for *interval* seconds instead. + + See Also + -------- + matplotlib.animation : Proper animations + show : Show all figures and optional block until all figures are closed. + """ + manager = _pylab_helpers.Gcf.get_active() + if manager is not None: + canvas = manager.canvas + if canvas.figure.stale: + canvas.draw_idle() + show(block=False) + canvas.start_event_loop(interval) + else: + time.sleep(interval) + + +@_copy_docstring_and_deprecators(matplotlib.rc) +def rc(group: str, **kwargs) -> None: + matplotlib.rc(group, **kwargs) + + +@_copy_docstring_and_deprecators(matplotlib.rc_context) +def rc_context( + rc: dict[str, Any] | None = None, + fname: str | pathlib.Path | os.PathLike | None = None, +) -> AbstractContextManager[None]: + return matplotlib.rc_context(rc, fname) + + +@_copy_docstring_and_deprecators(matplotlib.rcdefaults) +def rcdefaults() -> None: + matplotlib.rcdefaults() + if matplotlib.is_interactive(): + draw_all() + + +# getp/get/setp are explicitly reexported so that they show up in pyplot docs. + + +@_copy_docstring_and_deprecators(matplotlib.artist.getp) +def getp(obj, *args, **kwargs): + return matplotlib.artist.getp(obj, *args, **kwargs) + + +@_copy_docstring_and_deprecators(matplotlib.artist.get) +def get(obj, *args, **kwargs): + return matplotlib.artist.get(obj, *args, **kwargs) + + +@_copy_docstring_and_deprecators(matplotlib.artist.setp) +def setp(obj, *args, **kwargs): + return matplotlib.artist.setp(obj, *args, **kwargs) + + +def xkcd( + scale: float = 1, length: float = 100, randomness: float = 2 +) -> ExitStack: + """ + Turn on `xkcd `_ sketch-style drawing mode. + + This will only have an effect on things drawn after this function is called. + + For best results, install the `xkcd script `_ + font; xkcd fonts are not packaged with Matplotlib. + + Parameters + ---------- + scale : float, optional + The amplitude of the wiggle perpendicular to the source line. + length : float, optional + The length of the wiggle along the line. + randomness : float, optional + The scale factor by which the length is shrunken or expanded. + + Notes + ----- + This function works by a number of rcParams, so it will probably + override others you have set before. + + If you want the effects of this function to be temporary, it can + be used as a context manager, for example:: + + with plt.xkcd(): + # This figure will be in XKCD-style + fig1 = plt.figure() + # ... + + # This figure will be in regular style + fig2 = plt.figure() + """ + # This cannot be implemented in terms of contextmanager() or rc_context() + # because this needs to work as a non-contextmanager too. + + if rcParams['text.usetex']: + raise RuntimeError( + "xkcd mode is not compatible with text.usetex = True") + + stack = ExitStack() + stack.callback(rcParams._update_raw, rcParams.copy()) # type: ignore[arg-type] + + from matplotlib import patheffects + rcParams.update({ + 'font.family': ['xkcd', 'xkcd Script', 'Comic Neue', 'Comic Sans MS'], + 'font.size': 14.0, + 'path.sketch': (scale, length, randomness), + 'path.effects': [ + patheffects.withStroke(linewidth=4, foreground="w")], + 'axes.linewidth': 1.5, + 'lines.linewidth': 2.0, + 'figure.facecolor': 'white', + 'grid.linewidth': 0.0, + 'axes.grid': False, + 'axes.unicode_minus': False, + 'axes.edgecolor': 'black', + 'xtick.major.size': 8, + 'xtick.major.width': 3, + 'ytick.major.size': 8, + 'ytick.major.width': 3, + }) + + return stack + + +## Figures ## + +def figure( + # autoincrement if None, else integer from 1-N + num: int | str | Figure | SubFigure | None = None, + # defaults to rc figure.figsize + figsize: ArrayLike | None = None, + # defaults to rc figure.dpi + dpi: float | None = None, + *, + # defaults to rc figure.facecolor + facecolor: ColorType | None = None, + # defaults to rc figure.edgecolor + edgecolor: ColorType | None = None, + frameon: bool = True, + FigureClass: type[Figure] = Figure, + clear: bool = False, + **kwargs +) -> Figure: + """ + Create a new figure, or activate an existing figure. + + Parameters + ---------- + num : int or str or `.Figure` or `.SubFigure`, optional + A unique identifier for the figure. + + If a figure with that identifier already exists, this figure is made + active and returned. An integer refers to the ``Figure.number`` + attribute, a string refers to the figure label. + + If there is no figure with the identifier or *num* is not given, a new + figure is created, made active and returned. If *num* is an int, it + will be used for the ``Figure.number`` attribute, otherwise, an + auto-generated integer value is used (starting at 1 and incremented + for each new figure). If *num* is a string, the figure label and the + window title is set to this value. If num is a ``SubFigure``, its + parent ``Figure`` is activated. + + figsize : (float, float), default: :rc:`figure.figsize` + Width, height in inches. + + dpi : float, default: :rc:`figure.dpi` + The resolution of the figure in dots-per-inch. + + facecolor : :mpltype:`color`, default: :rc:`figure.facecolor` + The background color. + + edgecolor : :mpltype:`color`, default: :rc:`figure.edgecolor` + The border color. + + frameon : bool, default: True + If False, suppress drawing the figure frame. + + FigureClass : subclass of `~matplotlib.figure.Figure` + If set, an instance of this subclass will be created, rather than a + plain `.Figure`. + + clear : bool, default: False + If True and the figure already exists, then it is cleared. + + layout : {'constrained', 'compressed', 'tight', 'none', `.LayoutEngine`, None}, \ +default: None + The layout mechanism for positioning of plot elements to avoid + overlapping Axes decorations (labels, ticks, etc). Note that layout + managers can measurably slow down figure display. + + - 'constrained': The constrained layout solver adjusts Axes sizes + to avoid overlapping Axes decorations. Can handle complex plot + layouts and colorbars, and is thus recommended. + + See :ref:`constrainedlayout_guide` + for examples. + + - 'compressed': uses the same algorithm as 'constrained', but + removes extra space between fixed-aspect-ratio Axes. Best for + simple grids of Axes. + + - 'tight': Use the tight layout mechanism. This is a relatively + simple algorithm that adjusts the subplot parameters so that + decorations do not overlap. See `.Figure.set_tight_layout` for + further details. + + - 'none': Do not use a layout engine. + + - A `.LayoutEngine` instance. Builtin layout classes are + `.ConstrainedLayoutEngine` and `.TightLayoutEngine`, more easily + accessible by 'constrained' and 'tight'. Passing an instance + allows third parties to provide their own layout engine. + + If not given, fall back to using the parameters *tight_layout* and + *constrained_layout*, including their config defaults + :rc:`figure.autolayout` and :rc:`figure.constrained_layout.use`. + + **kwargs + Additional keyword arguments are passed to the `.Figure` constructor. + + Returns + ------- + `~matplotlib.figure.Figure` + + Notes + ----- + A newly created figure is passed to the `~.FigureCanvasBase.new_manager` + method or the `new_figure_manager` function provided by the current + backend, which install a canvas and a manager on the figure. + + Once this is done, :rc:`figure.hooks` are called, one at a time, on the + figure; these hooks allow arbitrary customization of the figure (e.g., + attaching callbacks) or of associated elements (e.g., modifying the + toolbar). See :doc:`/gallery/user_interfaces/mplcvd` for an example of + toolbar customization. + + If you are creating many figures, make sure you explicitly call + `.pyplot.close` on the figures you are not using, because this will + enable pyplot to properly clean up the memory. + + `~matplotlib.rcParams` defines the default values, which can be modified + in the matplotlibrc file. + """ + allnums = get_fignums() + + if isinstance(num, FigureBase): + # type narrowed to `Figure | SubFigure` by combination of input and isinstance + root_fig = num.get_figure(root=True) + if root_fig.canvas.manager is None: + raise ValueError("The passed figure is not managed by pyplot") + elif (any(param is not None for param in [figsize, dpi, facecolor, edgecolor]) + or not frameon or kwargs) and root_fig.canvas.manager.num in allnums: + _api.warn_external( + "Ignoring specified arguments in this call because figure " + f"with num: {root_fig.canvas.manager.num} already exists") + _pylab_helpers.Gcf.set_active(root_fig.canvas.manager) + return root_fig + + next_num = max(allnums) + 1 if allnums else 1 + fig_label = '' + if num is None: + num = next_num + else: + if (any(param is not None for param in [figsize, dpi, facecolor, edgecolor]) + or not frameon or kwargs) and num in allnums: + _api.warn_external( + "Ignoring specified arguments in this call " + f"because figure with num: {num} already exists") + if isinstance(num, str): + fig_label = num + all_labels = get_figlabels() + if fig_label not in all_labels: + if fig_label == 'all': + _api.warn_external("close('all') closes all existing figures.") + num = next_num + else: + inum = all_labels.index(fig_label) + num = allnums[inum] + else: + num = int(num) # crude validation of num argument + + # Type of "num" has narrowed to int, but mypy can't quite see it + manager = _pylab_helpers.Gcf.get_fig_manager(num) # type: ignore[arg-type] + if manager is None: + max_open_warning = rcParams['figure.max_open_warning'] + if len(allnums) == max_open_warning >= 1: + _api.warn_external( + f"More than {max_open_warning} figures have been opened. " + f"Figures created through the pyplot interface " + f"(`matplotlib.pyplot.figure`) are retained until explicitly " + f"closed and may consume too much memory. (To control this " + f"warning, see the rcParam `figure.max_open_warning`). " + f"Consider using `matplotlib.pyplot.close()`.", + RuntimeWarning) + + manager = new_figure_manager( + num, figsize=figsize, dpi=dpi, + facecolor=facecolor, edgecolor=edgecolor, frameon=frameon, + FigureClass=FigureClass, **kwargs) + fig = manager.canvas.figure + if fig_label: + fig.set_label(fig_label) + + for hookspecs in rcParams["figure.hooks"]: + module_name, dotted_name = hookspecs.split(":") + obj: Any = importlib.import_module(module_name) + for part in dotted_name.split("."): + obj = getattr(obj, part) + obj(fig) + + _pylab_helpers.Gcf._set_new_active_manager(manager) + + # make sure backends (inline) that we don't ship that expect this + # to be called in plotting commands to make the figure call show + # still work. There is probably a better way to do this in the + # FigureManager base class. + draw_if_interactive() + + if _REPL_DISPLAYHOOK is _ReplDisplayHook.PLAIN: + fig.stale_callback = _auto_draw_if_interactive + + if clear: + manager.canvas.figure.clear() + + return manager.canvas.figure + + +def _auto_draw_if_interactive(fig, val): + """ + An internal helper function for making sure that auto-redrawing + works as intended in the plain python repl. + + Parameters + ---------- + fig : Figure + A figure object which is assumed to be associated with a canvas + """ + if (val and matplotlib.is_interactive() + and not fig.canvas.is_saving() + and not fig.canvas._is_idle_drawing): + # Some artists can mark themselves as stale in the middle of drawing + # (e.g. axes position & tick labels being computed at draw time), but + # this shouldn't trigger a redraw because the current redraw will + # already take them into account. + with fig.canvas._idle_draw_cntx(): + fig.canvas.draw_idle() + + +def gcf() -> Figure: + """ + Get the current figure. + + If there is currently no figure on the pyplot figure stack, a new one is + created using `~.pyplot.figure()`. (To test whether there is currently a + figure on the pyplot figure stack, check whether `~.pyplot.get_fignums()` + is empty.) + """ + manager = _pylab_helpers.Gcf.get_active() + if manager is not None: + return manager.canvas.figure + else: + return figure() + + +def fignum_exists(num: int | str) -> bool: + """ + Return whether the figure with the given id exists. + + Parameters + ---------- + num : int or str + A figure identifier. + + Returns + ------- + bool + Whether or not a figure with id *num* exists. + """ + return ( + _pylab_helpers.Gcf.has_fignum(num) + if isinstance(num, int) + else num in get_figlabels() + ) + + +def get_fignums() -> list[int]: + """Return a list of existing figure numbers.""" + return sorted(_pylab_helpers.Gcf.figs) + + +def get_figlabels() -> list[Any]: + """Return a list of existing figure labels.""" + managers = _pylab_helpers.Gcf.get_all_fig_managers() + managers.sort(key=lambda m: m.num) + return [m.canvas.figure.get_label() for m in managers] + + +def get_current_fig_manager() -> FigureManagerBase | None: + """ + Return the figure manager of the current figure. + + The figure manager is a container for the actual backend-depended window + that displays the figure on screen. + + If no current figure exists, a new one is created, and its figure + manager is returned. + + Returns + ------- + `.FigureManagerBase` or backend-dependent subclass thereof + """ + return gcf().canvas.manager + + +@_copy_docstring_and_deprecators(FigureCanvasBase.mpl_connect) +def connect(s: str, func: Callable[[Event], Any]) -> int: + return gcf().canvas.mpl_connect(s, func) + + +@_copy_docstring_and_deprecators(FigureCanvasBase.mpl_disconnect) +def disconnect(cid: int) -> None: + gcf().canvas.mpl_disconnect(cid) + + +def close(fig: None | int | str | Figure | Literal["all"] = None) -> None: + """ + Close a figure window, and unregister it from pyplot. + + Parameters + ---------- + fig : None or int or str or `.Figure` + The figure to close. There are a number of ways to specify this: + + - *None*: the current figure + - `.Figure`: the given `.Figure` instance + - ``int``: a figure number + - ``str``: a figure name + - 'all': all figures + + Notes + ----- + pyplot maintains a reference to figures created with `figure()`. When + work on the figure is completed, it should be closed, i.e. deregistered + from pyplot, to free its memory (see also :rc:figure.max_open_warning). + Closing a figure window created by `show()` automatically deregisters the + figure. For all other use cases, most prominently `savefig()` without + `show()`, the figure must be deregistered explicitly using `close()`. + """ + if fig is None: + manager = _pylab_helpers.Gcf.get_active() + if manager is None: + return + else: + _pylab_helpers.Gcf.destroy(manager) + elif fig == 'all': + _pylab_helpers.Gcf.destroy_all() + elif isinstance(fig, int): + _pylab_helpers.Gcf.destroy(fig) + elif hasattr(fig, 'int'): + # if we are dealing with a type UUID, we + # can use its integer representation + _pylab_helpers.Gcf.destroy(fig.int) + elif isinstance(fig, str): + all_labels = get_figlabels() + if fig in all_labels: + num = get_fignums()[all_labels.index(fig)] + _pylab_helpers.Gcf.destroy(num) + elif isinstance(fig, Figure): + _pylab_helpers.Gcf.destroy_fig(fig) + else: + raise TypeError("close() argument must be a Figure, an int, a string, " + "or None, not %s" % type(fig)) + + +def clf() -> None: + """Clear the current figure.""" + gcf().clear() + + +def draw() -> None: + """ + Redraw the current figure. + + This is used to update a figure that has been altered, but not + automatically re-drawn. If interactive mode is on (via `.ion()`), this + should be only rarely needed, but there may be ways to modify the state of + a figure without marking it as "stale". Please report these cases as bugs. + + This is equivalent to calling ``fig.canvas.draw_idle()``, where ``fig`` is + the current figure. + + See Also + -------- + .FigureCanvasBase.draw_idle + .FigureCanvasBase.draw + """ + gcf().canvas.draw_idle() + + +@_copy_docstring_and_deprecators(Figure.savefig) +def savefig(*args, **kwargs) -> None: + fig = gcf() + # savefig default implementation has no return, so mypy is unhappy + # presumably this is here because subclasses can return? + res = fig.savefig(*args, **kwargs) # type: ignore[func-returns-value] + fig.canvas.draw_idle() # Need this if 'transparent=True', to reset colors. + return res + + +## Putting things in figures ## + + +def figlegend(*args, **kwargs) -> Legend: + return gcf().legend(*args, **kwargs) +if Figure.legend.__doc__: + figlegend.__doc__ = Figure.legend.__doc__ \ + .replace(" legend(", " figlegend(") \ + .replace("fig.legend(", "plt.figlegend(") \ + .replace("ax.plot(", "plt.plot(") + + +## Axes ## + +@_docstring.interpd +def axes( + arg: None | tuple[float, float, float, float] = None, + **kwargs +) -> matplotlib.axes.Axes: + """ + Add an Axes to the current figure and make it the current Axes. + + Call signatures:: + + plt.axes() + plt.axes(rect, projection=None, polar=False, **kwargs) + plt.axes(ax) + + Parameters + ---------- + arg : None or 4-tuple + The exact behavior of this function depends on the type: + + - *None*: A new full window Axes is added using + ``subplot(**kwargs)``. + - 4-tuple of floats *rect* = ``(left, bottom, width, height)``. + A new Axes is added with dimensions *rect* in normalized + (0, 1) units using `~.Figure.add_axes` on the current figure. + + projection : {None, 'aitoff', 'hammer', 'lambert', 'mollweide', \ +'polar', 'rectilinear', str}, optional + The projection type of the `~.axes.Axes`. *str* is the name of + a custom projection, see `~matplotlib.projections`. The default + None results in a 'rectilinear' projection. + + polar : bool, default: False + If True, equivalent to projection='polar'. + + sharex, sharey : `~matplotlib.axes.Axes`, optional + Share the x or y `~matplotlib.axis` with sharex and/or sharey. + The axis will have the same limits, ticks, and scale as the axis + of the shared Axes. + + label : str + A label for the returned Axes. + + Returns + ------- + `~.axes.Axes`, or a subclass of `~.axes.Axes` + The returned Axes class depends on the projection used. It is + `~.axes.Axes` if rectilinear projection is used and + `.projections.polar.PolarAxes` if polar projection is used. + + Other Parameters + ---------------- + **kwargs + This method also takes the keyword arguments for + the returned Axes class. The keyword arguments for the + rectilinear Axes class `~.axes.Axes` can be found in + the following table but there might also be other keyword + arguments if another projection is used, see the actual Axes + class. + + %(Axes:kwdoc)s + + See Also + -------- + .Figure.add_axes + .pyplot.subplot + .Figure.add_subplot + .Figure.subplots + .pyplot.subplots + + Examples + -------- + :: + + # Creating a new full window Axes + plt.axes() + + # Creating a new Axes with specified dimensions and a grey background + plt.axes((left, bottom, width, height), facecolor='grey') + """ + fig = gcf() + pos = kwargs.pop('position', None) + if arg is None: + if pos is None: + return fig.add_subplot(**kwargs) + else: + return fig.add_axes(pos, **kwargs) + else: + return fig.add_axes(arg, **kwargs) + + +def delaxes(ax: matplotlib.axes.Axes | None = None) -> None: + """ + Remove an `~.axes.Axes` (defaulting to the current Axes) from its figure. + """ + if ax is None: + ax = gca() + ax.remove() + + +def sca(ax: Axes) -> None: + """ + Set the current Axes to *ax* and the current Figure to the parent of *ax*. + """ + # Mypy sees ax.figure as potentially None, + # but if you are calling this, it won't be None + # Additionally the slight difference between `Figure` and `FigureBase` mypy catches + fig = ax.get_figure(root=False) + figure(fig) # type: ignore[arg-type] + fig.sca(ax) # type: ignore[union-attr] + + +def cla() -> None: + """Clear the current Axes.""" + # Not generated via boilerplate.py to allow a different docstring. + return gca().cla() + + +## More ways of creating Axes ## + +@_docstring.interpd +def subplot(*args, **kwargs) -> Axes: + """ + Add an Axes to the current figure or retrieve an existing Axes. + + This is a wrapper of `.Figure.add_subplot` which provides additional + behavior when working with the implicit API (see the notes section). + + Call signatures:: + + subplot(nrows, ncols, index, **kwargs) + subplot(pos, **kwargs) + subplot(**kwargs) + subplot(ax) + + Parameters + ---------- + *args : int, (int, int, *index*), or `.SubplotSpec`, default: (1, 1, 1) + The position of the subplot described by one of + + - Three integers (*nrows*, *ncols*, *index*). The subplot will take the + *index* position on a grid with *nrows* rows and *ncols* columns. + *index* starts at 1 in the upper left corner and increases to the + right. *index* can also be a two-tuple specifying the (*first*, + *last*) indices (1-based, and including *last*) of the subplot, e.g., + ``fig.add_subplot(3, 1, (1, 2))`` makes a subplot that spans the + upper 2/3 of the figure. + - A 3-digit integer. The digits are interpreted as if given separately + as three single-digit integers, i.e. ``fig.add_subplot(235)`` is the + same as ``fig.add_subplot(2, 3, 5)``. Note that this can only be used + if there are no more than 9 subplots. + - A `.SubplotSpec`. + + projection : {None, 'aitoff', 'hammer', 'lambert', 'mollweide', \ +'polar', 'rectilinear', str}, optional + The projection type of the subplot (`~.axes.Axes`). *str* is the name + of a custom projection, see `~matplotlib.projections`. The default + None results in a 'rectilinear' projection. + + polar : bool, default: False + If True, equivalent to projection='polar'. + + sharex, sharey : `~matplotlib.axes.Axes`, optional + Share the x or y `~matplotlib.axis` with sharex and/or sharey. The + axis will have the same limits, ticks, and scale as the axis of the + shared Axes. + + label : str + A label for the returned Axes. + + Returns + ------- + `~.axes.Axes` + + The Axes of the subplot. The returned Axes can actually be an instance + of a subclass, such as `.projections.polar.PolarAxes` for polar + projections. + + Other Parameters + ---------------- + **kwargs + This method also takes the keyword arguments for the returned Axes + base class; except for the *figure* argument. The keyword arguments + for the rectilinear base class `~.axes.Axes` can be found in + the following table but there might also be other keyword + arguments if another projection is used. + + %(Axes:kwdoc)s + + Notes + ----- + .. versionchanged:: 3.8 + In versions prior to 3.8, any preexisting Axes that overlap with the new Axes + beyond sharing a boundary was deleted. Deletion does not happen in more + recent versions anymore. Use `.Axes.remove` explicitly if needed. + + If you do not want this behavior, use the `.Figure.add_subplot` method + or the `.pyplot.axes` function instead. + + If no *kwargs* are passed and there exists an Axes in the location + specified by *args* then that Axes will be returned rather than a new + Axes being created. + + If *kwargs* are passed and there exists an Axes in the location + specified by *args*, the projection type is the same, and the + *kwargs* match with the existing Axes, then the existing Axes is + returned. Otherwise a new Axes is created with the specified + parameters. We save a reference to the *kwargs* which we use + for this comparison. If any of the values in *kwargs* are + mutable we will not detect the case where they are mutated. + In these cases we suggest using `.Figure.add_subplot` and the + explicit Axes API rather than the implicit pyplot API. + + See Also + -------- + .Figure.add_subplot + .pyplot.subplots + .pyplot.axes + .Figure.subplots + + Examples + -------- + :: + + plt.subplot(221) + + # equivalent but more general + ax1 = plt.subplot(2, 2, 1) + + # add a subplot with no frame + ax2 = plt.subplot(222, frameon=False) + + # add a polar subplot + plt.subplot(223, projection='polar') + + # add a red subplot that shares the x-axis with ax1 + plt.subplot(224, sharex=ax1, facecolor='red') + + # delete ax2 from the figure + plt.delaxes(ax2) + + # add ax2 to the figure again + plt.subplot(ax2) + + # make the first Axes "current" again + plt.subplot(221) + + """ + # Here we will only normalize `polar=True` vs `projection='polar'` and let + # downstream code deal with the rest. + unset = object() + projection = kwargs.get('projection', unset) + polar = kwargs.pop('polar', unset) + if polar is not unset and polar: + # if we got mixed messages from the user, raise + if projection is not unset and projection != 'polar': + raise ValueError( + f"polar={polar}, yet projection={projection!r}. " + "Only one of these arguments should be supplied." + ) + kwargs['projection'] = projection = 'polar' + + # if subplot called without arguments, create subplot(1, 1, 1) + if len(args) == 0: + args = (1, 1, 1) + + # This check was added because it is very easy to type subplot(1, 2, False) + # when subplots(1, 2, False) was intended (sharex=False, that is). In most + # cases, no error will ever occur, but mysterious behavior can result + # because what was intended to be the sharex argument is instead treated as + # a subplot index for subplot() + if len(args) >= 3 and isinstance(args[2], bool): + _api.warn_external("The subplot index argument to subplot() appears " + "to be a boolean. Did you intend to use " + "subplots()?") + # Check for nrows and ncols, which are not valid subplot args: + if 'nrows' in kwargs or 'ncols' in kwargs: + raise TypeError("subplot() got an unexpected keyword argument 'ncols' " + "and/or 'nrows'. Did you intend to call subplots()?") + + fig = gcf() + + # First, search for an existing subplot with a matching spec. + key = SubplotSpec._from_subplot_args(fig, args) + + for ax in fig.axes: + # If we found an Axes at the position, we can reuse it if the user passed no + # kwargs or if the Axes class and kwargs are identical. + if (ax.get_subplotspec() == key + and (kwargs == {} + or (ax._projection_init + == fig._process_projection_requirements(**kwargs)))): + break + else: + # we have exhausted the known Axes and none match, make a new one! + ax = fig.add_subplot(*args, **kwargs) + + fig.sca(ax) + + return ax + + +@overload +def subplots( + nrows: Literal[1] = ..., + ncols: Literal[1] = ..., + *, + sharex: bool | Literal["none", "all", "row", "col"] = ..., + sharey: bool | Literal["none", "all", "row", "col"] = ..., + squeeze: Literal[True] = ..., + width_ratios: Sequence[float] | None = ..., + height_ratios: Sequence[float] | None = ..., + subplot_kw: dict[str, Any] | None = ..., + gridspec_kw: dict[str, Any] | None = ..., + **fig_kw +) -> tuple[Figure, Axes]: + ... + + +@overload +def subplots( + nrows: int = ..., + ncols: int = ..., + *, + sharex: bool | Literal["none", "all", "row", "col"] = ..., + sharey: bool | Literal["none", "all", "row", "col"] = ..., + squeeze: Literal[False], + width_ratios: Sequence[float] | None = ..., + height_ratios: Sequence[float] | None = ..., + subplot_kw: dict[str, Any] | None = ..., + gridspec_kw: dict[str, Any] | None = ..., + **fig_kw +) -> tuple[Figure, np.ndarray]: # TODO numpy/numpy#24738 + ... + + +@overload +def subplots( + nrows: int = ..., + ncols: int = ..., + *, + sharex: bool | Literal["none", "all", "row", "col"] = ..., + sharey: bool | Literal["none", "all", "row", "col"] = ..., + squeeze: bool = ..., + width_ratios: Sequence[float] | None = ..., + height_ratios: Sequence[float] | None = ..., + subplot_kw: dict[str, Any] | None = ..., + gridspec_kw: dict[str, Any] | None = ..., + **fig_kw +) -> tuple[Figure, Any]: + ... + + +def subplots( + nrows: int = 1, ncols: int = 1, *, + sharex: bool | Literal["none", "all", "row", "col"] = False, + sharey: bool | Literal["none", "all", "row", "col"] = False, + squeeze: bool = True, + width_ratios: Sequence[float] | None = None, + height_ratios: Sequence[float] | None = None, + subplot_kw: dict[str, Any] | None = None, + gridspec_kw: dict[str, Any] | None = None, + **fig_kw +) -> tuple[Figure, Any]: + """ + Create a figure and a set of subplots. + + This utility wrapper makes it convenient to create common layouts of + subplots, including the enclosing figure object, in a single call. + + Parameters + ---------- + nrows, ncols : int, default: 1 + Number of rows/columns of the subplot grid. + + sharex, sharey : bool or {'none', 'all', 'row', 'col'}, default: False + Controls sharing of properties among x (*sharex*) or y (*sharey*) + axes: + + - True or 'all': x- or y-axis will be shared among all subplots. + - False or 'none': each subplot x- or y-axis will be independent. + - 'row': each subplot row will share an x- or y-axis. + - 'col': each subplot column will share an x- or y-axis. + + When subplots have a shared x-axis along a column, only the x tick + labels of the bottom subplot are created. Similarly, when subplots + have a shared y-axis along a row, only the y tick labels of the first + column subplot are created. To later turn other subplots' ticklabels + on, use `~matplotlib.axes.Axes.tick_params`. + + When subplots have a shared axis that has units, calling + `.Axis.set_units` will update each axis with the new units. + + Note that it is not possible to unshare axes. + + squeeze : bool, default: True + - If True, extra dimensions are squeezed out from the returned + array of `~matplotlib.axes.Axes`: + + - if only one subplot is constructed (nrows=ncols=1), the + resulting single Axes object is returned as a scalar. + - for Nx1 or 1xM subplots, the returned object is a 1D numpy + object array of Axes objects. + - for NxM, subplots with N>1 and M>1 are returned as a 2D array. + + - If False, no squeezing at all is done: the returned Axes object is + always a 2D array containing Axes instances, even if it ends up + being 1x1. + + width_ratios : array-like of length *ncols*, optional + Defines the relative widths of the columns. Each column gets a + relative width of ``width_ratios[i] / sum(width_ratios)``. + If not given, all columns will have the same width. Equivalent + to ``gridspec_kw={'width_ratios': [...]}``. + + height_ratios : array-like of length *nrows*, optional + Defines the relative heights of the rows. Each row gets a + relative height of ``height_ratios[i] / sum(height_ratios)``. + If not given, all rows will have the same height. Convenience + for ``gridspec_kw={'height_ratios': [...]}``. + + subplot_kw : dict, optional + Dict with keywords passed to the + `~matplotlib.figure.Figure.add_subplot` call used to create each + subplot. + + gridspec_kw : dict, optional + Dict with keywords passed to the `~matplotlib.gridspec.GridSpec` + constructor used to create the grid the subplots are placed on. + + **fig_kw + All additional keyword arguments are passed to the + `.pyplot.figure` call. + + Returns + ------- + fig : `.Figure` + + ax : `~matplotlib.axes.Axes` or array of Axes + *ax* can be either a single `~.axes.Axes` object, or an array of Axes + objects if more than one subplot was created. The dimensions of the + resulting array can be controlled with the squeeze keyword, see above. + + Typical idioms for handling the return value are:: + + # using the variable ax for single a Axes + fig, ax = plt.subplots() + + # using the variable axs for multiple Axes + fig, axs = plt.subplots(2, 2) + + # using tuple unpacking for multiple Axes + fig, (ax1, ax2) = plt.subplots(1, 2) + fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2) + + The names ``ax`` and pluralized ``axs`` are preferred over ``axes`` + because for the latter it's not clear if it refers to a single + `~.axes.Axes` instance or a collection of these. + + See Also + -------- + .pyplot.figure + .pyplot.subplot + .pyplot.axes + .Figure.subplots + .Figure.add_subplot + + Examples + -------- + :: + + # First create some toy data: + x = np.linspace(0, 2*np.pi, 400) + y = np.sin(x**2) + + # Create just a figure and only one subplot + fig, ax = plt.subplots() + ax.plot(x, y) + ax.set_title('Simple plot') + + # Create two subplots and unpack the output array immediately + f, (ax1, ax2) = plt.subplots(1, 2, sharey=True) + ax1.plot(x, y) + ax1.set_title('Sharing Y axis') + ax2.scatter(x, y) + + # Create four polar Axes and access them through the returned array + fig, axs = plt.subplots(2, 2, subplot_kw=dict(projection="polar")) + axs[0, 0].plot(x, y) + axs[1, 1].scatter(x, y) + + # Share a X axis with each column of subplots + plt.subplots(2, 2, sharex='col') + + # Share a Y axis with each row of subplots + plt.subplots(2, 2, sharey='row') + + # Share both X and Y axes with all subplots + plt.subplots(2, 2, sharex='all', sharey='all') + + # Note that this is the same as + plt.subplots(2, 2, sharex=True, sharey=True) + + # Create figure number 10 with a single subplot + # and clears it if it already exists. + fig, ax = plt.subplots(num=10, clear=True) + + """ + fig = figure(**fig_kw) + axs = fig.subplots(nrows=nrows, ncols=ncols, sharex=sharex, sharey=sharey, + squeeze=squeeze, subplot_kw=subplot_kw, + gridspec_kw=gridspec_kw, height_ratios=height_ratios, + width_ratios=width_ratios) + return fig, axs + + +@overload +def subplot_mosaic( + mosaic: str, + *, + sharex: bool = ..., + sharey: bool = ..., + width_ratios: ArrayLike | None = ..., + height_ratios: ArrayLike | None = ..., + empty_sentinel: str = ..., + subplot_kw: dict[str, Any] | None = ..., + gridspec_kw: dict[str, Any] | None = ..., + per_subplot_kw: dict[str | tuple[str, ...], dict[str, Any]] | None = ..., + **fig_kw: Any +) -> tuple[Figure, dict[str, matplotlib.axes.Axes]]: ... + + +@overload +def subplot_mosaic( + mosaic: list[HashableList[_T]], + *, + sharex: bool = ..., + sharey: bool = ..., + width_ratios: ArrayLike | None = ..., + height_ratios: ArrayLike | None = ..., + empty_sentinel: _T = ..., + subplot_kw: dict[str, Any] | None = ..., + gridspec_kw: dict[str, Any] | None = ..., + per_subplot_kw: dict[_T | tuple[_T, ...], dict[str, Any]] | None = ..., + **fig_kw: Any +) -> tuple[Figure, dict[_T, matplotlib.axes.Axes]]: ... + + +@overload +def subplot_mosaic( + mosaic: list[HashableList[Hashable]], + *, + sharex: bool = ..., + sharey: bool = ..., + width_ratios: ArrayLike | None = ..., + height_ratios: ArrayLike | None = ..., + empty_sentinel: Any = ..., + subplot_kw: dict[str, Any] | None = ..., + gridspec_kw: dict[str, Any] | None = ..., + per_subplot_kw: dict[Hashable | tuple[Hashable, ...], dict[str, Any]] | None = ..., + **fig_kw: Any +) -> tuple[Figure, dict[Hashable, matplotlib.axes.Axes]]: ... + + +def subplot_mosaic( + mosaic: str | list[HashableList[_T]] | list[HashableList[Hashable]], + *, + sharex: bool = False, + sharey: bool = False, + width_ratios: ArrayLike | None = None, + height_ratios: ArrayLike | None = None, + empty_sentinel: Any = '.', + subplot_kw: dict[str, Any] | None = None, + gridspec_kw: dict[str, Any] | None = None, + per_subplot_kw: dict[str | tuple[str, ...], dict[str, Any]] | + dict[_T | tuple[_T, ...], dict[str, Any]] | + dict[Hashable | tuple[Hashable, ...], dict[str, Any]] | None = None, + **fig_kw: Any +) -> tuple[Figure, dict[str, matplotlib.axes.Axes]] | \ + tuple[Figure, dict[_T, matplotlib.axes.Axes]] | \ + tuple[Figure, dict[Hashable, matplotlib.axes.Axes]]: + """ + Build a layout of Axes based on ASCII art or nested lists. + + This is a helper function to build complex GridSpec layouts visually. + + See :ref:`mosaic` + for an example and full API documentation + + Parameters + ---------- + mosaic : list of list of {hashable or nested} or str + + A visual layout of how you want your Axes to be arranged + labeled as strings. For example :: + + x = [['A panel', 'A panel', 'edge'], + ['C panel', '.', 'edge']] + + produces 4 Axes: + + - 'A panel' which is 1 row high and spans the first two columns + - 'edge' which is 2 rows high and is on the right edge + - 'C panel' which in 1 row and 1 column wide in the bottom left + - a blank space 1 row and 1 column wide in the bottom center + + Any of the entries in the layout can be a list of lists + of the same form to create nested layouts. + + If input is a str, then it must be of the form :: + + ''' + AAE + C.E + ''' + + where each character is a column and each line is a row. + This only allows only single character Axes labels and does + not allow nesting but is very terse. + + sharex, sharey : bool, default: False + If True, the x-axis (*sharex*) or y-axis (*sharey*) will be shared + among all subplots. In that case, tick label visibility and axis units + behave as for `subplots`. If False, each subplot's x- or y-axis will + be independent. + + width_ratios : array-like of length *ncols*, optional + Defines the relative widths of the columns. Each column gets a + relative width of ``width_ratios[i] / sum(width_ratios)``. + If not given, all columns will have the same width. Convenience + for ``gridspec_kw={'width_ratios': [...]}``. + + height_ratios : array-like of length *nrows*, optional + Defines the relative heights of the rows. Each row gets a + relative height of ``height_ratios[i] / sum(height_ratios)``. + If not given, all rows will have the same height. Convenience + for ``gridspec_kw={'height_ratios': [...]}``. + + empty_sentinel : object, optional + Entry in the layout to mean "leave this space empty". Defaults + to ``'.'``. Note, if *layout* is a string, it is processed via + `inspect.cleandoc` to remove leading white space, which may + interfere with using white-space as the empty sentinel. + + subplot_kw : dict, optional + Dictionary with keywords passed to the `.Figure.add_subplot` call + used to create each subplot. These values may be overridden by + values in *per_subplot_kw*. + + per_subplot_kw : dict, optional + A dictionary mapping the Axes identifiers or tuples of identifiers + to a dictionary of keyword arguments to be passed to the + `.Figure.add_subplot` call used to create each subplot. The values + in these dictionaries have precedence over the values in + *subplot_kw*. + + If *mosaic* is a string, and thus all keys are single characters, + it is possible to use a single string instead of a tuple as keys; + i.e. ``"AB"`` is equivalent to ``("A", "B")``. + + .. versionadded:: 3.7 + + gridspec_kw : dict, optional + Dictionary with keywords passed to the `.GridSpec` constructor used + to create the grid the subplots are placed on. + + **fig_kw + All additional keyword arguments are passed to the + `.pyplot.figure` call. + + Returns + ------- + fig : `.Figure` + The new figure + + dict[label, Axes] + A dictionary mapping the labels to the Axes objects. The order of + the Axes is left-to-right and top-to-bottom of their position in the + total layout. + + """ + fig = figure(**fig_kw) + ax_dict = fig.subplot_mosaic( # type: ignore[misc] + mosaic, # type: ignore[arg-type] + sharex=sharex, sharey=sharey, + height_ratios=height_ratios, width_ratios=width_ratios, + subplot_kw=subplot_kw, gridspec_kw=gridspec_kw, + empty_sentinel=empty_sentinel, + per_subplot_kw=per_subplot_kw, # type: ignore[arg-type] + ) + return fig, ax_dict + + +def subplot2grid( + shape: tuple[int, int], loc: tuple[int, int], + rowspan: int = 1, colspan: int = 1, + fig: Figure | None = None, + **kwargs +) -> matplotlib.axes.Axes: + """ + Create a subplot at a specific location inside a regular grid. + + Parameters + ---------- + shape : (int, int) + Number of rows and of columns of the grid in which to place axis. + loc : (int, int) + Row number and column number of the axis location within the grid. + rowspan : int, default: 1 + Number of rows for the axis to span downwards. + colspan : int, default: 1 + Number of columns for the axis to span to the right. + fig : `.Figure`, optional + Figure to place the subplot in. Defaults to the current figure. + **kwargs + Additional keyword arguments are handed to `~.Figure.add_subplot`. + + Returns + ------- + `~.axes.Axes` + + The Axes of the subplot. The returned Axes can actually be an instance + of a subclass, such as `.projections.polar.PolarAxes` for polar + projections. + + Notes + ----- + The following call :: + + ax = subplot2grid((nrows, ncols), (row, col), rowspan, colspan) + + is identical to :: + + fig = gcf() + gs = fig.add_gridspec(nrows, ncols) + ax = fig.add_subplot(gs[row:row+rowspan, col:col+colspan]) + """ + if fig is None: + fig = gcf() + rows, cols = shape + gs = GridSpec._check_gridspec_exists(fig, rows, cols) + subplotspec = gs.new_subplotspec(loc, rowspan=rowspan, colspan=colspan) + return fig.add_subplot(subplotspec, **kwargs) + + +def twinx(ax: matplotlib.axes.Axes | None = None) -> _AxesBase: + """ + Make and return a second Axes that shares the *x*-axis. The new Axes will + overlay *ax* (or the current Axes if *ax* is *None*), and its ticks will be + on the right. + + Examples + -------- + :doc:`/gallery/subplots_axes_and_figures/two_scales` + """ + if ax is None: + ax = gca() + ax1 = ax.twinx() + return ax1 + + +def twiny(ax: matplotlib.axes.Axes | None = None) -> _AxesBase: + """ + Make and return a second Axes that shares the *y*-axis. The new Axes will + overlay *ax* (or the current Axes if *ax* is *None*), and its ticks will be + on the top. + + Examples + -------- + :doc:`/gallery/subplots_axes_and_figures/two_scales` + """ + if ax is None: + ax = gca() + ax1 = ax.twiny() + return ax1 + + +def subplot_tool(targetfig: Figure | None = None) -> SubplotTool | None: + """ + Launch a subplot tool window for a figure. + + Returns + ------- + `matplotlib.widgets.SubplotTool` + """ + if targetfig is None: + targetfig = gcf() + tb = targetfig.canvas.manager.toolbar # type: ignore[union-attr] + if hasattr(tb, "configure_subplots"): # toolbar2 + from matplotlib.backend_bases import NavigationToolbar2 + return cast(NavigationToolbar2, tb).configure_subplots() + elif hasattr(tb, "trigger_tool"): # toolmanager + from matplotlib.backend_bases import ToolContainerBase + cast(ToolContainerBase, tb).trigger_tool("subplots") + return None + else: + raise ValueError("subplot_tool can only be launched for figures with " + "an associated toolbar") + + +def box(on: bool | None = None) -> None: + """ + Turn the Axes box on or off on the current Axes. + + Parameters + ---------- + on : bool or None + The new `~matplotlib.axes.Axes` box state. If ``None``, toggle + the state. + + See Also + -------- + :meth:`matplotlib.axes.Axes.set_frame_on` + :meth:`matplotlib.axes.Axes.get_frame_on` + """ + ax = gca() + if on is None: + on = not ax.get_frame_on() + ax.set_frame_on(on) + +## Axis ## + + +def xlim(*args, **kwargs) -> tuple[float, float]: + """ + Get or set the x limits of the current Axes. + + Call signatures:: + + left, right = xlim() # return the current xlim + xlim((left, right)) # set the xlim to left, right + xlim(left, right) # set the xlim to left, right + + If you do not specify args, you can pass *left* or *right* as kwargs, + i.e.:: + + xlim(right=3) # adjust the right leaving left unchanged + xlim(left=1) # adjust the left leaving right unchanged + + Setting limits turns autoscaling off for the x-axis. + + Returns + ------- + left, right + A tuple of the new x-axis limits. + + Notes + ----- + Calling this function with no arguments (e.g. ``xlim()``) is the pyplot + equivalent of calling `~.Axes.get_xlim` on the current Axes. + Calling this function with arguments is the pyplot equivalent of calling + `~.Axes.set_xlim` on the current Axes. All arguments are passed though. + """ + ax = gca() + if not args and not kwargs: + return ax.get_xlim() + ret = ax.set_xlim(*args, **kwargs) + return ret + + +def ylim(*args, **kwargs) -> tuple[float, float]: + """ + Get or set the y-limits of the current Axes. + + Call signatures:: + + bottom, top = ylim() # return the current ylim + ylim((bottom, top)) # set the ylim to bottom, top + ylim(bottom, top) # set the ylim to bottom, top + + If you do not specify args, you can alternatively pass *bottom* or + *top* as kwargs, i.e.:: + + ylim(top=3) # adjust the top leaving bottom unchanged + ylim(bottom=1) # adjust the bottom leaving top unchanged + + Setting limits turns autoscaling off for the y-axis. + + Returns + ------- + bottom, top + A tuple of the new y-axis limits. + + Notes + ----- + Calling this function with no arguments (e.g. ``ylim()``) is the pyplot + equivalent of calling `~.Axes.get_ylim` on the current Axes. + Calling this function with arguments is the pyplot equivalent of calling + `~.Axes.set_ylim` on the current Axes. All arguments are passed though. + """ + ax = gca() + if not args and not kwargs: + return ax.get_ylim() + ret = ax.set_ylim(*args, **kwargs) + return ret + + +def xticks( + ticks: ArrayLike | None = None, + labels: Sequence[str] | None = None, + *, + minor: bool = False, + **kwargs +) -> tuple[list[Tick] | np.ndarray, list[Text]]: + """ + Get or set the current tick locations and labels of the x-axis. + + Pass no arguments to return the current values without modifying them. + + Parameters + ---------- + ticks : array-like, optional + The list of xtick locations. Passing an empty list removes all xticks. + labels : array-like, optional + The labels to place at the given *ticks* locations. This argument can + only be passed if *ticks* is passed as well. + minor : bool, default: False + If ``False``, get/set the major ticks/labels; if ``True``, the minor + ticks/labels. + **kwargs + `.Text` properties can be used to control the appearance of the labels. + + .. warning:: + + This only sets the properties of the current ticks, which is + only sufficient if you either pass *ticks*, resulting in a + fixed list of ticks, or if the plot is static. + + Ticks are not guaranteed to be persistent. Various operations + can create, delete and modify the Tick instances. There is an + imminent risk that these settings can get lost if you work on + the figure further (including also panning/zooming on a + displayed figure). + + Use `~.pyplot.tick_params` instead if possible. + + + Returns + ------- + locs + The list of xtick locations. + labels + The list of xlabel `.Text` objects. + + Notes + ----- + Calling this function with no arguments (e.g. ``xticks()``) is the pyplot + equivalent of calling `~.Axes.get_xticks` and `~.Axes.get_xticklabels` on + the current Axes. + Calling this function with arguments is the pyplot equivalent of calling + `~.Axes.set_xticks` and `~.Axes.set_xticklabels` on the current Axes. + + Examples + -------- + >>> locs, labels = xticks() # Get the current locations and labels. + >>> xticks(np.arange(0, 1, step=0.2)) # Set label locations. + >>> xticks(np.arange(3), ['Tom', 'Dick', 'Sue']) # Set text labels. + >>> xticks([0, 1, 2], ['January', 'February', 'March'], + ... rotation=20) # Set text labels and properties. + >>> xticks([]) # Disable xticks. + """ + ax = gca() + + locs: list[Tick] | np.ndarray + if ticks is None: + locs = ax.get_xticks(minor=minor) + if labels is not None: + raise TypeError("xticks(): Parameter 'labels' can't be set " + "without setting 'ticks'") + else: + locs = ax.set_xticks(ticks, minor=minor) + + labels_out: list[Text] = [] + if labels is None: + labels_out = ax.get_xticklabels(minor=minor) + for l in labels_out: + l._internal_update(kwargs) + else: + labels_out = ax.set_xticklabels(labels, minor=minor, **kwargs) + + return locs, labels_out + + +def yticks( + ticks: ArrayLike | None = None, + labels: Sequence[str] | None = None, + *, + minor: bool = False, + **kwargs +) -> tuple[list[Tick] | np.ndarray, list[Text]]: + """ + Get or set the current tick locations and labels of the y-axis. + + Pass no arguments to return the current values without modifying them. + + Parameters + ---------- + ticks : array-like, optional + The list of ytick locations. Passing an empty list removes all yticks. + labels : array-like, optional + The labels to place at the given *ticks* locations. This argument can + only be passed if *ticks* is passed as well. + minor : bool, default: False + If ``False``, get/set the major ticks/labels; if ``True``, the minor + ticks/labels. + **kwargs + `.Text` properties can be used to control the appearance of the labels. + + .. warning:: + + This only sets the properties of the current ticks, which is + only sufficient if you either pass *ticks*, resulting in a + fixed list of ticks, or if the plot is static. + + Ticks are not guaranteed to be persistent. Various operations + can create, delete and modify the Tick instances. There is an + imminent risk that these settings can get lost if you work on + the figure further (including also panning/zooming on a + displayed figure). + + Use `~.pyplot.tick_params` instead if possible. + + Returns + ------- + locs + The list of ytick locations. + labels + The list of ylabel `.Text` objects. + + Notes + ----- + Calling this function with no arguments (e.g. ``yticks()``) is the pyplot + equivalent of calling `~.Axes.get_yticks` and `~.Axes.get_yticklabels` on + the current Axes. + Calling this function with arguments is the pyplot equivalent of calling + `~.Axes.set_yticks` and `~.Axes.set_yticklabels` on the current Axes. + + Examples + -------- + >>> locs, labels = yticks() # Get the current locations and labels. + >>> yticks(np.arange(0, 1, step=0.2)) # Set label locations. + >>> yticks(np.arange(3), ['Tom', 'Dick', 'Sue']) # Set text labels. + >>> yticks([0, 1, 2], ['January', 'February', 'March'], + ... rotation=45) # Set text labels and properties. + >>> yticks([]) # Disable yticks. + """ + ax = gca() + + locs: list[Tick] | np.ndarray + if ticks is None: + locs = ax.get_yticks(minor=minor) + if labels is not None: + raise TypeError("yticks(): Parameter 'labels' can't be set " + "without setting 'ticks'") + else: + locs = ax.set_yticks(ticks, minor=minor) + + labels_out: list[Text] = [] + if labels is None: + labels_out = ax.get_yticklabels(minor=minor) + for l in labels_out: + l._internal_update(kwargs) + else: + labels_out = ax.set_yticklabels(labels, minor=minor, **kwargs) + + return locs, labels_out + + +def rgrids( + radii: ArrayLike | None = None, + labels: Sequence[str | Text] | None = None, + angle: float | None = None, + fmt: str | None = None, + **kwargs +) -> tuple[list[Line2D], list[Text]]: + """ + Get or set the radial gridlines on the current polar plot. + + Call signatures:: + + lines, labels = rgrids() + lines, labels = rgrids(radii, labels=None, angle=22.5, fmt=None, **kwargs) + + When called with no arguments, `.rgrids` simply returns the tuple + (*lines*, *labels*). When called with arguments, the labels will + appear at the specified radial distances and angle. + + Parameters + ---------- + radii : tuple with floats + The radii for the radial gridlines + + labels : tuple with strings or None + The labels to use at each radial gridline. The + `matplotlib.ticker.ScalarFormatter` will be used if None. + + angle : float + The angular position of the radius labels in degrees. + + fmt : str or None + Format string used in `matplotlib.ticker.FormatStrFormatter`. + For example '%f'. + + Returns + ------- + lines : list of `.lines.Line2D` + The radial gridlines. + + labels : list of `.text.Text` + The tick labels. + + Other Parameters + ---------------- + **kwargs + *kwargs* are optional `.Text` properties for the labels. + + See Also + -------- + .pyplot.thetagrids + .projections.polar.PolarAxes.set_rgrids + .Axis.get_gridlines + .Axis.get_ticklabels + + Examples + -------- + :: + + # set the locations of the radial gridlines + lines, labels = rgrids( (0.25, 0.5, 1.0) ) + + # set the locations and labels of the radial gridlines + lines, labels = rgrids( (0.25, 0.5, 1.0), ('Tom', 'Dick', 'Harry' )) + """ + ax = gca() + if not isinstance(ax, PolarAxes): + raise RuntimeError('rgrids only defined for polar Axes') + if all(p is None for p in [radii, labels, angle, fmt]) and not kwargs: + lines_out: list[Line2D] = ax.yaxis.get_gridlines() + labels_out: list[Text] = ax.yaxis.get_ticklabels() + elif radii is None: + raise TypeError("'radii' cannot be None when other parameters are passed") + else: + lines_out, labels_out = ax.set_rgrids( + radii, labels=labels, angle=angle, fmt=fmt, **kwargs) + return lines_out, labels_out + + +def thetagrids( + angles: ArrayLike | None = None, + labels: Sequence[str | Text] | None = None, + fmt: str | None = None, + **kwargs +) -> tuple[list[Line2D], list[Text]]: + """ + Get or set the theta gridlines on the current polar plot. + + Call signatures:: + + lines, labels = thetagrids() + lines, labels = thetagrids(angles, labels=None, fmt=None, **kwargs) + + When called with no arguments, `.thetagrids` simply returns the tuple + (*lines*, *labels*). When called with arguments, the labels will + appear at the specified angles. + + Parameters + ---------- + angles : tuple with floats, degrees + The angles of the theta gridlines. + + labels : tuple with strings or None + The labels to use at each radial gridline. The + `.projections.polar.ThetaFormatter` will be used if None. + + fmt : str or None + Format string used in `matplotlib.ticker.FormatStrFormatter`. + For example '%f'. Note that the angle in radians will be used. + + Returns + ------- + lines : list of `.lines.Line2D` + The theta gridlines. + + labels : list of `.text.Text` + The tick labels. + + Other Parameters + ---------------- + **kwargs + *kwargs* are optional `.Text` properties for the labels. + + See Also + -------- + .pyplot.rgrids + .projections.polar.PolarAxes.set_thetagrids + .Axis.get_gridlines + .Axis.get_ticklabels + + Examples + -------- + :: + + # set the locations of the angular gridlines + lines, labels = thetagrids(range(45, 360, 90)) + + # set the locations and labels of the angular gridlines + lines, labels = thetagrids(range(45, 360, 90), ('NE', 'NW', 'SW', 'SE')) + """ + ax = gca() + if not isinstance(ax, PolarAxes): + raise RuntimeError('thetagrids only defined for polar Axes') + if all(param is None for param in [angles, labels, fmt]) and not kwargs: + lines_out: list[Line2D] = ax.xaxis.get_ticklines() + labels_out: list[Text] = ax.xaxis.get_ticklabels() + elif angles is None: + raise TypeError("'angles' cannot be None when other parameters are passed") + else: + lines_out, labels_out = ax.set_thetagrids(angles, + labels=labels, fmt=fmt, + **kwargs) + return lines_out, labels_out + + +@_api.deprecated("3.7", pending=True) +def get_plot_commands() -> list[str]: + """ + Get a sorted list of all of the plotting commands. + """ + NON_PLOT_COMMANDS = { + 'connect', 'disconnect', 'get_current_fig_manager', 'ginput', + 'new_figure_manager', 'waitforbuttonpress'} + return [name for name in _get_pyplot_commands() + if name not in NON_PLOT_COMMANDS] + + +def _get_pyplot_commands() -> list[str]: + # This works by searching for all functions in this module and removing + # a few hard-coded exclusions, as well as all of the colormap-setting + # functions, and anything marked as private with a preceding underscore. + exclude = {'colormaps', 'colors', 'get_plot_commands', *colormaps} + this_module = inspect.getmodule(get_plot_commands) + return sorted( + name for name, obj in globals().items() + if not name.startswith('_') and name not in exclude + and inspect.isfunction(obj) + and inspect.getmodule(obj) is this_module) + + +## Plotting part 1: manually generated functions and wrappers ## + + +@_copy_docstring_and_deprecators(Figure.colorbar) +def colorbar( + mappable: ScalarMappable | ColorizingArtist | None = None, + cax: matplotlib.axes.Axes | None = None, + ax: matplotlib.axes.Axes | Iterable[matplotlib.axes.Axes] | None = None, + **kwargs +) -> Colorbar: + if mappable is None: + mappable = gci() + if mappable is None: + raise RuntimeError('No mappable was found to use for colorbar ' + 'creation. First define a mappable such as ' + 'an image (with imshow) or a contour set (' + 'with contourf).') + ret = gcf().colorbar(mappable, cax=cax, ax=ax, **kwargs) + return ret + + +def clim(vmin: float | None = None, vmax: float | None = None) -> None: + """ + Set the color limits of the current image. + + If either *vmin* or *vmax* is None, the image min/max respectively + will be used for color scaling. + + If you want to set the clim of multiple images, use + `~.ScalarMappable.set_clim` on every image, for example:: + + for im in gca().get_images(): + im.set_clim(0, 0.5) + + """ + im = gci() + if im is None: + raise RuntimeError('You must first define an image, e.g., with imshow') + + im.set_clim(vmin, vmax) + + +def get_cmap(name: Colormap | str | None = None, lut: int | None = None) -> Colormap: + """ + Get a colormap instance, defaulting to rc values if *name* is None. + + Parameters + ---------- + name : `~matplotlib.colors.Colormap` or str or None, default: None + If a `.Colormap` instance, it will be returned. Otherwise, the name of + a colormap known to Matplotlib, which will be resampled by *lut*. The + default, None, means :rc:`image.cmap`. + lut : int or None, default: None + If *name* is not already a Colormap instance and *lut* is not None, the + colormap will be resampled to have *lut* entries in the lookup table. + + Returns + ------- + Colormap + """ + if name is None: + name = rcParams['image.cmap'] + if isinstance(name, Colormap): + return name + _api.check_in_list(sorted(_colormaps), name=name) + if lut is None: + return _colormaps[name] + else: + return _colormaps[name].resampled(lut) + + +def set_cmap(cmap: Colormap | str) -> None: + """ + Set the default colormap, and applies it to the current image if any. + + Parameters + ---------- + cmap : `~matplotlib.colors.Colormap` or str + A colormap instance or the name of a registered colormap. + + See Also + -------- + colormaps + get_cmap + """ + cmap = get_cmap(cmap) + + rc('image', cmap=cmap.name) + im = gci() + + if im is not None: + im.set_cmap(cmap) + + +@_copy_docstring_and_deprecators(matplotlib.image.imread) +def imread( + fname: str | pathlib.Path | BinaryIO, format: str | None = None +) -> np.ndarray: + return matplotlib.image.imread(fname, format) + + +@_copy_docstring_and_deprecators(matplotlib.image.imsave) +def imsave( + fname: str | os.PathLike | BinaryIO, arr: ArrayLike, **kwargs +) -> None: + matplotlib.image.imsave(fname, arr, **kwargs) + + +def matshow(A: ArrayLike, fignum: None | int = None, **kwargs) -> AxesImage: + """ + Display a 2D array as a matrix in a new figure window. + + The origin is set at the upper left hand corner. + The indexing is ``(row, column)`` so that the first index runs vertically + and the second index runs horizontally in the figure: + + .. code-block:: none + + A[0, 0] ⋯ A[0, M-1] + ⋮ ⋮ + A[N-1, 0] ⋯ A[N-1, M-1] + + The aspect ratio of the figure window is that of the array, + unless this would make an excessively short or narrow figure. + + Tick labels for the xaxis are placed on top. + + Parameters + ---------- + A : 2D array-like + The matrix to be displayed. + + fignum : None or int + If *None*, create a new, appropriately sized figure window. + + If 0, use the current Axes (creating one if there is none, without ever + adjusting the figure size). + + Otherwise, create a new Axes on the figure with the given number + (creating it at the appropriate size if it does not exist, but not + adjusting the figure size otherwise). Note that this will be drawn on + top of any preexisting Axes on the figure. + + Returns + ------- + `~matplotlib.image.AxesImage` + + Other Parameters + ---------------- + **kwargs : `~matplotlib.axes.Axes.imshow` arguments + + """ + A = np.asanyarray(A) + if fignum == 0: + ax = gca() + else: + if fignum is not None and fignum_exists(fignum): + # Do not try to set a figure size. + figsize = None + else: + # Extract actual aspect ratio of array and make appropriately sized figure. + figsize = figaspect(A) + fig = figure(fignum, figsize=figsize) + ax = fig.add_axes((0.15, 0.09, 0.775, 0.775)) + im = ax.matshow(A, **kwargs) + sci(im) + return im + + +def polar(*args, **kwargs) -> list[Line2D]: + """ + Make a polar plot. + + call signature:: + + polar(theta, r, [fmt], **kwargs) + + This is a convenience wrapper around `.pyplot.plot`. It ensures that the + current Axes is polar (or creates one if needed) and then passes all parameters + to ``.pyplot.plot``. + + .. note:: + When making polar plots using the :ref:`pyplot API `, + ``polar()`` should typically be the first command because that makes sure + a polar Axes is created. Using other commands such as ``plt.title()`` + before this can lead to the implicit creation of a rectangular Axes, in which + case a subsequent ``polar()`` call will fail. + """ + # If an axis already exists, check if it has a polar projection + if gcf().get_axes(): + ax = gca() + if not isinstance(ax, PolarAxes): + _api.warn_deprecated( + "3.10", + message="There exists a non-polar current Axes. Therefore, the " + "resulting plot from 'polar()' is non-polar. You likely " + "should call 'polar()' before any other pyplot plotting " + "commands. " + "Support for this scenario is deprecated in %(since)s and " + "will raise an error in %(removal)s" + ) + else: + ax = axes(projection="polar") + return ax.plot(*args, **kwargs) + + +# If rcParams['backend_fallback'] is true, and an interactive backend is +# requested, ignore rcParams['backend'] and force selection of a backend that +# is compatible with the current running interactive framework. +if rcParams["backend_fallback"]: + requested_backend = rcParams._get_backend_or_none() # type: ignore[attr-defined] + requested_backend = None if requested_backend is None else requested_backend.lower() + available_backends = backend_registry.list_builtin(BackendFilter.INTERACTIVE) + if ( + requested_backend in (set(available_backends) - {'webagg', 'nbagg'}) + and cbook._get_running_interactive_framework() + ): + rcParams._set("backend", rcsetup._auto_backend_sentinel) + +# fmt: on + +################# REMAINING CONTENT GENERATED BY boilerplate.py ############## + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Figure.figimage) +def figimage( + X: ArrayLike, + xo: int = 0, + yo: int = 0, + alpha: float | None = None, + norm: str | Normalize | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + origin: Literal["upper", "lower"] | None = None, + resize: bool = False, + *, + colorizer: Colorizer | None = None, + **kwargs, +) -> FigureImage: + return gcf().figimage( + X, + xo=xo, + yo=yo, + alpha=alpha, + norm=norm, + cmap=cmap, + vmin=vmin, + vmax=vmax, + origin=origin, + resize=resize, + colorizer=colorizer, + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Figure.text) +def figtext( + x: float, y: float, s: str, fontdict: dict[str, Any] | None = None, **kwargs +) -> Text: + return gcf().text(x, y, s, fontdict=fontdict, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Figure.gca) +def gca() -> Axes: + return gcf().gca() + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Figure._gci) +def gci() -> ColorizingArtist | None: + return gcf()._gci() + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Figure.ginput) +def ginput( + n: int = 1, + timeout: float = 30, + show_clicks: bool = True, + mouse_add: MouseButton = MouseButton.LEFT, + mouse_pop: MouseButton = MouseButton.RIGHT, + mouse_stop: MouseButton = MouseButton.MIDDLE, +) -> list[tuple[int, int]]: + return gcf().ginput( + n=n, + timeout=timeout, + show_clicks=show_clicks, + mouse_add=mouse_add, + mouse_pop=mouse_pop, + mouse_stop=mouse_stop, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Figure.subplots_adjust) +def subplots_adjust( + left: float | None = None, + bottom: float | None = None, + right: float | None = None, + top: float | None = None, + wspace: float | None = None, + hspace: float | None = None, +) -> None: + gcf().subplots_adjust( + left=left, bottom=bottom, right=right, top=top, wspace=wspace, hspace=hspace + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Figure.suptitle) +def suptitle(t: str, **kwargs) -> Text: + return gcf().suptitle(t, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Figure.tight_layout) +def tight_layout( + *, + pad: float = 1.08, + h_pad: float | None = None, + w_pad: float | None = None, + rect: tuple[float, float, float, float] | None = None, +) -> None: + gcf().tight_layout(pad=pad, h_pad=h_pad, w_pad=w_pad, rect=rect) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Figure.waitforbuttonpress) +def waitforbuttonpress(timeout: float = -1) -> None | bool: + return gcf().waitforbuttonpress(timeout=timeout) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.acorr) +def acorr( + x: ArrayLike, *, data=None, **kwargs +) -> tuple[np.ndarray, np.ndarray, LineCollection | Line2D, Line2D | None]: + return gca().acorr(x, **({"data": data} if data is not None else {}), **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.angle_spectrum) +def angle_spectrum( + x: ArrayLike, + Fs: float | None = None, + Fc: int | None = None, + window: Callable[[ArrayLike], ArrayLike] | ArrayLike | None = None, + pad_to: int | None = None, + sides: Literal["default", "onesided", "twosided"] | None = None, + *, + data=None, + **kwargs, +) -> tuple[np.ndarray, np.ndarray, Line2D]: + return gca().angle_spectrum( + x, + Fs=Fs, + Fc=Fc, + window=window, + pad_to=pad_to, + sides=sides, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.annotate) +def annotate( + text: str, + xy: tuple[float, float], + xytext: tuple[float, float] | None = None, + xycoords: CoordsType = "data", + textcoords: CoordsType | None = None, + arrowprops: dict[str, Any] | None = None, + annotation_clip: bool | None = None, + **kwargs, +) -> Annotation: + return gca().annotate( + text, + xy, + xytext=xytext, + xycoords=xycoords, + textcoords=textcoords, + arrowprops=arrowprops, + annotation_clip=annotation_clip, + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.arrow) +def arrow(x: float, y: float, dx: float, dy: float, **kwargs) -> FancyArrow: + return gca().arrow(x, y, dx, dy, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.autoscale) +def autoscale( + enable: bool = True, + axis: Literal["both", "x", "y"] = "both", + tight: bool | None = None, +) -> None: + gca().autoscale(enable=enable, axis=axis, tight=tight) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.axhline) +def axhline(y: float = 0, xmin: float = 0, xmax: float = 1, **kwargs) -> Line2D: + return gca().axhline(y=y, xmin=xmin, xmax=xmax, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.axhspan) +def axhspan( + ymin: float, ymax: float, xmin: float = 0, xmax: float = 1, **kwargs +) -> Rectangle: + return gca().axhspan(ymin, ymax, xmin=xmin, xmax=xmax, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.axis) +def axis( + arg: tuple[float, float, float, float] | bool | str | None = None, + /, + *, + emit: bool = True, + **kwargs, +) -> tuple[float, float, float, float]: + return gca().axis(arg, emit=emit, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.axline) +def axline( + xy1: tuple[float, float], + xy2: tuple[float, float] | None = None, + *, + slope: float | None = None, + **kwargs, +) -> AxLine: + return gca().axline(xy1, xy2=xy2, slope=slope, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.axvline) +def axvline(x: float = 0, ymin: float = 0, ymax: float = 1, **kwargs) -> Line2D: + return gca().axvline(x=x, ymin=ymin, ymax=ymax, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.axvspan) +def axvspan( + xmin: float, xmax: float, ymin: float = 0, ymax: float = 1, **kwargs +) -> Rectangle: + return gca().axvspan(xmin, xmax, ymin=ymin, ymax=ymax, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.bar) +def bar( + x: float | ArrayLike, + height: float | ArrayLike, + width: float | ArrayLike = 0.8, + bottom: float | ArrayLike | None = None, + *, + align: Literal["center", "edge"] = "center", + data=None, + **kwargs, +) -> BarContainer: + return gca().bar( + x, + height, + width=width, + bottom=bottom, + align=align, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.barbs) +def barbs(*args, data=None, **kwargs) -> Barbs: + return gca().barbs(*args, **({"data": data} if data is not None else {}), **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.barh) +def barh( + y: float | ArrayLike, + width: float | ArrayLike, + height: float | ArrayLike = 0.8, + left: float | ArrayLike | None = None, + *, + align: Literal["center", "edge"] = "center", + data=None, + **kwargs, +) -> BarContainer: + return gca().barh( + y, + width, + height=height, + left=left, + align=align, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.bar_label) +def bar_label( + container: BarContainer, + labels: ArrayLike | None = None, + *, + fmt: str | Callable[[float], str] = "%g", + label_type: Literal["center", "edge"] = "edge", + padding: float = 0, + **kwargs, +) -> list[Annotation]: + return gca().bar_label( + container, + labels=labels, + fmt=fmt, + label_type=label_type, + padding=padding, + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.boxplot) +def boxplot( + x: ArrayLike | Sequence[ArrayLike], + notch: bool | None = None, + sym: str | None = None, + vert: bool | None = None, + orientation: Literal["vertical", "horizontal"] = "vertical", + whis: float | tuple[float, float] | None = None, + positions: ArrayLike | None = None, + widths: float | ArrayLike | None = None, + patch_artist: bool | None = None, + bootstrap: int | None = None, + usermedians: ArrayLike | None = None, + conf_intervals: ArrayLike | None = None, + meanline: bool | None = None, + showmeans: bool | None = None, + showcaps: bool | None = None, + showbox: bool | None = None, + showfliers: bool | None = None, + boxprops: dict[str, Any] | None = None, + tick_labels: Sequence[str] | None = None, + flierprops: dict[str, Any] | None = None, + medianprops: dict[str, Any] | None = None, + meanprops: dict[str, Any] | None = None, + capprops: dict[str, Any] | None = None, + whiskerprops: dict[str, Any] | None = None, + manage_ticks: bool = True, + autorange: bool = False, + zorder: float | None = None, + capwidths: float | ArrayLike | None = None, + label: Sequence[str] | None = None, + *, + data=None, +) -> dict[str, Any]: + return gca().boxplot( + x, + notch=notch, + sym=sym, + vert=vert, + orientation=orientation, + whis=whis, + positions=positions, + widths=widths, + patch_artist=patch_artist, + bootstrap=bootstrap, + usermedians=usermedians, + conf_intervals=conf_intervals, + meanline=meanline, + showmeans=showmeans, + showcaps=showcaps, + showbox=showbox, + showfliers=showfliers, + boxprops=boxprops, + tick_labels=tick_labels, + flierprops=flierprops, + medianprops=medianprops, + meanprops=meanprops, + capprops=capprops, + whiskerprops=whiskerprops, + manage_ticks=manage_ticks, + autorange=autorange, + zorder=zorder, + capwidths=capwidths, + label=label, + **({"data": data} if data is not None else {}), + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.broken_barh) +def broken_barh( + xranges: Sequence[tuple[float, float]], + yrange: tuple[float, float], + *, + data=None, + **kwargs, +) -> PolyCollection: + return gca().broken_barh( + xranges, yrange, **({"data": data} if data is not None else {}), **kwargs + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.clabel) +def clabel(CS: ContourSet, levels: ArrayLike | None = None, **kwargs) -> list[Text]: + return gca().clabel(CS, levels=levels, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.cohere) +def cohere( + x: ArrayLike, + y: ArrayLike, + NFFT: int = 256, + Fs: float = 2, + Fc: int = 0, + detrend: Literal["none", "mean", "linear"] + | Callable[[ArrayLike], ArrayLike] = mlab.detrend_none, + window: Callable[[ArrayLike], ArrayLike] | ArrayLike = mlab.window_hanning, + noverlap: int = 0, + pad_to: int | None = None, + sides: Literal["default", "onesided", "twosided"] = "default", + scale_by_freq: bool | None = None, + *, + data=None, + **kwargs, +) -> tuple[np.ndarray, np.ndarray]: + return gca().cohere( + x, + y, + NFFT=NFFT, + Fs=Fs, + Fc=Fc, + detrend=detrend, + window=window, + noverlap=noverlap, + pad_to=pad_to, + sides=sides, + scale_by_freq=scale_by_freq, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.contour) +def contour(*args, data=None, **kwargs) -> QuadContourSet: + __ret = gca().contour( + *args, **({"data": data} if data is not None else {}), **kwargs + ) + if __ret._A is not None: # type: ignore[attr-defined] + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.contourf) +def contourf(*args, data=None, **kwargs) -> QuadContourSet: + __ret = gca().contourf( + *args, **({"data": data} if data is not None else {}), **kwargs + ) + if __ret._A is not None: # type: ignore[attr-defined] + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.csd) +def csd( + x: ArrayLike, + y: ArrayLike, + NFFT: int | None = None, + Fs: float | None = None, + Fc: int | None = None, + detrend: Literal["none", "mean", "linear"] + | Callable[[ArrayLike], ArrayLike] + | None = None, + window: Callable[[ArrayLike], ArrayLike] | ArrayLike | None = None, + noverlap: int | None = None, + pad_to: int | None = None, + sides: Literal["default", "onesided", "twosided"] | None = None, + scale_by_freq: bool | None = None, + return_line: bool | None = None, + *, + data=None, + **kwargs, +) -> tuple[np.ndarray, np.ndarray] | tuple[np.ndarray, np.ndarray, Line2D]: + return gca().csd( + x, + y, + NFFT=NFFT, + Fs=Fs, + Fc=Fc, + detrend=detrend, + window=window, + noverlap=noverlap, + pad_to=pad_to, + sides=sides, + scale_by_freq=scale_by_freq, + return_line=return_line, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.ecdf) +def ecdf( + x: ArrayLike, + weights: ArrayLike | None = None, + *, + complementary: bool = False, + orientation: Literal["vertical", "horizontal"] = "vertical", + compress: bool = False, + data=None, + **kwargs, +) -> Line2D: + return gca().ecdf( + x, + weights=weights, + complementary=complementary, + orientation=orientation, + compress=compress, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.errorbar) +def errorbar( + x: float | ArrayLike, + y: float | ArrayLike, + yerr: float | ArrayLike | None = None, + xerr: float | ArrayLike | None = None, + fmt: str = "", + ecolor: ColorType | None = None, + elinewidth: float | None = None, + capsize: float | None = None, + barsabove: bool = False, + lolims: bool | ArrayLike = False, + uplims: bool | ArrayLike = False, + xlolims: bool | ArrayLike = False, + xuplims: bool | ArrayLike = False, + errorevery: int | tuple[int, int] = 1, + capthick: float | None = None, + *, + data=None, + **kwargs, +) -> ErrorbarContainer: + return gca().errorbar( + x, + y, + yerr=yerr, + xerr=xerr, + fmt=fmt, + ecolor=ecolor, + elinewidth=elinewidth, + capsize=capsize, + barsabove=barsabove, + lolims=lolims, + uplims=uplims, + xlolims=xlolims, + xuplims=xuplims, + errorevery=errorevery, + capthick=capthick, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.eventplot) +def eventplot( + positions: ArrayLike | Sequence[ArrayLike], + orientation: Literal["horizontal", "vertical"] = "horizontal", + lineoffsets: float | Sequence[float] = 1, + linelengths: float | Sequence[float] = 1, + linewidths: float | Sequence[float] | None = None, + colors: ColorType | Sequence[ColorType] | None = None, + alpha: float | Sequence[float] | None = None, + linestyles: LineStyleType | Sequence[LineStyleType] = "solid", + *, + data=None, + **kwargs, +) -> EventCollection: + return gca().eventplot( + positions, + orientation=orientation, + lineoffsets=lineoffsets, + linelengths=linelengths, + linewidths=linewidths, + colors=colors, + alpha=alpha, + linestyles=linestyles, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.fill) +def fill(*args, data=None, **kwargs) -> list[Polygon]: + return gca().fill(*args, **({"data": data} if data is not None else {}), **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.fill_between) +def fill_between( + x: ArrayLike, + y1: ArrayLike | float, + y2: ArrayLike | float = 0, + where: Sequence[bool] | None = None, + interpolate: bool = False, + step: Literal["pre", "post", "mid"] | None = None, + *, + data=None, + **kwargs, +) -> FillBetweenPolyCollection: + return gca().fill_between( + x, + y1, + y2=y2, + where=where, + interpolate=interpolate, + step=step, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.fill_betweenx) +def fill_betweenx( + y: ArrayLike, + x1: ArrayLike | float, + x2: ArrayLike | float = 0, + where: Sequence[bool] | None = None, + step: Literal["pre", "post", "mid"] | None = None, + interpolate: bool = False, + *, + data=None, + **kwargs, +) -> FillBetweenPolyCollection: + return gca().fill_betweenx( + y, + x1, + x2=x2, + where=where, + step=step, + interpolate=interpolate, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.grid) +def grid( + visible: bool | None = None, + which: Literal["major", "minor", "both"] = "major", + axis: Literal["both", "x", "y"] = "both", + **kwargs, +) -> None: + gca().grid(visible=visible, which=which, axis=axis, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.hexbin) +def hexbin( + x: ArrayLike, + y: ArrayLike, + C: ArrayLike | None = None, + gridsize: int | tuple[int, int] = 100, + bins: Literal["log"] | int | Sequence[float] | None = None, + xscale: Literal["linear", "log"] = "linear", + yscale: Literal["linear", "log"] = "linear", + extent: tuple[float, float, float, float] | None = None, + cmap: str | Colormap | None = None, + norm: str | Normalize | None = None, + vmin: float | None = None, + vmax: float | None = None, + alpha: float | None = None, + linewidths: float | None = None, + edgecolors: Literal["face", "none"] | ColorType = "face", + reduce_C_function: Callable[[np.ndarray | list[float]], float] = np.mean, + mincnt: int | None = None, + marginals: bool = False, + colorizer: Colorizer | None = None, + *, + data=None, + **kwargs, +) -> PolyCollection: + __ret = gca().hexbin( + x, + y, + C=C, + gridsize=gridsize, + bins=bins, + xscale=xscale, + yscale=yscale, + extent=extent, + cmap=cmap, + norm=norm, + vmin=vmin, + vmax=vmax, + alpha=alpha, + linewidths=linewidths, + edgecolors=edgecolors, + reduce_C_function=reduce_C_function, + mincnt=mincnt, + marginals=marginals, + colorizer=colorizer, + **({"data": data} if data is not None else {}), + **kwargs, + ) + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.hist) +def hist( + x: ArrayLike | Sequence[ArrayLike], + bins: int | Sequence[float] | str | None = None, + range: tuple[float, float] | None = None, + density: bool = False, + weights: ArrayLike | None = None, + cumulative: bool | float = False, + bottom: ArrayLike | float | None = None, + histtype: Literal["bar", "barstacked", "step", "stepfilled"] = "bar", + align: Literal["left", "mid", "right"] = "mid", + orientation: Literal["vertical", "horizontal"] = "vertical", + rwidth: float | None = None, + log: bool = False, + color: ColorType | Sequence[ColorType] | None = None, + label: str | Sequence[str] | None = None, + stacked: bool = False, + *, + data=None, + **kwargs, +) -> tuple[ + np.ndarray | list[np.ndarray], + np.ndarray, + BarContainer | Polygon | list[BarContainer | Polygon], +]: + return gca().hist( + x, + bins=bins, + range=range, + density=density, + weights=weights, + cumulative=cumulative, + bottom=bottom, + histtype=histtype, + align=align, + orientation=orientation, + rwidth=rwidth, + log=log, + color=color, + label=label, + stacked=stacked, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.stairs) +def stairs( + values: ArrayLike, + edges: ArrayLike | None = None, + *, + orientation: Literal["vertical", "horizontal"] = "vertical", + baseline: float | ArrayLike | None = 0, + fill: bool = False, + data=None, + **kwargs, +) -> StepPatch: + return gca().stairs( + values, + edges=edges, + orientation=orientation, + baseline=baseline, + fill=fill, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.hist2d) +def hist2d( + x: ArrayLike, + y: ArrayLike, + bins: None | int | tuple[int, int] | ArrayLike | tuple[ArrayLike, ArrayLike] = 10, + range: ArrayLike | None = None, + density: bool = False, + weights: ArrayLike | None = None, + cmin: float | None = None, + cmax: float | None = None, + *, + data=None, + **kwargs, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, QuadMesh]: + __ret = gca().hist2d( + x, + y, + bins=bins, + range=range, + density=density, + weights=weights, + cmin=cmin, + cmax=cmax, + **({"data": data} if data is not None else {}), + **kwargs, + ) + sci(__ret[-1]) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.hlines) +def hlines( + y: float | ArrayLike, + xmin: float | ArrayLike, + xmax: float | ArrayLike, + colors: ColorType | Sequence[ColorType] | None = None, + linestyles: LineStyleType = "solid", + label: str = "", + *, + data=None, + **kwargs, +) -> LineCollection: + return gca().hlines( + y, + xmin, + xmax, + colors=colors, + linestyles=linestyles, + label=label, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.imshow) +def imshow( + X: ArrayLike | PIL.Image.Image, + cmap: str | Colormap | None = None, + norm: str | Normalize | None = None, + *, + aspect: Literal["equal", "auto"] | float | None = None, + interpolation: str | None = None, + alpha: float | ArrayLike | None = None, + vmin: float | None = None, + vmax: float | None = None, + colorizer: Colorizer | None = None, + origin: Literal["upper", "lower"] | None = None, + extent: tuple[float, float, float, float] | None = None, + interpolation_stage: Literal["data", "rgba", "auto"] | None = None, + filternorm: bool = True, + filterrad: float = 4.0, + resample: bool | None = None, + url: str | None = None, + data=None, + **kwargs, +) -> AxesImage: + __ret = gca().imshow( + X, + cmap=cmap, + norm=norm, + aspect=aspect, + interpolation=interpolation, + alpha=alpha, + vmin=vmin, + vmax=vmax, + colorizer=colorizer, + origin=origin, + extent=extent, + interpolation_stage=interpolation_stage, + filternorm=filternorm, + filterrad=filterrad, + resample=resample, + url=url, + **({"data": data} if data is not None else {}), + **kwargs, + ) + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.legend) +def legend(*args, **kwargs) -> Legend: + return gca().legend(*args, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.locator_params) +def locator_params( + axis: Literal["both", "x", "y"] = "both", tight: bool | None = None, **kwargs +) -> None: + gca().locator_params(axis=axis, tight=tight, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.loglog) +def loglog(*args, **kwargs) -> list[Line2D]: + return gca().loglog(*args, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.magnitude_spectrum) +def magnitude_spectrum( + x: ArrayLike, + Fs: float | None = None, + Fc: int | None = None, + window: Callable[[ArrayLike], ArrayLike] | ArrayLike | None = None, + pad_to: int | None = None, + sides: Literal["default", "onesided", "twosided"] | None = None, + scale: Literal["default", "linear", "dB"] | None = None, + *, + data=None, + **kwargs, +) -> tuple[np.ndarray, np.ndarray, Line2D]: + return gca().magnitude_spectrum( + x, + Fs=Fs, + Fc=Fc, + window=window, + pad_to=pad_to, + sides=sides, + scale=scale, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.margins) +def margins( + *margins: float, + x: float | None = None, + y: float | None = None, + tight: bool | None = True, +) -> tuple[float, float] | None: + return gca().margins(*margins, x=x, y=y, tight=tight) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.minorticks_off) +def minorticks_off() -> None: + gca().minorticks_off() + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.minorticks_on) +def minorticks_on() -> None: + gca().minorticks_on() + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.pcolor) +def pcolor( + *args: ArrayLike, + shading: Literal["flat", "nearest", "auto"] | None = None, + alpha: float | None = None, + norm: str | Normalize | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + colorizer: Colorizer | None = None, + data=None, + **kwargs, +) -> Collection: + __ret = gca().pcolor( + *args, + shading=shading, + alpha=alpha, + norm=norm, + cmap=cmap, + vmin=vmin, + vmax=vmax, + colorizer=colorizer, + **({"data": data} if data is not None else {}), + **kwargs, + ) + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.pcolormesh) +def pcolormesh( + *args: ArrayLike, + alpha: float | None = None, + norm: str | Normalize | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + colorizer: Colorizer | None = None, + shading: Literal["flat", "nearest", "gouraud", "auto"] | None = None, + antialiased: bool = False, + data=None, + **kwargs, +) -> QuadMesh: + __ret = gca().pcolormesh( + *args, + alpha=alpha, + norm=norm, + cmap=cmap, + vmin=vmin, + vmax=vmax, + colorizer=colorizer, + shading=shading, + antialiased=antialiased, + **({"data": data} if data is not None else {}), + **kwargs, + ) + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.phase_spectrum) +def phase_spectrum( + x: ArrayLike, + Fs: float | None = None, + Fc: int | None = None, + window: Callable[[ArrayLike], ArrayLike] | ArrayLike | None = None, + pad_to: int | None = None, + sides: Literal["default", "onesided", "twosided"] | None = None, + *, + data=None, + **kwargs, +) -> tuple[np.ndarray, np.ndarray, Line2D]: + return gca().phase_spectrum( + x, + Fs=Fs, + Fc=Fc, + window=window, + pad_to=pad_to, + sides=sides, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.pie) +def pie( + x: ArrayLike, + explode: ArrayLike | None = None, + labels: Sequence[str] | None = None, + colors: ColorType | Sequence[ColorType] | None = None, + autopct: str | Callable[[float], str] | None = None, + pctdistance: float = 0.6, + shadow: bool = False, + labeldistance: float | None = 1.1, + startangle: float = 0, + radius: float = 1, + counterclock: bool = True, + wedgeprops: dict[str, Any] | None = None, + textprops: dict[str, Any] | None = None, + center: tuple[float, float] = (0, 0), + frame: bool = False, + rotatelabels: bool = False, + *, + normalize: bool = True, + hatch: str | Sequence[str] | None = None, + data=None, +) -> tuple[list[Wedge], list[Text]] | tuple[list[Wedge], list[Text], list[Text]]: + return gca().pie( + x, + explode=explode, + labels=labels, + colors=colors, + autopct=autopct, + pctdistance=pctdistance, + shadow=shadow, + labeldistance=labeldistance, + startangle=startangle, + radius=radius, + counterclock=counterclock, + wedgeprops=wedgeprops, + textprops=textprops, + center=center, + frame=frame, + rotatelabels=rotatelabels, + normalize=normalize, + hatch=hatch, + **({"data": data} if data is not None else {}), + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.plot) +def plot( + *args: float | ArrayLike | str, + scalex: bool = True, + scaley: bool = True, + data=None, + **kwargs, +) -> list[Line2D]: + return gca().plot( + *args, + scalex=scalex, + scaley=scaley, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.plot_date) +def plot_date( + x: ArrayLike, + y: ArrayLike, + fmt: str = "o", + tz: str | datetime.tzinfo | None = None, + xdate: bool = True, + ydate: bool = False, + *, + data=None, + **kwargs, +) -> list[Line2D]: + return gca().plot_date( + x, + y, + fmt=fmt, + tz=tz, + xdate=xdate, + ydate=ydate, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.psd) +def psd( + x: ArrayLike, + NFFT: int | None = None, + Fs: float | None = None, + Fc: int | None = None, + detrend: Literal["none", "mean", "linear"] + | Callable[[ArrayLike], ArrayLike] + | None = None, + window: Callable[[ArrayLike], ArrayLike] | ArrayLike | None = None, + noverlap: int | None = None, + pad_to: int | None = None, + sides: Literal["default", "onesided", "twosided"] | None = None, + scale_by_freq: bool | None = None, + return_line: bool | None = None, + *, + data=None, + **kwargs, +) -> tuple[np.ndarray, np.ndarray] | tuple[np.ndarray, np.ndarray, Line2D]: + return gca().psd( + x, + NFFT=NFFT, + Fs=Fs, + Fc=Fc, + detrend=detrend, + window=window, + noverlap=noverlap, + pad_to=pad_to, + sides=sides, + scale_by_freq=scale_by_freq, + return_line=return_line, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.quiver) +def quiver(*args, data=None, **kwargs) -> Quiver: + __ret = gca().quiver( + *args, **({"data": data} if data is not None else {}), **kwargs + ) + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.quiverkey) +def quiverkey( + Q: Quiver, X: float, Y: float, U: float, label: str, **kwargs +) -> QuiverKey: + return gca().quiverkey(Q, X, Y, U, label, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.scatter) +def scatter( + x: float | ArrayLike, + y: float | ArrayLike, + s: float | ArrayLike | None = None, + c: ArrayLike | Sequence[ColorType] | ColorType | None = None, + marker: MarkerType | None = None, + cmap: str | Colormap | None = None, + norm: str | Normalize | None = None, + vmin: float | None = None, + vmax: float | None = None, + alpha: float | None = None, + linewidths: float | Sequence[float] | None = None, + *, + edgecolors: Literal["face", "none"] | ColorType | Sequence[ColorType] | None = None, + colorizer: Colorizer | None = None, + plotnonfinite: bool = False, + data=None, + **kwargs, +) -> PathCollection: + __ret = gca().scatter( + x, + y, + s=s, + c=c, + marker=marker, + cmap=cmap, + norm=norm, + vmin=vmin, + vmax=vmax, + alpha=alpha, + linewidths=linewidths, + edgecolors=edgecolors, + colorizer=colorizer, + plotnonfinite=plotnonfinite, + **({"data": data} if data is not None else {}), + **kwargs, + ) + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.semilogx) +def semilogx(*args, **kwargs) -> list[Line2D]: + return gca().semilogx(*args, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.semilogy) +def semilogy(*args, **kwargs) -> list[Line2D]: + return gca().semilogy(*args, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.specgram) +def specgram( + x: ArrayLike, + NFFT: int | None = None, + Fs: float | None = None, + Fc: int | None = None, + detrend: Literal["none", "mean", "linear"] + | Callable[[ArrayLike], ArrayLike] + | None = None, + window: Callable[[ArrayLike], ArrayLike] | ArrayLike | None = None, + noverlap: int | None = None, + cmap: str | Colormap | None = None, + xextent: tuple[float, float] | None = None, + pad_to: int | None = None, + sides: Literal["default", "onesided", "twosided"] | None = None, + scale_by_freq: bool | None = None, + mode: Literal["default", "psd", "magnitude", "angle", "phase"] | None = None, + scale: Literal["default", "linear", "dB"] | None = None, + vmin: float | None = None, + vmax: float | None = None, + *, + data=None, + **kwargs, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, AxesImage]: + __ret = gca().specgram( + x, + NFFT=NFFT, + Fs=Fs, + Fc=Fc, + detrend=detrend, + window=window, + noverlap=noverlap, + cmap=cmap, + xextent=xextent, + pad_to=pad_to, + sides=sides, + scale_by_freq=scale_by_freq, + mode=mode, + scale=scale, + vmin=vmin, + vmax=vmax, + **({"data": data} if data is not None else {}), + **kwargs, + ) + sci(__ret[-1]) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.spy) +def spy( + Z: ArrayLike, + precision: float | Literal["present"] = 0, + marker: str | None = None, + markersize: float | None = None, + aspect: Literal["equal", "auto"] | float | None = "equal", + origin: Literal["upper", "lower"] = "upper", + **kwargs, +) -> AxesImage: + __ret = gca().spy( + Z, + precision=precision, + marker=marker, + markersize=markersize, + aspect=aspect, + origin=origin, + **kwargs, + ) + if isinstance(__ret, _ColorizerInterface): + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.stackplot) +def stackplot( + x, *args, labels=(), colors=None, hatch=None, baseline="zero", data=None, **kwargs +): + return gca().stackplot( + x, + *args, + labels=labels, + colors=colors, + hatch=hatch, + baseline=baseline, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.stem) +def stem( + *args: ArrayLike | str, + linefmt: str | None = None, + markerfmt: str | None = None, + basefmt: str | None = None, + bottom: float = 0, + label: str | None = None, + orientation: Literal["vertical", "horizontal"] = "vertical", + data=None, +) -> StemContainer: + return gca().stem( + *args, + linefmt=linefmt, + markerfmt=markerfmt, + basefmt=basefmt, + bottom=bottom, + label=label, + orientation=orientation, + **({"data": data} if data is not None else {}), + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.step) +def step( + x: ArrayLike, + y: ArrayLike, + *args, + where: Literal["pre", "post", "mid"] = "pre", + data=None, + **kwargs, +) -> list[Line2D]: + return gca().step( + x, + y, + *args, + where=where, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.streamplot) +def streamplot( + x, + y, + u, + v, + density=1, + linewidth=None, + color=None, + cmap=None, + norm=None, + arrowsize=1, + arrowstyle="-|>", + minlength=0.1, + transform=None, + zorder=None, + start_points=None, + maxlength=4.0, + integration_direction="both", + broken_streamlines=True, + *, + data=None, +): + __ret = gca().streamplot( + x, + y, + u, + v, + density=density, + linewidth=linewidth, + color=color, + cmap=cmap, + norm=norm, + arrowsize=arrowsize, + arrowstyle=arrowstyle, + minlength=minlength, + transform=transform, + zorder=zorder, + start_points=start_points, + maxlength=maxlength, + integration_direction=integration_direction, + broken_streamlines=broken_streamlines, + **({"data": data} if data is not None else {}), + ) + sci(__ret.lines) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.table) +def table( + cellText=None, + cellColours=None, + cellLoc="right", + colWidths=None, + rowLabels=None, + rowColours=None, + rowLoc="left", + colLabels=None, + colColours=None, + colLoc="center", + loc="bottom", + bbox=None, + edges="closed", + **kwargs, +): + return gca().table( + cellText=cellText, + cellColours=cellColours, + cellLoc=cellLoc, + colWidths=colWidths, + rowLabels=rowLabels, + rowColours=rowColours, + rowLoc=rowLoc, + colLabels=colLabels, + colColours=colColours, + colLoc=colLoc, + loc=loc, + bbox=bbox, + edges=edges, + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.text) +def text( + x: float, y: float, s: str, fontdict: dict[str, Any] | None = None, **kwargs +) -> Text: + return gca().text(x, y, s, fontdict=fontdict, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.tick_params) +def tick_params(axis: Literal["both", "x", "y"] = "both", **kwargs) -> None: + gca().tick_params(axis=axis, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.ticklabel_format) +def ticklabel_format( + *, + axis: Literal["both", "x", "y"] = "both", + style: Literal["", "sci", "scientific", "plain"] | None = None, + scilimits: tuple[int, int] | None = None, + useOffset: bool | float | None = None, + useLocale: bool | None = None, + useMathText: bool | None = None, +) -> None: + gca().ticklabel_format( + axis=axis, + style=style, + scilimits=scilimits, + useOffset=useOffset, + useLocale=useLocale, + useMathText=useMathText, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.tricontour) +def tricontour(*args, **kwargs): + __ret = gca().tricontour(*args, **kwargs) + if __ret._A is not None: # type: ignore[attr-defined] + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.tricontourf) +def tricontourf(*args, **kwargs): + __ret = gca().tricontourf(*args, **kwargs) + if __ret._A is not None: # type: ignore[attr-defined] + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.tripcolor) +def tripcolor( + *args, + alpha=1.0, + norm=None, + cmap=None, + vmin=None, + vmax=None, + shading="flat", + facecolors=None, + **kwargs, +): + __ret = gca().tripcolor( + *args, + alpha=alpha, + norm=norm, + cmap=cmap, + vmin=vmin, + vmax=vmax, + shading=shading, + facecolors=facecolors, + **kwargs, + ) + sci(__ret) + return __ret + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.triplot) +def triplot(*args, **kwargs): + return gca().triplot(*args, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.violinplot) +def violinplot( + dataset: ArrayLike | Sequence[ArrayLike], + positions: ArrayLike | None = None, + vert: bool | None = None, + orientation: Literal["vertical", "horizontal"] = "vertical", + widths: float | ArrayLike = 0.5, + showmeans: bool = False, + showextrema: bool = True, + showmedians: bool = False, + quantiles: Sequence[float | Sequence[float]] | None = None, + points: int = 100, + bw_method: Literal["scott", "silverman"] + | float + | Callable[[GaussianKDE], float] + | None = None, + side: Literal["both", "low", "high"] = "both", + *, + data=None, +) -> dict[str, Collection]: + return gca().violinplot( + dataset, + positions=positions, + vert=vert, + orientation=orientation, + widths=widths, + showmeans=showmeans, + showextrema=showextrema, + showmedians=showmedians, + quantiles=quantiles, + points=points, + bw_method=bw_method, + side=side, + **({"data": data} if data is not None else {}), + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.vlines) +def vlines( + x: float | ArrayLike, + ymin: float | ArrayLike, + ymax: float | ArrayLike, + colors: ColorType | Sequence[ColorType] | None = None, + linestyles: LineStyleType = "solid", + label: str = "", + *, + data=None, + **kwargs, +) -> LineCollection: + return gca().vlines( + x, + ymin, + ymax, + colors=colors, + linestyles=linestyles, + label=label, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.xcorr) +def xcorr( + x: ArrayLike, + y: ArrayLike, + normed: bool = True, + detrend: Callable[[ArrayLike], ArrayLike] = mlab.detrend_none, + usevlines: bool = True, + maxlags: int = 10, + *, + data=None, + **kwargs, +) -> tuple[np.ndarray, np.ndarray, LineCollection | Line2D, Line2D | None]: + return gca().xcorr( + x, + y, + normed=normed, + detrend=detrend, + usevlines=usevlines, + maxlags=maxlags, + **({"data": data} if data is not None else {}), + **kwargs, + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes._sci) +def sci(im: ColorizingArtist) -> None: + gca()._sci(im) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.set_title) +def title( + label: str, + fontdict: dict[str, Any] | None = None, + loc: Literal["left", "center", "right"] | None = None, + pad: float | None = None, + *, + y: float | None = None, + **kwargs, +) -> Text: + return gca().set_title(label, fontdict=fontdict, loc=loc, pad=pad, y=y, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.set_xlabel) +def xlabel( + xlabel: str, + fontdict: dict[str, Any] | None = None, + labelpad: float | None = None, + *, + loc: Literal["left", "center", "right"] | None = None, + **kwargs, +) -> Text: + return gca().set_xlabel( + xlabel, fontdict=fontdict, labelpad=labelpad, loc=loc, **kwargs + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.set_ylabel) +def ylabel( + ylabel: str, + fontdict: dict[str, Any] | None = None, + labelpad: float | None = None, + *, + loc: Literal["bottom", "center", "top"] | None = None, + **kwargs, +) -> Text: + return gca().set_ylabel( + ylabel, fontdict=fontdict, labelpad=labelpad, loc=loc, **kwargs + ) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.set_xscale) +def xscale(value: str | ScaleBase, **kwargs) -> None: + gca().set_xscale(value, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +@_copy_docstring_and_deprecators(Axes.set_yscale) +def yscale(value: str | ScaleBase, **kwargs) -> None: + gca().set_yscale(value, **kwargs) + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def autumn() -> None: + """ + Set the colormap to 'autumn'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("autumn") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def bone() -> None: + """ + Set the colormap to 'bone'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("bone") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def cool() -> None: + """ + Set the colormap to 'cool'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("cool") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def copper() -> None: + """ + Set the colormap to 'copper'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("copper") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def flag() -> None: + """ + Set the colormap to 'flag'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("flag") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def gray() -> None: + """ + Set the colormap to 'gray'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("gray") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def hot() -> None: + """ + Set the colormap to 'hot'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("hot") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def hsv() -> None: + """ + Set the colormap to 'hsv'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("hsv") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def jet() -> None: + """ + Set the colormap to 'jet'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("jet") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def pink() -> None: + """ + Set the colormap to 'pink'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("pink") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def prism() -> None: + """ + Set the colormap to 'prism'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("prism") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def spring() -> None: + """ + Set the colormap to 'spring'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("spring") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def summer() -> None: + """ + Set the colormap to 'summer'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("summer") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def winter() -> None: + """ + Set the colormap to 'winter'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("winter") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def magma() -> None: + """ + Set the colormap to 'magma'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("magma") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def inferno() -> None: + """ + Set the colormap to 'inferno'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("inferno") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def plasma() -> None: + """ + Set the colormap to 'plasma'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("plasma") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def viridis() -> None: + """ + Set the colormap to 'viridis'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("viridis") + + +# Autogenerated by boilerplate.py. Do not edit as changes will be lost. +def nipy_spectral() -> None: + """ + Set the colormap to 'nipy_spectral'. + + This changes the default colormap as well as the colormap of the current + image if there is one. See ``help(colormaps)`` for more information. + """ + set_cmap("nipy_spectral") diff --git a/lib/python3.12/site-packages/matplotlib/sankey.pyi b/lib/python3.12/site-packages/matplotlib/sankey.pyi new file mode 100644 index 0000000000000000000000000000000000000000..33565b998a9caf00ead62bdfa0adb50e272dd6ac --- /dev/null +++ b/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]: ... diff --git a/lib/python3.12/site-packages/matplotlib/scale.py b/lib/python3.12/site-packages/matplotlib/scale.py new file mode 100644 index 0000000000000000000000000000000000000000..7630ba6e54b91dacf533aa196ab12c5570cc76d2 --- /dev/null +++ b/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 ` 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(), + ) diff --git a/lib/python3.12/site-packages/matplotlib/texmanager.py b/lib/python3.12/site-packages/matplotlib/texmanager.py new file mode 100644 index 0000000000000000000000000000000000000000..2083510396f1af8325dafd98396f62a68c142f7d --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/matplotlib/transforms.py b/lib/python3.12/site-packages/matplotlib/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..2934b0a77809cb160b516557f6eca1ca5cadca8f --- /dev/null +++ b/lib/python3.12/site-packages/matplotlib/transforms.py @@ -0,0 +1,2995 @@ +""" +Matplotlib includes a framework for arbitrary geometric transformations that is used to +determine the final position of all elements drawn on the canvas. + +Transforms are composed into trees of `TransformNode` objects +whose actual value depends on their children. When the contents of +children change, their parents are automatically invalidated. The +next time an invalidated transform is accessed, it is recomputed to +reflect those changes. This invalidation/caching approach prevents +unnecessary recomputations of transforms, and contributes to better +interactive performance. + +For example, here is a graph of the transform tree used to plot data to the figure: + +.. graphviz:: /api/transforms.dot + :alt: Diagram of transform tree from data to figure coordinates. + +The framework can be used for both affine and non-affine +transformations. However, for speed, we want to use the backend +renderers to perform affine transformations whenever possible. +Therefore, it is possible to perform just the affine or non-affine +part of a transformation on a set of data. The affine is always +assumed to occur after the non-affine. For any transform:: + + full transform == non-affine part + affine part + +The backends are not expected to handle non-affine transformations +themselves. + +See the tutorial :ref:`transforms_tutorial` for examples +of how to use transforms. +""" + +# Note: There are a number of places in the code where we use `np.min` or +# `np.minimum` instead of the builtin `min`, and likewise for `max`. This is +# done so that `nan`s are propagated, instead of being silently dropped. + +import copy +import functools +import itertools +import textwrap +import weakref +import math + +import numpy as np +from numpy.linalg import inv + +from matplotlib import _api +from matplotlib._path import ( + affine_transform, count_bboxes_overlapping_bbox, update_path_extents) +from .path import Path + +DEBUG = False + + +def _make_str_method(*args, **kwargs): + """ + Generate a ``__str__`` method for a `.Transform` subclass. + + After :: + + class T: + __str__ = _make_str_method("attr", key="other") + + ``str(T(...))`` will be + + .. code-block:: text + + {type(T).__name__}( + {self.attr}, + key={self.other}) + """ + indent = functools.partial(textwrap.indent, prefix=" " * 4) + def strrepr(x): return repr(x) if isinstance(x, str) else str(x) + return lambda self: ( + type(self).__name__ + "(" + + ",".join([*(indent("\n" + strrepr(getattr(self, arg))) + for arg in args), + *(indent("\n" + k + "=" + strrepr(getattr(self, arg))) + for k, arg in kwargs.items())]) + + ")") + + +class TransformNode: + """ + The base class for anything that participates in the transform tree + and needs to invalidate its parents or be invalidated. This includes + classes that are not really transforms, such as bounding boxes, since some + transforms depend on bounding boxes to compute their values. + """ + + # Invalidation may affect only the affine part. If the + # invalidation was "affine-only", the _invalid member is set to + # INVALID_AFFINE_ONLY + + # Possible values for the _invalid attribute. + _VALID, _INVALID_AFFINE_ONLY, _INVALID_FULL = range(3) + + # Some metadata about the transform, used to determine whether an + # invalidation is affine-only + is_affine = False + is_bbox = _api.deprecated("3.9")(_api.classproperty(lambda cls: False)) + + pass_through = False + """ + If pass_through is True, all ancestors will always be + invalidated, even if 'self' is already invalid. + """ + + def __init__(self, shorthand_name=None): + """ + Parameters + ---------- + shorthand_name : str + A string representing the "name" of the transform. The name carries + no significance other than to improve the readability of + ``str(transform)`` when DEBUG=True. + """ + self._parents = {} + # Initially invalid, until first computation. + self._invalid = self._INVALID_FULL + self._shorthand_name = shorthand_name or '' + + if DEBUG: + def __str__(self): + # either just return the name of this TransformNode, or its repr + return self._shorthand_name or repr(self) + + def __getstate__(self): + # turn the dictionary with weak values into a normal dictionary + return {**self.__dict__, + '_parents': {k: v() for k, v in self._parents.items()}} + + def __setstate__(self, data_dict): + self.__dict__ = data_dict + # turn the normal dictionary back into a dictionary with weak values + # The extra lambda is to provide a callback to remove dead + # weakrefs from the dictionary when garbage collection is done. + self._parents = { + k: weakref.ref(v, lambda _, pop=self._parents.pop, k=k: pop(k)) + for k, v in self._parents.items() if v is not None} + + def __copy__(self): + other = copy.copy(super()) + # If `c = a + b; a1 = copy(a)`, then modifications to `a1` do not + # propagate back to `c`, i.e. we need to clear the parents of `a1`. + other._parents = {} + # If `c = a + b; c1 = copy(c)`, then modifications to `a` also need to + # be propagated to `c1`. + for key, val in vars(self).items(): + if isinstance(val, TransformNode) and id(self) in val._parents: + other.set_children(val) # val == getattr(other, key) + return other + + def invalidate(self): + """ + Invalidate this `TransformNode` and triggers an invalidation of its + ancestors. Should be called any time the transform changes. + """ + return self._invalidate_internal( + level=self._INVALID_AFFINE_ONLY if self.is_affine else self._INVALID_FULL, + invalidating_node=self) + + def _invalidate_internal(self, level, invalidating_node): + """ + Called by :meth:`invalidate` and subsequently ascends the transform + stack calling each TransformNode's _invalidate_internal method. + """ + # If we are already more invalid than the currently propagated invalidation, + # then we don't need to do anything. + if level <= self._invalid and not self.pass_through: + return + self._invalid = level + for parent in list(self._parents.values()): + parent = parent() # Dereference the weak reference. + if parent is not None: + parent._invalidate_internal(level=level, invalidating_node=self) + + def set_children(self, *children): + """ + Set the children of the transform, to let the invalidation + system know which transforms can invalidate this transform. + Should be called from the constructor of any transforms that + depend on other transforms. + """ + # Parents are stored as weak references, so that if the + # parents are destroyed, references from the children won't + # keep them alive. + id_self = id(self) + for child in children: + # Use weak references so this dictionary won't keep obsolete nodes + # alive; the callback deletes the dictionary entry. This is a + # performance improvement over using WeakValueDictionary. + ref = weakref.ref( + self, lambda _, pop=child._parents.pop, k=id_self: pop(k)) + child._parents[id_self] = ref + + def frozen(self): + """ + Return a frozen copy of this transform node. The frozen copy will not + be updated when its children change. Useful for storing a previously + known state of a transform where ``copy.deepcopy()`` might normally be + used. + """ + return self + + +class BboxBase(TransformNode): + """ + The base class of all bounding boxes. + + This class is immutable; `Bbox` is a mutable subclass. + + The canonical representation is as two points, with no + restrictions on their ordering. Convenience properties are + provided to get the left, bottom, right and top edges and width + and height, but these are not stored explicitly. + """ + + is_bbox = _api.deprecated("3.9")(_api.classproperty(lambda cls: True)) + is_affine = True + + if DEBUG: + @staticmethod + def _check(points): + if isinstance(points, np.ma.MaskedArray): + _api.warn_external("Bbox bounds are a masked array.") + points = np.asarray(points) + if any((points[1, :] - points[0, :]) == 0): + _api.warn_external("Singular Bbox.") + + def frozen(self): + return Bbox(self.get_points().copy()) + frozen.__doc__ = TransformNode.__doc__ + + def __array__(self, *args, **kwargs): + return self.get_points() + + @property + def x0(self): + """ + The first of the pair of *x* coordinates that define the bounding box. + + This is not guaranteed to be less than :attr:`x1` (for that, use + :attr:`xmin`). + """ + return self.get_points()[0, 0] + + @property + def y0(self): + """ + The first of the pair of *y* coordinates that define the bounding box. + + This is not guaranteed to be less than :attr:`y1` (for that, use + :attr:`ymin`). + """ + return self.get_points()[0, 1] + + @property + def x1(self): + """ + The second of the pair of *x* coordinates that define the bounding box. + + This is not guaranteed to be greater than :attr:`x0` (for that, use + :attr:`xmax`). + """ + return self.get_points()[1, 0] + + @property + def y1(self): + """ + The second of the pair of *y* coordinates that define the bounding box. + + This is not guaranteed to be greater than :attr:`y0` (for that, use + :attr:`ymax`). + """ + return self.get_points()[1, 1] + + @property + def p0(self): + """ + The first pair of (*x*, *y*) coordinates that define the bounding box. + + This is not guaranteed to be the bottom-left corner (for that, use + :attr:`min`). + """ + return self.get_points()[0] + + @property + def p1(self): + """ + The second pair of (*x*, *y*) coordinates that define the bounding box. + + This is not guaranteed to be the top-right corner (for that, use + :attr:`max`). + """ + return self.get_points()[1] + + @property + def xmin(self): + """The left edge of the bounding box.""" + return np.min(self.get_points()[:, 0]) + + @property + def ymin(self): + """The bottom edge of the bounding box.""" + return np.min(self.get_points()[:, 1]) + + @property + def xmax(self): + """The right edge of the bounding box.""" + return np.max(self.get_points()[:, 0]) + + @property + def ymax(self): + """The top edge of the bounding box.""" + return np.max(self.get_points()[:, 1]) + + @property + def min(self): + """The bottom-left corner of the bounding box.""" + return np.min(self.get_points(), axis=0) + + @property + def max(self): + """The top-right corner of the bounding box.""" + return np.max(self.get_points(), axis=0) + + @property + def intervalx(self): + """ + The pair of *x* coordinates that define the bounding box. + + This is not guaranteed to be sorted from left to right. + """ + return self.get_points()[:, 0] + + @property + def intervaly(self): + """ + The pair of *y* coordinates that define the bounding box. + + This is not guaranteed to be sorted from bottom to top. + """ + return self.get_points()[:, 1] + + @property + def width(self): + """The (signed) width of the bounding box.""" + points = self.get_points() + return points[1, 0] - points[0, 0] + + @property + def height(self): + """The (signed) height of the bounding box.""" + points = self.get_points() + return points[1, 1] - points[0, 1] + + @property + def size(self): + """The (signed) width and height of the bounding box.""" + points = self.get_points() + return points[1] - points[0] + + @property + def bounds(self): + """Return (:attr:`x0`, :attr:`y0`, :attr:`width`, :attr:`height`).""" + (x0, y0), (x1, y1) = self.get_points() + return (x0, y0, x1 - x0, y1 - y0) + + @property + def extents(self): + """Return (:attr:`x0`, :attr:`y0`, :attr:`x1`, :attr:`y1`).""" + return self.get_points().flatten() # flatten returns a copy. + + def get_points(self): + raise NotImplementedError + + def containsx(self, x): + """ + Return whether *x* is in the closed (:attr:`x0`, :attr:`x1`) interval. + """ + x0, x1 = self.intervalx + return x0 <= x <= x1 or x0 >= x >= x1 + + def containsy(self, y): + """ + Return whether *y* is in the closed (:attr:`y0`, :attr:`y1`) interval. + """ + y0, y1 = self.intervaly + return y0 <= y <= y1 or y0 >= y >= y1 + + def contains(self, x, y): + """ + Return whether ``(x, y)`` is in the bounding box or on its edge. + """ + return self.containsx(x) and self.containsy(y) + + def overlaps(self, other): + """ + Return whether this bounding box overlaps with the other bounding box. + + Parameters + ---------- + other : `.BboxBase` + """ + ax1, ay1, ax2, ay2 = self.extents + bx1, by1, bx2, by2 = other.extents + if ax2 < ax1: + ax2, ax1 = ax1, ax2 + if ay2 < ay1: + ay2, ay1 = ay1, ay2 + if bx2 < bx1: + bx2, bx1 = bx1, bx2 + if by2 < by1: + by2, by1 = by1, by2 + return ax1 <= bx2 and bx1 <= ax2 and ay1 <= by2 and by1 <= ay2 + + def fully_containsx(self, x): + """ + Return whether *x* is in the open (:attr:`x0`, :attr:`x1`) interval. + """ + x0, x1 = self.intervalx + return x0 < x < x1 or x0 > x > x1 + + def fully_containsy(self, y): + """ + Return whether *y* is in the open (:attr:`y0`, :attr:`y1`) interval. + """ + y0, y1 = self.intervaly + return y0 < y < y1 or y0 > y > y1 + + def fully_contains(self, x, y): + """ + Return whether ``x, y`` is in the bounding box, but not on its edge. + """ + return self.fully_containsx(x) and self.fully_containsy(y) + + def fully_overlaps(self, other): + """ + Return whether this bounding box overlaps with the other bounding box, + not including the edges. + + Parameters + ---------- + other : `.BboxBase` + """ + ax1, ay1, ax2, ay2 = self.extents + bx1, by1, bx2, by2 = other.extents + if ax2 < ax1: + ax2, ax1 = ax1, ax2 + if ay2 < ay1: + ay2, ay1 = ay1, ay2 + if bx2 < bx1: + bx2, bx1 = bx1, bx2 + if by2 < by1: + by2, by1 = by1, by2 + return ax1 < bx2 and bx1 < ax2 and ay1 < by2 and by1 < ay2 + + def transformed(self, transform): + """ + Construct a `Bbox` by statically transforming this one by *transform*. + """ + pts = self.get_points() + ll, ul, lr = transform.transform(np.array( + [pts[0], [pts[0, 0], pts[1, 1]], [pts[1, 0], pts[0, 1]]])) + return Bbox([ll, [lr[0], ul[1]]]) + + coefs = {'C': (0.5, 0.5), + 'SW': (0, 0), + 'S': (0.5, 0), + 'SE': (1.0, 0), + 'E': (1.0, 0.5), + 'NE': (1.0, 1.0), + 'N': (0.5, 1.0), + 'NW': (0, 1.0), + 'W': (0, 0.5)} + + def anchored(self, c, container): + """ + Return a copy of the `Bbox` anchored to *c* within *container*. + + Parameters + ---------- + c : (float, float) or {'C', 'SW', 'S', 'SE', 'E', 'NE', ...} + Either an (*x*, *y*) pair of relative coordinates (0 is left or + bottom, 1 is right or top), 'C' (center), or a cardinal direction + ('SW', southwest, is bottom left, etc.). + container : `Bbox` + The box within which the `Bbox` is positioned. + + See Also + -------- + .Axes.set_anchor + """ + l, b, w, h = container.bounds + L, B, W, H = self.bounds + cx, cy = self.coefs[c] if isinstance(c, str) else c + return Bbox(self._points + + [(l + cx * (w - W)) - L, + (b + cy * (h - H)) - B]) + + def shrunk(self, mx, my): + """ + Return a copy of the `Bbox`, shrunk by the factor *mx* + in the *x* direction and the factor *my* in the *y* direction. + The lower left corner of the box remains unchanged. Normally + *mx* and *my* will be less than 1, but this is not enforced. + """ + w, h = self.size + return Bbox([self._points[0], + self._points[0] + [mx * w, my * h]]) + + def shrunk_to_aspect(self, box_aspect, container=None, fig_aspect=1.0): + """ + Return a copy of the `Bbox`, shrunk so that it is as + large as it can be while having the desired aspect ratio, + *box_aspect*. If the box coordinates are relative (i.e. + fractions of a larger box such as a figure) then the + physical aspect ratio of that figure is specified with + *fig_aspect*, so that *box_aspect* can also be given as a + ratio of the absolute dimensions, not the relative dimensions. + """ + if box_aspect <= 0 or fig_aspect <= 0: + raise ValueError("'box_aspect' and 'fig_aspect' must be positive") + if container is None: + container = self + w, h = container.size + H = w * box_aspect / fig_aspect + if H <= h: + W = w + else: + W = h * fig_aspect / box_aspect + H = h + return Bbox([self._points[0], + self._points[0] + (W, H)]) + + def splitx(self, *args): + """ + Return a list of new `Bbox` objects formed by splitting the original + one with vertical lines at fractional positions given by *args*. + """ + xf = [0, *args, 1] + x0, y0, x1, y1 = self.extents + w = x1 - x0 + return [Bbox([[x0 + xf0 * w, y0], [x0 + xf1 * w, y1]]) + for xf0, xf1 in itertools.pairwise(xf)] + + def splity(self, *args): + """ + Return a list of new `Bbox` objects formed by splitting the original + one with horizontal lines at fractional positions given by *args*. + """ + yf = [0, *args, 1] + x0, y0, x1, y1 = self.extents + h = y1 - y0 + return [Bbox([[x0, y0 + yf0 * h], [x1, y0 + yf1 * h]]) + for yf0, yf1 in itertools.pairwise(yf)] + + def count_contains(self, vertices): + """ + Count the number of vertices contained in the `Bbox`. + Any vertices with a non-finite x or y value are ignored. + + Parameters + ---------- + vertices : (N, 2) array + """ + if len(vertices) == 0: + return 0 + vertices = np.asarray(vertices) + with np.errstate(invalid='ignore'): + return (((self.min < vertices) & + (vertices < self.max)).all(axis=1).sum()) + + def count_overlaps(self, bboxes): + """ + Count the number of bounding boxes that overlap this one. + + Parameters + ---------- + bboxes : sequence of `.BboxBase` + """ + return count_bboxes_overlapping_bbox( + self, np.atleast_3d([np.array(x) for x in bboxes])) + + def expanded(self, sw, sh): + """ + Construct a `Bbox` by expanding this one around its center by the + factors *sw* and *sh*. + """ + width = self.width + height = self.height + deltaw = (sw * width - width) / 2.0 + deltah = (sh * height - height) / 2.0 + a = np.array([[-deltaw, -deltah], [deltaw, deltah]]) + return Bbox(self._points + a) + + def padded(self, w_pad, h_pad=None): + """ + Construct a `Bbox` by padding this one on all four sides. + + Parameters + ---------- + w_pad : float + Width pad + h_pad : float, optional + Height pad. Defaults to *w_pad*. + + """ + points = self.get_points() + if h_pad is None: + h_pad = w_pad + return Bbox(points + [[-w_pad, -h_pad], [w_pad, h_pad]]) + + def translated(self, tx, ty): + """Construct a `Bbox` by translating this one by *tx* and *ty*.""" + return Bbox(self._points + (tx, ty)) + + def corners(self): + """ + Return the corners of this rectangle as an array of points. + + Specifically, this returns the array + ``[[x0, y0], [x0, y1], [x1, y0], [x1, y1]]``. + """ + (x0, y0), (x1, y1) = self.get_points() + return np.array([[x0, y0], [x0, y1], [x1, y0], [x1, y1]]) + + def rotated(self, radians): + """ + Return the axes-aligned bounding box that bounds the result of rotating + this `Bbox` by an angle of *radians*. + """ + corners = self.corners() + corners_rotated = Affine2D().rotate(radians).transform(corners) + bbox = Bbox.unit() + bbox.update_from_data_xy(corners_rotated, ignore=True) + return bbox + + @staticmethod + def union(bboxes): + """Return a `Bbox` that contains all of the given *bboxes*.""" + if not len(bboxes): + raise ValueError("'bboxes' cannot be empty") + x0 = np.min([bbox.xmin for bbox in bboxes]) + x1 = np.max([bbox.xmax for bbox in bboxes]) + y0 = np.min([bbox.ymin for bbox in bboxes]) + y1 = np.max([bbox.ymax for bbox in bboxes]) + return Bbox([[x0, y0], [x1, y1]]) + + @staticmethod + def intersection(bbox1, bbox2): + """ + Return the intersection of *bbox1* and *bbox2* if they intersect, or + None if they don't. + """ + x0 = np.maximum(bbox1.xmin, bbox2.xmin) + x1 = np.minimum(bbox1.xmax, bbox2.xmax) + y0 = np.maximum(bbox1.ymin, bbox2.ymin) + y1 = np.minimum(bbox1.ymax, bbox2.ymax) + return Bbox([[x0, y0], [x1, y1]]) if x0 <= x1 and y0 <= y1 else None + + +_default_minpos = np.array([np.inf, np.inf]) + + +class Bbox(BboxBase): + """ + A mutable bounding box. + + Examples + -------- + **Create from known bounds** + + The default constructor takes the boundary "points" ``[[xmin, ymin], + [xmax, ymax]]``. + + >>> Bbox([[1, 1], [3, 7]]) + Bbox([[1.0, 1.0], [3.0, 7.0]]) + + Alternatively, a Bbox can be created from the flattened points array, the + so-called "extents" ``(xmin, ymin, xmax, ymax)`` + + >>> Bbox.from_extents(1, 1, 3, 7) + Bbox([[1.0, 1.0], [3.0, 7.0]]) + + or from the "bounds" ``(xmin, ymin, width, height)``. + + >>> Bbox.from_bounds(1, 1, 2, 6) + Bbox([[1.0, 1.0], [3.0, 7.0]]) + + **Create from collections of points** + + The "empty" object for accumulating Bboxs is the null bbox, which is a + stand-in for the empty set. + + >>> Bbox.null() + Bbox([[inf, inf], [-inf, -inf]]) + + Adding points to the null bbox will give you the bbox of those points. + + >>> box = Bbox.null() + >>> box.update_from_data_xy([[1, 1]]) + >>> box + Bbox([[1.0, 1.0], [1.0, 1.0]]) + >>> box.update_from_data_xy([[2, 3], [3, 2]], ignore=False) + >>> box + Bbox([[1.0, 1.0], [3.0, 3.0]]) + + Setting ``ignore=True`` is equivalent to starting over from a null bbox. + + >>> box.update_from_data_xy([[1, 1]], ignore=True) + >>> box + Bbox([[1.0, 1.0], [1.0, 1.0]]) + + .. warning:: + + It is recommended to always specify ``ignore`` explicitly. If not, the + default value of ``ignore`` can be changed at any time by code with + access to your Bbox, for example using the method `~.Bbox.ignore`. + + **Properties of the ``null`` bbox** + + .. note:: + + The current behavior of `Bbox.null()` may be surprising as it does + not have all of the properties of the "empty set", and as such does + not behave like a "zero" object in the mathematical sense. We may + change that in the future (with a deprecation period). + + The null bbox is the identity for intersections + + >>> Bbox.intersection(Bbox([[1, 1], [3, 7]]), Bbox.null()) + Bbox([[1.0, 1.0], [3.0, 7.0]]) + + except with itself, where it returns the full space. + + >>> Bbox.intersection(Bbox.null(), Bbox.null()) + Bbox([[-inf, -inf], [inf, inf]]) + + A union containing null will always return the full space (not the other + set!) + + >>> Bbox.union([Bbox([[0, 0], [0, 0]]), Bbox.null()]) + Bbox([[-inf, -inf], [inf, inf]]) + """ + + def __init__(self, points, **kwargs): + """ + Parameters + ---------- + points : `~numpy.ndarray` + A (2, 2) array of the form ``[[x0, y0], [x1, y1]]``. + """ + super().__init__(**kwargs) + points = np.asarray(points, float) + if points.shape != (2, 2): + raise ValueError('Bbox points must be of the form ' + '"[[x0, y0], [x1, y1]]".') + self._points = points + self._minpos = _default_minpos.copy() + self._ignore = True + # it is helpful in some contexts to know if the bbox is a + # default or has been mutated; we store the orig points to + # support the mutated methods + self._points_orig = self._points.copy() + if DEBUG: + ___init__ = __init__ + + def __init__(self, points, **kwargs): + self._check(points) + self.___init__(points, **kwargs) + + def invalidate(self): + self._check(self._points) + super().invalidate() + + def frozen(self): + # docstring inherited + frozen_bbox = super().frozen() + frozen_bbox._minpos = self.minpos.copy() + return frozen_bbox + + @staticmethod + def unit(): + """Create a new unit `Bbox` from (0, 0) to (1, 1).""" + return Bbox([[0, 0], [1, 1]]) + + @staticmethod + def null(): + """Create a new null `Bbox` from (inf, inf) to (-inf, -inf).""" + return Bbox([[np.inf, np.inf], [-np.inf, -np.inf]]) + + @staticmethod + def from_bounds(x0, y0, width, height): + """ + Create a new `Bbox` from *x0*, *y0*, *width* and *height*. + + *width* and *height* may be negative. + """ + return Bbox.from_extents(x0, y0, x0 + width, y0 + height) + + @staticmethod + def from_extents(*args, minpos=None): + """ + Create a new Bbox from *left*, *bottom*, *right* and *top*. + + The *y*-axis increases upwards. + + Parameters + ---------- + left, bottom, right, top : float + The four extents of the bounding box. + minpos : float or None + If this is supplied, the Bbox will have a minimum positive value + set. This is useful when dealing with logarithmic scales and other + scales where negative bounds result in floating point errors. + """ + bbox = Bbox(np.reshape(args, (2, 2))) + if minpos is not None: + bbox._minpos[:] = minpos + return bbox + + def __format__(self, fmt): + return ( + 'Bbox(x0={0.x0:{1}}, y0={0.y0:{1}}, x1={0.x1:{1}}, y1={0.y1:{1}})'. + format(self, fmt)) + + def __str__(self): + return format(self, '') + + def __repr__(self): + return 'Bbox([[{0.x0}, {0.y0}], [{0.x1}, {0.y1}]])'.format(self) + + def ignore(self, value): + """ + Set whether the existing bounds of the box should be ignored + by subsequent calls to :meth:`update_from_data_xy`. + + value : bool + - When ``True``, subsequent calls to `update_from_data_xy` will + ignore the existing bounds of the `Bbox`. + - When ``False``, subsequent calls to `update_from_data_xy` will + include the existing bounds of the `Bbox`. + """ + self._ignore = value + + def update_from_path(self, path, ignore=None, updatex=True, updatey=True): + """ + Update the bounds of the `Bbox` to contain the vertices of the + provided path. After updating, the bounds will have positive *width* + and *height*; *x0* and *y0* will be the minimal values. + + Parameters + ---------- + path : `~matplotlib.path.Path` + ignore : bool, optional + - When ``True``, ignore the existing bounds of the `Bbox`. + - When ``False``, include the existing bounds of the `Bbox`. + - When ``None``, use the last value passed to :meth:`ignore`. + updatex, updatey : bool, default: True + When ``True``, update the x/y values. + """ + if ignore is None: + ignore = self._ignore + + if path.vertices.size == 0: + return + + points, minpos, changed = update_path_extents( + path, None, self._points, self._minpos, ignore) + + if changed: + self.invalidate() + if updatex: + self._points[:, 0] = points[:, 0] + self._minpos[0] = minpos[0] + if updatey: + self._points[:, 1] = points[:, 1] + self._minpos[1] = minpos[1] + + def update_from_data_x(self, x, ignore=None): + """ + Update the x-bounds of the `Bbox` based on the passed in data. After + updating, the bounds will have positive *width*, and *x0* will be the + minimal value. + + Parameters + ---------- + x : `~numpy.ndarray` + Array of x-values. + ignore : bool, optional + - When ``True``, ignore the existing bounds of the `Bbox`. + - When ``False``, include the existing bounds of the `Bbox`. + - When ``None``, use the last value passed to :meth:`ignore`. + """ + x = np.ravel(x) + self.update_from_data_xy(np.column_stack([x, np.ones(x.size)]), + ignore=ignore, updatey=False) + + def update_from_data_y(self, y, ignore=None): + """ + Update the y-bounds of the `Bbox` based on the passed in data. After + updating, the bounds will have positive *height*, and *y0* will be the + minimal value. + + Parameters + ---------- + y : `~numpy.ndarray` + Array of y-values. + ignore : bool, optional + - When ``True``, ignore the existing bounds of the `Bbox`. + - When ``False``, include the existing bounds of the `Bbox`. + - When ``None``, use the last value passed to :meth:`ignore`. + """ + y = np.ravel(y) + self.update_from_data_xy(np.column_stack([np.ones(y.size), y]), + ignore=ignore, updatex=False) + + def update_from_data_xy(self, xy, ignore=None, updatex=True, updatey=True): + """ + Update the `Bbox` bounds based on the passed in *xy* coordinates. + + After updating, the bounds will have positive *width* and *height*; + *x0* and *y0* will be the minimal values. + + Parameters + ---------- + xy : (N, 2) array-like + The (x, y) coordinates. + ignore : bool, optional + - When ``True``, ignore the existing bounds of the `Bbox`. + - When ``False``, include the existing bounds of the `Bbox`. + - When ``None``, use the last value passed to :meth:`ignore`. + updatex, updatey : bool, default: True + When ``True``, update the x/y values. + """ + if len(xy) == 0: + return + + path = Path(xy) + self.update_from_path(path, ignore=ignore, + updatex=updatex, updatey=updatey) + + @BboxBase.x0.setter + def x0(self, val): + self._points[0, 0] = val + self.invalidate() + + @BboxBase.y0.setter + def y0(self, val): + self._points[0, 1] = val + self.invalidate() + + @BboxBase.x1.setter + def x1(self, val): + self._points[1, 0] = val + self.invalidate() + + @BboxBase.y1.setter + def y1(self, val): + self._points[1, 1] = val + self.invalidate() + + @BboxBase.p0.setter + def p0(self, val): + self._points[0] = val + self.invalidate() + + @BboxBase.p1.setter + def p1(self, val): + self._points[1] = val + self.invalidate() + + @BboxBase.intervalx.setter + def intervalx(self, interval): + self._points[:, 0] = interval + self.invalidate() + + @BboxBase.intervaly.setter + def intervaly(self, interval): + self._points[:, 1] = interval + self.invalidate() + + @BboxBase.bounds.setter + def bounds(self, bounds): + l, b, w, h = bounds + points = np.array([[l, b], [l + w, b + h]], float) + if np.any(self._points != points): + self._points = points + self.invalidate() + + @property + def minpos(self): + """ + The minimum positive value in both directions within the Bbox. + + This is useful when dealing with logarithmic scales and other scales + where negative bounds result in floating point errors, and will be used + as the minimum extent instead of *p0*. + """ + return self._minpos + + @minpos.setter + def minpos(self, val): + self._minpos[:] = val + + @property + def minposx(self): + """ + The minimum positive value in the *x*-direction within the Bbox. + + This is useful when dealing with logarithmic scales and other scales + where negative bounds result in floating point errors, and will be used + as the minimum *x*-extent instead of *x0*. + """ + return self._minpos[0] + + @minposx.setter + def minposx(self, val): + self._minpos[0] = val + + @property + def minposy(self): + """ + The minimum positive value in the *y*-direction within the Bbox. + + This is useful when dealing with logarithmic scales and other scales + where negative bounds result in floating point errors, and will be used + as the minimum *y*-extent instead of *y0*. + """ + return self._minpos[1] + + @minposy.setter + def minposy(self, val): + self._minpos[1] = val + + def get_points(self): + """ + Get the points of the bounding box as an array of the form + ``[[x0, y0], [x1, y1]]``. + """ + self._invalid = 0 + return self._points + + def set_points(self, points): + """ + Set the points of the bounding box directly from an array of the form + ``[[x0, y0], [x1, y1]]``. No error checking is performed, as this + method is mainly for internal use. + """ + if np.any(self._points != points): + self._points = points + self.invalidate() + + def set(self, other): + """ + Set this bounding box from the "frozen" bounds of another `Bbox`. + """ + if np.any(self._points != other.get_points()): + self._points = other.get_points() + self.invalidate() + + def mutated(self): + """Return whether the bbox has changed since init.""" + return self.mutatedx() or self.mutatedy() + + def mutatedx(self): + """Return whether the x-limits have changed since init.""" + return (self._points[0, 0] != self._points_orig[0, 0] or + self._points[1, 0] != self._points_orig[1, 0]) + + def mutatedy(self): + """Return whether the y-limits have changed since init.""" + return (self._points[0, 1] != self._points_orig[0, 1] or + self._points[1, 1] != self._points_orig[1, 1]) + + +class TransformedBbox(BboxBase): + """ + A `Bbox` that is automatically transformed by a given + transform. When either the child bounding box or transform + changes, the bounds of this bbox will update accordingly. + """ + + def __init__(self, bbox, transform, **kwargs): + """ + Parameters + ---------- + bbox : `Bbox` + transform : `Transform` + """ + _api.check_isinstance(BboxBase, bbox=bbox) + _api.check_isinstance(Transform, transform=transform) + if transform.input_dims != 2 or transform.output_dims != 2: + raise ValueError( + "The input and output dimensions of 'transform' must be 2") + + super().__init__(**kwargs) + self._bbox = bbox + self._transform = transform + self.set_children(bbox, transform) + self._points = None + + __str__ = _make_str_method("_bbox", "_transform") + + def get_points(self): + # docstring inherited + if self._invalid: + p = self._bbox.get_points() + # Transform all four points, then make a new bounding box + # from the result, taking care to make the orientation the + # same. + points = self._transform.transform( + [[p[0, 0], p[0, 1]], + [p[1, 0], p[0, 1]], + [p[0, 0], p[1, 1]], + [p[1, 0], p[1, 1]]]) + points = np.ma.filled(points, 0.0) + + xs = min(points[:, 0]), max(points[:, 0]) + if p[0, 0] > p[1, 0]: + xs = xs[::-1] + + ys = min(points[:, 1]), max(points[:, 1]) + if p[0, 1] > p[1, 1]: + ys = ys[::-1] + + self._points = np.array([ + [xs[0], ys[0]], + [xs[1], ys[1]] + ]) + + self._invalid = 0 + return self._points + + if DEBUG: + _get_points = get_points + + def get_points(self): + points = self._get_points() + self._check(points) + return points + + def contains(self, x, y): + # Docstring inherited. + return self._bbox.contains(*self._transform.inverted().transform((x, y))) + + def fully_contains(self, x, y): + # Docstring inherited. + return self._bbox.fully_contains(*self._transform.inverted().transform((x, y))) + + +class LockableBbox(BboxBase): + """ + A `Bbox` where some elements may be locked at certain values. + + When the child bounding box changes, the bounds of this bbox will update + accordingly with the exception of the locked elements. + """ + def __init__(self, bbox, x0=None, y0=None, x1=None, y1=None, **kwargs): + """ + Parameters + ---------- + bbox : `Bbox` + The child bounding box to wrap. + + x0 : float or None + The locked value for x0, or None to leave unlocked. + + y0 : float or None + The locked value for y0, or None to leave unlocked. + + x1 : float or None + The locked value for x1, or None to leave unlocked. + + y1 : float or None + The locked value for y1, or None to leave unlocked. + + """ + _api.check_isinstance(BboxBase, bbox=bbox) + super().__init__(**kwargs) + self._bbox = bbox + self.set_children(bbox) + self._points = None + fp = [x0, y0, x1, y1] + mask = [val is None for val in fp] + self._locked_points = np.ma.array(fp, float, mask=mask).reshape((2, 2)) + + __str__ = _make_str_method("_bbox", "_locked_points") + + def get_points(self): + # docstring inherited + if self._invalid: + points = self._bbox.get_points() + self._points = np.where(self._locked_points.mask, + points, + self._locked_points) + self._invalid = 0 + return self._points + + if DEBUG: + _get_points = get_points + + def get_points(self): + points = self._get_points() + self._check(points) + return points + + @property + def locked_x0(self): + """ + float or None: The value used for the locked x0. + """ + if self._locked_points.mask[0, 0]: + return None + else: + return self._locked_points[0, 0] + + @locked_x0.setter + def locked_x0(self, x0): + self._locked_points.mask[0, 0] = x0 is None + self._locked_points.data[0, 0] = x0 + self.invalidate() + + @property + def locked_y0(self): + """ + float or None: The value used for the locked y0. + """ + if self._locked_points.mask[0, 1]: + return None + else: + return self._locked_points[0, 1] + + @locked_y0.setter + def locked_y0(self, y0): + self._locked_points.mask[0, 1] = y0 is None + self._locked_points.data[0, 1] = y0 + self.invalidate() + + @property + def locked_x1(self): + """ + float or None: The value used for the locked x1. + """ + if self._locked_points.mask[1, 0]: + return None + else: + return self._locked_points[1, 0] + + @locked_x1.setter + def locked_x1(self, x1): + self._locked_points.mask[1, 0] = x1 is None + self._locked_points.data[1, 0] = x1 + self.invalidate() + + @property + def locked_y1(self): + """ + float or None: The value used for the locked y1. + """ + if self._locked_points.mask[1, 1]: + return None + else: + return self._locked_points[1, 1] + + @locked_y1.setter + def locked_y1(self, y1): + self._locked_points.mask[1, 1] = y1 is None + self._locked_points.data[1, 1] = y1 + self.invalidate() + + +class Transform(TransformNode): + """ + The base class of all `TransformNode` instances that + actually perform a transformation. + + All non-affine transformations should be subclasses of this class. + New affine transformations should be subclasses of `Affine2D`. + + Subclasses of this class should override the following members (at + minimum): + + - :attr:`input_dims` + - :attr:`output_dims` + - :meth:`transform` + - :meth:`inverted` (if an inverse exists) + + The following attributes may be overridden if the default is unsuitable: + + - :attr:`is_separable` (defaults to True for 1D -> 1D transforms, False + otherwise) + - :attr:`has_inverse` (defaults to True if :meth:`inverted` is overridden, + False otherwise) + + If the transform needs to do something non-standard with + `matplotlib.path.Path` objects, such as adding curves + where there were once line segments, it should override: + + - :meth:`transform_path` + """ + + input_dims = None + """ + The number of input dimensions of this transform. + Must be overridden (with integers) in the subclass. + """ + + output_dims = None + """ + The number of output dimensions of this transform. + Must be overridden (with integers) in the subclass. + """ + + is_separable = False + """True if this transform is separable in the x- and y- dimensions.""" + + has_inverse = False + """True if this transform has a corresponding inverse transform.""" + + def __init_subclass__(cls): + # 1d transforms are always separable; we assume higher-dimensional ones + # are not but subclasses can also directly set is_separable -- this is + # verified by checking whether "is_separable" appears more than once in + # the class's MRO (it appears once in Transform). + if (sum("is_separable" in vars(parent) for parent in cls.__mro__) == 1 + and cls.input_dims == cls.output_dims == 1): + cls.is_separable = True + # Transform.inverted raises NotImplementedError; we assume that if this + # is overridden then the transform is invertible but subclass can also + # directly set has_inverse. + if (sum("has_inverse" in vars(parent) for parent in cls.__mro__) == 1 + and hasattr(cls, "inverted") + and cls.inverted is not Transform.inverted): + cls.has_inverse = True + + def __add__(self, other): + """ + Compose two transforms together so that *self* is followed by *other*. + + ``A + B`` returns a transform ``C`` so that + ``C.transform(x) == B.transform(A.transform(x))``. + """ + return (composite_transform_factory(self, other) + if isinstance(other, Transform) else + NotImplemented) + + # Equality is based on object identity for `Transform`s (so we don't + # override `__eq__`), but some subclasses, such as TransformWrapper & + # AffineBase, override this behavior. + + def _iter_break_from_left_to_right(self): + """ + Return an iterator breaking down this transform stack from left to + right recursively. If self == ((A, N), A) then the result will be an + iterator which yields I : ((A, N), A), followed by A : (N, A), + followed by (A, N) : (A), but not ((A, N), A) : I. + + This is equivalent to flattening the stack then yielding + ``flat_stack[:i], flat_stack[i:]`` where i=0..(n-1). + """ + yield IdentityTransform(), self + + @property + def depth(self): + """ + Return the number of transforms which have been chained + together to form this Transform instance. + + .. note:: + + For the special case of a Composite transform, the maximum depth + of the two is returned. + + """ + return 1 + + def contains_branch(self, other): + """ + Return whether the given transform is a sub-tree of this transform. + + This routine uses transform equality to identify sub-trees, therefore + in many situations it is object id which will be used. + + For the case where the given transform represents the whole + of this transform, returns True. + """ + if self.depth < other.depth: + return False + + # check that a subtree is equal to other (starting from self) + for _, sub_tree in self._iter_break_from_left_to_right(): + if sub_tree == other: + return True + return False + + def contains_branch_seperately(self, other_transform): + """ + Return whether the given branch is a sub-tree of this transform on + each separate dimension. + + A common use for this method is to identify if a transform is a blended + transform containing an Axes' data transform. e.g.:: + + x_isdata, y_isdata = trans.contains_branch_seperately(ax.transData) + + """ + if self.output_dims != 2: + raise ValueError('contains_branch_seperately only supports ' + 'transforms with 2 output dimensions') + # for a non-blended transform each separate dimension is the same, so + # just return the appropriate shape. + return (self.contains_branch(other_transform), ) * 2 + + def __sub__(self, other): + """ + Compose *self* with the inverse of *other*, cancelling identical terms + if any:: + + # In general: + A - B == A + B.inverted() + # (but see note regarding frozen transforms below). + + # If A "ends with" B (i.e. A == A' + B for some A') we can cancel + # out B: + (A' + B) - B == A' + + # Likewise, if B "starts with" A (B = A + B'), we can cancel out A: + A - (A + B') == B'.inverted() == B'^-1 + + Cancellation (rather than naively returning ``A + B.inverted()``) is + important for multiple reasons: + + - It avoids floating-point inaccuracies when computing the inverse of + B: ``B - B`` is guaranteed to cancel out exactly (resulting in the + identity transform), whereas ``B + B.inverted()`` may differ by a + small epsilon. + - ``B.inverted()`` always returns a frozen transform: if one computes + ``A + B + B.inverted()`` and later mutates ``B``, then + ``B.inverted()`` won't be updated and the last two terms won't cancel + out anymore; on the other hand, ``A + B - B`` will always be equal to + ``A`` even if ``B`` is mutated. + """ + # we only know how to do this operation if other is a Transform. + if not isinstance(other, Transform): + return NotImplemented + for remainder, sub_tree in self._iter_break_from_left_to_right(): + if sub_tree == other: + return remainder + for remainder, sub_tree in other._iter_break_from_left_to_right(): + if sub_tree == self: + if not remainder.has_inverse: + raise ValueError( + "The shortcut cannot be computed since 'other' " + "includes a non-invertible component") + return remainder.inverted() + # if we have got this far, then there was no shortcut possible + if other.has_inverse: + return self + other.inverted() + else: + raise ValueError('It is not possible to compute transA - transB ' + 'since transB cannot be inverted and there is no ' + 'shortcut possible.') + + def __array__(self, *args, **kwargs): + """Array interface to get at this Transform's affine matrix.""" + return self.get_affine().get_matrix() + + def transform(self, values): + """ + Apply this transformation on the given array of *values*. + + Parameters + ---------- + values : array-like + The input values as an array of length :attr:`input_dims` or + shape (N, :attr:`input_dims`). + + Returns + ------- + array + The output values as an array of length :attr:`output_dims` or + shape (N, :attr:`output_dims`), depending on the input. + """ + # Ensure that values is a 2d array (but remember whether + # we started with a 1d or 2d array). + values = np.asanyarray(values) + ndim = values.ndim + values = values.reshape((-1, self.input_dims)) + + # Transform the values + res = self.transform_affine(self.transform_non_affine(values)) + + # Convert the result back to the shape of the input values. + if ndim == 0: + assert not np.ma.is_masked(res) # just to be on the safe side + return res[0, 0] + if ndim == 1: + return res.reshape(-1) + elif ndim == 2: + return res + raise ValueError( + "Input values must have shape (N, {dims}) or ({dims},)" + .format(dims=self.input_dims)) + + def transform_affine(self, values): + """ + Apply only the affine part of this transformation on the + given array of values. + + ``transform(values)`` is always equivalent to + ``transform_affine(transform_non_affine(values))``. + + In non-affine transformations, this is generally a no-op. In + affine transformations, this is equivalent to + ``transform(values)``. + + Parameters + ---------- + values : array + The input values as an array of length :attr:`input_dims` or + shape (N, :attr:`input_dims`). + + Returns + ------- + array + The output values as an array of length :attr:`output_dims` or + shape (N, :attr:`output_dims`), depending on the input. + """ + return self.get_affine().transform(values) + + def transform_non_affine(self, values): + """ + Apply only the non-affine part of this transformation. + + ``transform(values)`` is always equivalent to + ``transform_affine(transform_non_affine(values))``. + + In non-affine transformations, this is generally equivalent to + ``transform(values)``. In affine transformations, this is + always a no-op. + + Parameters + ---------- + values : array + The input values as an array of length :attr:`input_dims` or + shape (N, :attr:`input_dims`). + + Returns + ------- + array + The output values as an array of length :attr:`output_dims` or + shape (N, :attr:`output_dims`), depending on the input. + """ + return values + + def transform_bbox(self, bbox): + """ + Transform the given bounding box. + + For smarter transforms including caching (a common requirement in + Matplotlib), see `TransformedBbox`. + """ + return Bbox(self.transform(bbox.get_points())) + + def get_affine(self): + """Get the affine part of this transform.""" + return IdentityTransform() + + def get_matrix(self): + """Get the matrix for the affine part of this transform.""" + return self.get_affine().get_matrix() + + def transform_point(self, point): + """ + Return a transformed point. + + This function is only kept for backcompatibility; the more general + `.transform` method is capable of transforming both a list of points + and a single point. + + The point is given as a sequence of length :attr:`input_dims`. + The transformed point is returned as a sequence of length + :attr:`output_dims`. + """ + if len(point) != self.input_dims: + raise ValueError("The length of 'point' must be 'self.input_dims'") + return self.transform(point) + + def transform_path(self, path): + """ + Apply the transform to `.Path` *path*, returning a new `.Path`. + + In some cases, this transform may insert curves into the path + that began as line segments. + """ + return self.transform_path_affine(self.transform_path_non_affine(path)) + + def transform_path_affine(self, path): + """ + Apply the affine part of this transform to `.Path` *path*, returning a + new `.Path`. + + ``transform_path(path)`` is equivalent to + ``transform_path_affine(transform_path_non_affine(values))``. + """ + return self.get_affine().transform_path_affine(path) + + def transform_path_non_affine(self, path): + """ + Apply the non-affine part of this transform to `.Path` *path*, + returning a new `.Path`. + + ``transform_path(path)`` is equivalent to + ``transform_path_affine(transform_path_non_affine(values))``. + """ + x = self.transform_non_affine(path.vertices) + return Path._fast_from_codes_and_verts(x, path.codes, path) + + def transform_angles(self, angles, pts, radians=False, pushoff=1e-5): + """ + Transform a set of angles anchored at specific locations. + + Parameters + ---------- + angles : (N,) array-like + The angles to transform. + pts : (N, 2) array-like + The points where the angles are anchored. + radians : bool, default: False + Whether *angles* are radians or degrees. + pushoff : float + For each point in *pts* and angle in *angles*, the transformed + angle is computed by transforming a segment of length *pushoff* + starting at that point and making that angle relative to the + horizontal axis, and measuring the angle between the horizontal + axis and the transformed segment. + + Returns + ------- + (N,) array + """ + # Must be 2D + if self.input_dims != 2 or self.output_dims != 2: + raise NotImplementedError('Only defined in 2D') + angles = np.asarray(angles) + pts = np.asarray(pts) + _api.check_shape((None, 2), pts=pts) + _api.check_shape((None,), angles=angles) + if len(angles) != len(pts): + raise ValueError("There must be as many 'angles' as 'pts'") + # Convert to radians if desired + if not radians: + angles = np.deg2rad(angles) + # Move a short distance away + pts2 = pts + pushoff * np.column_stack([np.cos(angles), + np.sin(angles)]) + # Transform both sets of points + tpts = self.transform(pts) + tpts2 = self.transform(pts2) + # Calculate transformed angles + d = tpts2 - tpts + a = np.arctan2(d[:, 1], d[:, 0]) + # Convert back to degrees if desired + if not radians: + a = np.rad2deg(a) + return a + + def inverted(self): + """ + Return the corresponding inverse transformation. + + It holds ``x == self.inverted().transform(self.transform(x))``. + + The return value of this method should be treated as + temporary. An update to *self* does not cause a corresponding + update to its inverted copy. + """ + raise NotImplementedError() + + +class TransformWrapper(Transform): + """ + A helper class that holds a single child transform and acts + equivalently to it. + + This is useful if a node of the transform tree must be replaced at + run time with a transform of a different type. This class allows + that replacement to correctly trigger invalidation. + + `TransformWrapper` instances must have the same input and output dimensions + during their entire lifetime, so the child transform may only be replaced + with another child transform of the same dimensions. + """ + + pass_through = True + + def __init__(self, child): + """ + *child*: A `Transform` instance. This child may later + be replaced with :meth:`set`. + """ + _api.check_isinstance(Transform, child=child) + super().__init__() + self.set(child) + + def __eq__(self, other): + return self._child.__eq__(other) + + __str__ = _make_str_method("_child") + + def frozen(self): + # docstring inherited + return self._child.frozen() + + def set(self, child): + """ + Replace the current child of this transform with another one. + + The new child must have the same number of input and output + dimensions as the current child. + """ + if hasattr(self, "_child"): # Absent during init. + self.invalidate() + new_dims = (child.input_dims, child.output_dims) + old_dims = (self._child.input_dims, self._child.output_dims) + if new_dims != old_dims: + raise ValueError( + f"The input and output dims of the new child {new_dims} " + f"do not match those of current child {old_dims}") + self._child._parents.pop(id(self), None) + + self._child = child + self.set_children(child) + + self.transform = child.transform + self.transform_affine = child.transform_affine + self.transform_non_affine = child.transform_non_affine + self.transform_path = child.transform_path + self.transform_path_affine = child.transform_path_affine + self.transform_path_non_affine = child.transform_path_non_affine + self.get_affine = child.get_affine + self.inverted = child.inverted + self.get_matrix = child.get_matrix + # note we do not wrap other properties here since the transform's + # child can be changed with WrappedTransform.set and so checking + # is_affine and other such properties may be dangerous. + + self._invalid = 0 + self.invalidate() + self._invalid = 0 + + input_dims = property(lambda self: self._child.input_dims) + output_dims = property(lambda self: self._child.output_dims) + is_affine = property(lambda self: self._child.is_affine) + is_separable = property(lambda self: self._child.is_separable) + has_inverse = property(lambda self: self._child.has_inverse) + + +class AffineBase(Transform): + """ + The base class of all affine transformations of any number of dimensions. + """ + is_affine = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._inverted = None + + def __array__(self, *args, **kwargs): + # optimises the access of the transform matrix vs. the superclass + return self.get_matrix() + + def __eq__(self, other): + if getattr(other, "is_affine", False) and hasattr(other, "get_matrix"): + return (self.get_matrix() == other.get_matrix()).all() + return NotImplemented + + def transform(self, values): + # docstring inherited + return self.transform_affine(values) + + def transform_affine(self, values): + # docstring inherited + raise NotImplementedError('Affine subclasses should override this ' + 'method.') + + def transform_non_affine(self, values): + # docstring inherited + return values + + def transform_path(self, path): + # docstring inherited + return self.transform_path_affine(path) + + def transform_path_affine(self, path): + # docstring inherited + return Path(self.transform_affine(path.vertices), + path.codes, path._interpolation_steps) + + def transform_path_non_affine(self, path): + # docstring inherited + return path + + def get_affine(self): + # docstring inherited + return self + + +class Affine2DBase(AffineBase): + """ + The base class of all 2D affine transformations. + + 2D affine transformations are performed using a 3x3 numpy array:: + + a c e + b d f + 0 0 1 + + This class provides the read-only interface. For a mutable 2D + affine transformation, use `Affine2D`. + + Subclasses of this class will generally only need to override a + constructor and `~.Transform.get_matrix` that generates a custom 3x3 matrix. + """ + input_dims = 2 + output_dims = 2 + + def frozen(self): + # docstring inherited + return Affine2D(self.get_matrix().copy()) + + @property + def is_separable(self): + mtx = self.get_matrix() + return mtx[0, 1] == mtx[1, 0] == 0.0 + + def to_values(self): + """ + Return the values of the matrix as an ``(a, b, c, d, e, f)`` tuple. + """ + mtx = self.get_matrix() + return tuple(mtx[:2].swapaxes(0, 1).flat) + + def transform_affine(self, values): + mtx = self.get_matrix() + if isinstance(values, np.ma.MaskedArray): + tpoints = affine_transform(values.data, mtx) + return np.ma.MaskedArray(tpoints, mask=np.ma.getmask(values)) + return affine_transform(values, mtx) + + if DEBUG: + _transform_affine = transform_affine + + def transform_affine(self, values): + # docstring inherited + # The major speed trap here is just converting to the + # points to an array in the first place. If we can use + # more arrays upstream, that should help here. + if not isinstance(values, np.ndarray): + _api.warn_external( + f'A non-numpy array of type {type(values)} was passed in ' + f'for transformation, which results in poor performance.') + return self._transform_affine(values) + + def inverted(self): + # docstring inherited + if self._inverted is None or self._invalid: + mtx = self.get_matrix() + shorthand_name = None + if self._shorthand_name: + shorthand_name = '(%s)-1' % self._shorthand_name + self._inverted = Affine2D(inv(mtx), shorthand_name=shorthand_name) + self._invalid = 0 + return self._inverted + + +class Affine2D(Affine2DBase): + """ + A mutable 2D affine transformation. + """ + + def __init__(self, matrix=None, **kwargs): + """ + Initialize an Affine transform from a 3x3 numpy float array:: + + a c e + b d f + 0 0 1 + + If *matrix* is None, initialize with the identity transform. + """ + super().__init__(**kwargs) + if matrix is None: + # A bit faster than np.identity(3). + matrix = IdentityTransform._mtx + self._mtx = matrix.copy() + self._invalid = 0 + + _base_str = _make_str_method("_mtx") + + def __str__(self): + return (self._base_str() + if (self._mtx != np.diag(np.diag(self._mtx))).any() + else f"Affine2D().scale({self._mtx[0, 0]}, {self._mtx[1, 1]})" + if self._mtx[0, 0] != self._mtx[1, 1] + else f"Affine2D().scale({self._mtx[0, 0]})") + + @staticmethod + def from_values(a, b, c, d, e, f): + """ + Create a new Affine2D instance from the given values:: + + a c e + b d f + 0 0 1 + + . + """ + return Affine2D( + np.array([a, c, e, b, d, f, 0.0, 0.0, 1.0], float).reshape((3, 3))) + + def get_matrix(self): + """ + Get the underlying transformation matrix as a 3x3 array:: + + a c e + b d f + 0 0 1 + + . + """ + if self._invalid: + self._inverted = None + self._invalid = 0 + return self._mtx + + def set_matrix(self, mtx): + """ + Set the underlying transformation matrix from a 3x3 array:: + + a c e + b d f + 0 0 1 + + . + """ + self._mtx = mtx + self.invalidate() + + def set(self, other): + """ + Set this transformation from the frozen copy of another + `Affine2DBase` object. + """ + _api.check_isinstance(Affine2DBase, other=other) + self._mtx = other.get_matrix() + self.invalidate() + + def clear(self): + """ + Reset the underlying matrix to the identity transform. + """ + # A bit faster than np.identity(3). + self._mtx = IdentityTransform._mtx.copy() + self.invalidate() + return self + + def rotate(self, theta): + """ + Add a rotation (in radians) to this transform in place. + + Returns *self*, so this method can easily be chained with more + calls to :meth:`rotate`, :meth:`rotate_deg`, :meth:`translate` + and :meth:`scale`. + """ + a = math.cos(theta) + b = math.sin(theta) + mtx = self._mtx + # Operating and assigning one scalar at a time is much faster. + (xx, xy, x0), (yx, yy, y0), _ = mtx.tolist() + # mtx = [[a -b 0], [b a 0], [0 0 1]] * mtx + mtx[0, 0] = a * xx - b * yx + mtx[0, 1] = a * xy - b * yy + mtx[0, 2] = a * x0 - b * y0 + mtx[1, 0] = b * xx + a * yx + mtx[1, 1] = b * xy + a * yy + mtx[1, 2] = b * x0 + a * y0 + self.invalidate() + return self + + def rotate_deg(self, degrees): + """ + Add a rotation (in degrees) to this transform in place. + + Returns *self*, so this method can easily be chained with more + calls to :meth:`rotate`, :meth:`rotate_deg`, :meth:`translate` + and :meth:`scale`. + """ + return self.rotate(math.radians(degrees)) + + def rotate_around(self, x, y, theta): + """ + Add a rotation (in radians) around the point (x, y) in place. + + Returns *self*, so this method can easily be chained with more + calls to :meth:`rotate`, :meth:`rotate_deg`, :meth:`translate` + and :meth:`scale`. + """ + return self.translate(-x, -y).rotate(theta).translate(x, y) + + def rotate_deg_around(self, x, y, degrees): + """ + Add a rotation (in degrees) around the point (x, y) in place. + + Returns *self*, so this method can easily be chained with more + calls to :meth:`rotate`, :meth:`rotate_deg`, :meth:`translate` + and :meth:`scale`. + """ + # Cast to float to avoid wraparound issues with uint8's + x, y = float(x), float(y) + return self.translate(-x, -y).rotate_deg(degrees).translate(x, y) + + def translate(self, tx, ty): + """ + Add a translation in place. + + Returns *self*, so this method can easily be chained with more + calls to :meth:`rotate`, :meth:`rotate_deg`, :meth:`translate` + and :meth:`scale`. + """ + self._mtx[0, 2] += tx + self._mtx[1, 2] += ty + self.invalidate() + return self + + def scale(self, sx, sy=None): + """ + Add a scale in place. + + If *sy* is None, the same scale is applied in both the *x*- and + *y*-directions. + + Returns *self*, so this method can easily be chained with more + calls to :meth:`rotate`, :meth:`rotate_deg`, :meth:`translate` + and :meth:`scale`. + """ + if sy is None: + sy = sx + # explicit element-wise scaling is fastest + self._mtx[0, 0] *= sx + self._mtx[0, 1] *= sx + self._mtx[0, 2] *= sx + self._mtx[1, 0] *= sy + self._mtx[1, 1] *= sy + self._mtx[1, 2] *= sy + self.invalidate() + return self + + def skew(self, xShear, yShear): + """ + Add a skew in place. + + *xShear* and *yShear* are the shear angles along the *x*- and + *y*-axes, respectively, in radians. + + Returns *self*, so this method can easily be chained with more + calls to :meth:`rotate`, :meth:`rotate_deg`, :meth:`translate` + and :meth:`scale`. + """ + rx = math.tan(xShear) + ry = math.tan(yShear) + mtx = self._mtx + # Operating and assigning one scalar at a time is much faster. + (xx, xy, x0), (yx, yy, y0), _ = mtx.tolist() + # mtx = [[1 rx 0], [ry 1 0], [0 0 1]] * mtx + mtx[0, 0] += rx * yx + mtx[0, 1] += rx * yy + mtx[0, 2] += rx * y0 + mtx[1, 0] += ry * xx + mtx[1, 1] += ry * xy + mtx[1, 2] += ry * x0 + self.invalidate() + return self + + def skew_deg(self, xShear, yShear): + """ + Add a skew in place. + + *xShear* and *yShear* are the shear angles along the *x*- and + *y*-axes, respectively, in degrees. + + Returns *self*, so this method can easily be chained with more + calls to :meth:`rotate`, :meth:`rotate_deg`, :meth:`translate` + and :meth:`scale`. + """ + return self.skew(math.radians(xShear), math.radians(yShear)) + + +class IdentityTransform(Affine2DBase): + """ + A special class that does one thing, the identity transform, in a + fast way. + """ + _mtx = np.identity(3) + + def frozen(self): + # docstring inherited + return self + + __str__ = _make_str_method() + + def get_matrix(self): + # docstring inherited + return self._mtx + + def transform(self, values): + # docstring inherited + return np.asanyarray(values) + + def transform_affine(self, values): + # docstring inherited + return np.asanyarray(values) + + def transform_non_affine(self, values): + # docstring inherited + return np.asanyarray(values) + + def transform_path(self, path): + # docstring inherited + return path + + def transform_path_affine(self, path): + # docstring inherited + return path + + def transform_path_non_affine(self, path): + # docstring inherited + return path + + def get_affine(self): + # docstring inherited + return self + + def inverted(self): + # docstring inherited + return self + + +class _BlendedMixin: + """Common methods for `BlendedGenericTransform` and `BlendedAffine2D`.""" + + def __eq__(self, other): + if isinstance(other, (BlendedAffine2D, BlendedGenericTransform)): + return (self._x == other._x) and (self._y == other._y) + elif self._x == self._y: + return self._x == other + else: + return NotImplemented + + def contains_branch_seperately(self, transform): + return (self._x.contains_branch(transform), + self._y.contains_branch(transform)) + + __str__ = _make_str_method("_x", "_y") + + +class BlendedGenericTransform(_BlendedMixin, Transform): + """ + A "blended" transform uses one transform for the *x*-direction, and + another transform for the *y*-direction. + + This "generic" version can handle any given child transform in the + *x*- and *y*-directions. + """ + input_dims = 2 + output_dims = 2 + is_separable = True + pass_through = True + + def __init__(self, x_transform, y_transform, **kwargs): + """ + Create a new "blended" transform using *x_transform* to transform the + *x*-axis and *y_transform* to transform the *y*-axis. + + You will generally not call this constructor directly but use the + `blended_transform_factory` function instead, which can determine + automatically which kind of blended transform to create. + """ + Transform.__init__(self, **kwargs) + self._x = x_transform + self._y = y_transform + self.set_children(x_transform, y_transform) + self._affine = None + + @property + def depth(self): + return max(self._x.depth, self._y.depth) + + def contains_branch(self, other): + # A blended transform cannot possibly contain a branch from two + # different transforms. + return False + + is_affine = property(lambda self: self._x.is_affine and self._y.is_affine) + has_inverse = property( + lambda self: self._x.has_inverse and self._y.has_inverse) + + def frozen(self): + # docstring inherited + return blended_transform_factory(self._x.frozen(), self._y.frozen()) + + def transform_non_affine(self, values): + # docstring inherited + if self._x.is_affine and self._y.is_affine: + return values + x = self._x + y = self._y + + if x == y and x.input_dims == 2: + return x.transform_non_affine(values) + + if x.input_dims == 2: + x_points = x.transform_non_affine(values)[:, 0:1] + else: + x_points = x.transform_non_affine(values[:, 0]) + x_points = x_points.reshape((len(x_points), 1)) + + if y.input_dims == 2: + y_points = y.transform_non_affine(values)[:, 1:] + else: + y_points = y.transform_non_affine(values[:, 1]) + y_points = y_points.reshape((len(y_points), 1)) + + if (isinstance(x_points, np.ma.MaskedArray) or + isinstance(y_points, np.ma.MaskedArray)): + return np.ma.concatenate((x_points, y_points), 1) + else: + return np.concatenate((x_points, y_points), 1) + + def inverted(self): + # docstring inherited + return BlendedGenericTransform(self._x.inverted(), self._y.inverted()) + + def get_affine(self): + # docstring inherited + if self._invalid or self._affine is None: + if self._x == self._y: + self._affine = self._x.get_affine() + else: + x_mtx = self._x.get_affine().get_matrix() + y_mtx = self._y.get_affine().get_matrix() + # We already know the transforms are separable, so we can skip + # setting b and c to zero. + mtx = np.array([x_mtx[0], y_mtx[1], [0.0, 0.0, 1.0]]) + self._affine = Affine2D(mtx) + self._invalid = 0 + return self._affine + + +class BlendedAffine2D(_BlendedMixin, Affine2DBase): + """ + A "blended" transform uses one transform for the *x*-direction, and + another transform for the *y*-direction. + + This version is an optimization for the case where both child + transforms are of type `Affine2DBase`. + """ + + is_separable = True + + def __init__(self, x_transform, y_transform, **kwargs): + """ + Create a new "blended" transform using *x_transform* to transform the + *x*-axis and *y_transform* to transform the *y*-axis. + + Both *x_transform* and *y_transform* must be 2D affine transforms. + + You will generally not call this constructor directly but use the + `blended_transform_factory` function instead, which can determine + automatically which kind of blended transform to create. + """ + is_affine = x_transform.is_affine and y_transform.is_affine + is_separable = x_transform.is_separable and y_transform.is_separable + is_correct = is_affine and is_separable + if not is_correct: + raise ValueError("Both *x_transform* and *y_transform* must be 2D " + "affine transforms") + + Transform.__init__(self, **kwargs) + self._x = x_transform + self._y = y_transform + self.set_children(x_transform, y_transform) + + Affine2DBase.__init__(self) + self._mtx = None + + def get_matrix(self): + # docstring inherited + if self._invalid: + if self._x == self._y: + self._mtx = self._x.get_matrix() + else: + x_mtx = self._x.get_matrix() + y_mtx = self._y.get_matrix() + # We already know the transforms are separable, so we can skip + # setting b and c to zero. + self._mtx = np.array([x_mtx[0], y_mtx[1], [0.0, 0.0, 1.0]]) + self._inverted = None + self._invalid = 0 + return self._mtx + + +def blended_transform_factory(x_transform, y_transform): + """ + Create a new "blended" transform using *x_transform* to transform + the *x*-axis and *y_transform* to transform the *y*-axis. + + A faster version of the blended transform is returned for the case + where both child transforms are affine. + """ + if (isinstance(x_transform, Affine2DBase) and + isinstance(y_transform, Affine2DBase)): + return BlendedAffine2D(x_transform, y_transform) + return BlendedGenericTransform(x_transform, y_transform) + + +class CompositeGenericTransform(Transform): + """ + A composite transform formed by applying transform *a* then + transform *b*. + + This "generic" version can handle any two arbitrary + transformations. + """ + pass_through = True + + def __init__(self, a, b, **kwargs): + """ + Create a new composite transform that is the result of + applying transform *a* then transform *b*. + + You will generally not call this constructor directly but write ``a + + b`` instead, which will automatically choose the best kind of composite + transform instance to create. + """ + if a.output_dims != b.input_dims: + raise ValueError("The output dimension of 'a' must be equal to " + "the input dimensions of 'b'") + self.input_dims = a.input_dims + self.output_dims = b.output_dims + + super().__init__(**kwargs) + self._a = a + self._b = b + self.set_children(a, b) + + def frozen(self): + # docstring inherited + self._invalid = 0 + frozen = composite_transform_factory( + self._a.frozen(), self._b.frozen()) + if not isinstance(frozen, CompositeGenericTransform): + return frozen.frozen() + return frozen + + def _invalidate_internal(self, level, invalidating_node): + # When the left child is invalidated at AFFINE_ONLY level and the right child is + # non-affine, the composite transform is FULLY invalidated. + if invalidating_node is self._a and not self._b.is_affine: + level = Transform._INVALID_FULL + super()._invalidate_internal(level, invalidating_node) + + def __eq__(self, other): + if isinstance(other, (CompositeGenericTransform, CompositeAffine2D)): + return self is other or (self._a == other._a + and self._b == other._b) + else: + return False + + def _iter_break_from_left_to_right(self): + for left, right in self._a._iter_break_from_left_to_right(): + yield left, right + self._b + for left, right in self._b._iter_break_from_left_to_right(): + yield self._a + left, right + + def contains_branch_seperately(self, other_transform): + # docstring inherited + if self.output_dims != 2: + raise ValueError('contains_branch_seperately only supports ' + 'transforms with 2 output dimensions') + if self == other_transform: + return (True, True) + return self._b.contains_branch_seperately(other_transform) + + depth = property(lambda self: self._a.depth + self._b.depth) + is_affine = property(lambda self: self._a.is_affine and self._b.is_affine) + is_separable = property( + lambda self: self._a.is_separable and self._b.is_separable) + has_inverse = property( + lambda self: self._a.has_inverse and self._b.has_inverse) + + __str__ = _make_str_method("_a", "_b") + + def transform_affine(self, values): + # docstring inherited + return self.get_affine().transform(values) + + def transform_non_affine(self, values): + # docstring inherited + if self._a.is_affine and self._b.is_affine: + return values + elif not self._a.is_affine and self._b.is_affine: + return self._a.transform_non_affine(values) + else: + return self._b.transform_non_affine(self._a.transform(values)) + + def transform_path_non_affine(self, path): + # docstring inherited + if self._a.is_affine and self._b.is_affine: + return path + elif not self._a.is_affine and self._b.is_affine: + return self._a.transform_path_non_affine(path) + else: + return self._b.transform_path_non_affine( + self._a.transform_path(path)) + + def get_affine(self): + # docstring inherited + if not self._b.is_affine: + return self._b.get_affine() + else: + return Affine2D(np.dot(self._b.get_affine().get_matrix(), + self._a.get_affine().get_matrix())) + + def inverted(self): + # docstring inherited + return CompositeGenericTransform( + self._b.inverted(), self._a.inverted()) + + +class CompositeAffine2D(Affine2DBase): + """ + A composite transform formed by applying transform *a* then transform *b*. + + This version is an optimization that handles the case where both *a* + and *b* are 2D affines. + """ + def __init__(self, a, b, **kwargs): + """ + Create a new composite transform that is the result of + applying `Affine2DBase` *a* then `Affine2DBase` *b*. + + You will generally not call this constructor directly but write ``a + + b`` instead, which will automatically choose the best kind of composite + transform instance to create. + """ + if not a.is_affine or not b.is_affine: + raise ValueError("'a' and 'b' must be affine transforms") + if a.output_dims != b.input_dims: + raise ValueError("The output dimension of 'a' must be equal to " + "the input dimensions of 'b'") + self.input_dims = a.input_dims + self.output_dims = b.output_dims + + super().__init__(**kwargs) + self._a = a + self._b = b + self.set_children(a, b) + self._mtx = None + + @property + def depth(self): + return self._a.depth + self._b.depth + + def _iter_break_from_left_to_right(self): + for left, right in self._a._iter_break_from_left_to_right(): + yield left, right + self._b + for left, right in self._b._iter_break_from_left_to_right(): + yield self._a + left, right + + __str__ = _make_str_method("_a", "_b") + + def get_matrix(self): + # docstring inherited + if self._invalid: + self._mtx = np.dot( + self._b.get_matrix(), + self._a.get_matrix()) + self._inverted = None + self._invalid = 0 + return self._mtx + + +def composite_transform_factory(a, b): + """ + Create a new composite transform that is the result of applying + transform a then transform b. + + Shortcut versions of the blended transform are provided for the + case where both child transforms are affine, or one or the other + is the identity transform. + + Composite transforms may also be created using the '+' operator, + e.g.:: + + c = a + b + """ + # check to see if any of a or b are IdentityTransforms. We use + # isinstance here to guarantee that the transforms will *always* + # be IdentityTransforms. Since TransformWrappers are mutable, + # use of equality here would be wrong. + if isinstance(a, IdentityTransform): + return b + elif isinstance(b, IdentityTransform): + return a + elif isinstance(a, Affine2D) and isinstance(b, Affine2D): + return CompositeAffine2D(a, b) + return CompositeGenericTransform(a, b) + + +class BboxTransform(Affine2DBase): + """ + `BboxTransform` linearly transforms points from one `Bbox` to another. + """ + + is_separable = True + + def __init__(self, boxin, boxout, **kwargs): + """ + Create a new `BboxTransform` that linearly transforms + points from *boxin* to *boxout*. + """ + _api.check_isinstance(BboxBase, boxin=boxin, boxout=boxout) + + super().__init__(**kwargs) + self._boxin = boxin + self._boxout = boxout + self.set_children(boxin, boxout) + self._mtx = None + self._inverted = None + + __str__ = _make_str_method("_boxin", "_boxout") + + def get_matrix(self): + # docstring inherited + if self._invalid: + inl, inb, inw, inh = self._boxin.bounds + outl, outb, outw, outh = self._boxout.bounds + x_scale = outw / inw + y_scale = outh / inh + if DEBUG and (x_scale == 0 or y_scale == 0): + raise ValueError( + "Transforming from or to a singular bounding box") + self._mtx = np.array([[x_scale, 0.0, -inl*x_scale+outl], + [ 0.0, y_scale, -inb*y_scale+outb], + [ 0.0, 0.0, 1.0]], + float) + self._inverted = None + self._invalid = 0 + return self._mtx + + +class BboxTransformTo(Affine2DBase): + """ + `BboxTransformTo` is a transformation that linearly transforms points from + the unit bounding box to a given `Bbox`. + """ + + is_separable = True + + def __init__(self, boxout, **kwargs): + """ + Create a new `BboxTransformTo` that linearly transforms + points from the unit bounding box to *boxout*. + """ + _api.check_isinstance(BboxBase, boxout=boxout) + + super().__init__(**kwargs) + self._boxout = boxout + self.set_children(boxout) + self._mtx = None + self._inverted = None + + __str__ = _make_str_method("_boxout") + + def get_matrix(self): + # docstring inherited + if self._invalid: + outl, outb, outw, outh = self._boxout.bounds + if DEBUG and (outw == 0 or outh == 0): + raise ValueError("Transforming to a singular bounding box.") + self._mtx = np.array([[outw, 0.0, outl], + [ 0.0, outh, outb], + [ 0.0, 0.0, 1.0]], + float) + self._inverted = None + self._invalid = 0 + return self._mtx + + +@_api.deprecated("3.9") +class BboxTransformToMaxOnly(BboxTransformTo): + """ + `BboxTransformToMaxOnly` is a transformation that linearly transforms points from + the unit bounding box to a given `Bbox` with a fixed upper left of (0, 0). + """ + def get_matrix(self): + # docstring inherited + if self._invalid: + xmax, ymax = self._boxout.max + if DEBUG and (xmax == 0 or ymax == 0): + raise ValueError("Transforming to a singular bounding box.") + self._mtx = np.array([[xmax, 0.0, 0.0], + [ 0.0, ymax, 0.0], + [ 0.0, 0.0, 1.0]], + float) + self._inverted = None + self._invalid = 0 + return self._mtx + + +class BboxTransformFrom(Affine2DBase): + """ + `BboxTransformFrom` linearly transforms points from a given `Bbox` to the + unit bounding box. + """ + is_separable = True + + def __init__(self, boxin, **kwargs): + _api.check_isinstance(BboxBase, boxin=boxin) + + super().__init__(**kwargs) + self._boxin = boxin + self.set_children(boxin) + self._mtx = None + self._inverted = None + + __str__ = _make_str_method("_boxin") + + def get_matrix(self): + # docstring inherited + if self._invalid: + inl, inb, inw, inh = self._boxin.bounds + if DEBUG and (inw == 0 or inh == 0): + raise ValueError("Transforming from a singular bounding box.") + x_scale = 1.0 / inw + y_scale = 1.0 / inh + self._mtx = np.array([[x_scale, 0.0, -inl*x_scale], + [ 0.0, y_scale, -inb*y_scale], + [ 0.0, 0.0, 1.0]], + float) + self._inverted = None + self._invalid = 0 + return self._mtx + + +class ScaledTranslation(Affine2DBase): + """ + A transformation that translates by *xt* and *yt*, after *xt* and *yt* + have been transformed by *scale_trans*. + """ + def __init__(self, xt, yt, scale_trans, **kwargs): + super().__init__(**kwargs) + self._t = (xt, yt) + self._scale_trans = scale_trans + self.set_children(scale_trans) + self._mtx = None + self._inverted = None + + __str__ = _make_str_method("_t") + + def get_matrix(self): + # docstring inherited + if self._invalid: + # A bit faster than np.identity(3). + self._mtx = IdentityTransform._mtx.copy() + self._mtx[:2, 2] = self._scale_trans.transform(self._t) + self._invalid = 0 + self._inverted = None + return self._mtx + + +class _ScaledRotation(Affine2DBase): + """ + A transformation that applies rotation by *theta*, after transform by *trans_shift*. + """ + def __init__(self, theta, trans_shift): + super().__init__() + self._theta = theta + self._trans_shift = trans_shift + self._mtx = None + + def get_matrix(self): + if self._invalid: + transformed_coords = self._trans_shift.transform([[self._theta, 0]])[0] + adjusted_theta = transformed_coords[0] + rotation = Affine2D().rotate(adjusted_theta) + self._mtx = rotation.get_matrix() + return self._mtx + + +class AffineDeltaTransform(Affine2DBase): + r""" + A transform wrapper for transforming displacements between pairs of points. + + This class is intended to be used to transform displacements ("position + deltas") between pairs of points (e.g., as the ``offset_transform`` + of `.Collection`\s): given a transform ``t`` such that ``t = + AffineDeltaTransform(t) + offset``, ``AffineDeltaTransform`` + satisfies ``AffineDeltaTransform(a - b) == AffineDeltaTransform(a) - + AffineDeltaTransform(b)``. + + This is implemented by forcing the offset components of the transform + matrix to zero. + + This class is experimental as of 3.3, and the API may change. + """ + + pass_through = True + + def __init__(self, transform, **kwargs): + super().__init__(**kwargs) + self._base_transform = transform + self.set_children(transform) + + __str__ = _make_str_method("_base_transform") + + def get_matrix(self): + if self._invalid: + self._mtx = self._base_transform.get_matrix().copy() + self._mtx[:2, -1] = 0 + return self._mtx + + +class TransformedPath(TransformNode): + """ + A `TransformedPath` caches a non-affine transformed copy of the + `~.path.Path`. This cached copy is automatically updated when the + non-affine part of the transform changes. + + .. note:: + + Paths are considered immutable by this class. Any update to the + path's vertices/codes will not trigger a transform recomputation. + + """ + def __init__(self, path, transform): + """ + Parameters + ---------- + path : `~.path.Path` + transform : `Transform` + """ + _api.check_isinstance(Transform, transform=transform) + super().__init__() + self._path = path + self._transform = transform + self.set_children(transform) + self._transformed_path = None + self._transformed_points = None + + def _revalidate(self): + # only recompute if the invalidation includes the non_affine part of + # the transform + if (self._invalid == self._INVALID_FULL + or self._transformed_path is None): + self._transformed_path = \ + self._transform.transform_path_non_affine(self._path) + self._transformed_points = \ + Path._fast_from_codes_and_verts( + self._transform.transform_non_affine(self._path.vertices), + None, self._path) + self._invalid = 0 + + def get_transformed_points_and_affine(self): + """ + Return a copy of the child path, with the non-affine part of + the transform already applied, along with the affine part of + the path necessary to complete the transformation. Unlike + :meth:`get_transformed_path_and_affine`, no interpolation will + be performed. + """ + self._revalidate() + return self._transformed_points, self.get_affine() + + def get_transformed_path_and_affine(self): + """ + Return a copy of the child path, with the non-affine part of + the transform already applied, along with the affine part of + the path necessary to complete the transformation. + """ + self._revalidate() + return self._transformed_path, self.get_affine() + + def get_fully_transformed_path(self): + """ + Return a fully-transformed copy of the child path. + """ + self._revalidate() + return self._transform.transform_path_affine(self._transformed_path) + + def get_affine(self): + return self._transform.get_affine() + + +class TransformedPatchPath(TransformedPath): + """ + A `TransformedPatchPath` caches a non-affine transformed copy of the + `~.patches.Patch`. This cached copy is automatically updated when the + non-affine part of the transform or the patch changes. + """ + + def __init__(self, patch): + """ + Parameters + ---------- + patch : `~.patches.Patch` + """ + # Defer to TransformedPath.__init__. + super().__init__(patch.get_path(), patch.get_transform()) + self._patch = patch + + def _revalidate(self): + patch_path = self._patch.get_path() + # Force invalidation if the patch path changed; otherwise, let base + # class check invalidation. + if patch_path != self._path: + self._path = patch_path + self._transformed_path = None + super()._revalidate() + + +def nonsingular(vmin, vmax, expander=0.001, tiny=1e-15, increasing=True): + """ + Modify the endpoints of a range as needed to avoid singularities. + + Parameters + ---------- + vmin, vmax : float + The initial endpoints. + expander : float, default: 0.001 + Fractional amount by which *vmin* and *vmax* are expanded if + the original interval is too small, based on *tiny*. + tiny : float, default: 1e-15 + Threshold for the ratio of the interval to the maximum absolute + value of its endpoints. If the interval is smaller than + this, it will be expanded. This value should be around + 1e-15 or larger; otherwise the interval will be approaching + the double precision resolution limit. + increasing : bool, default: True + If True, swap *vmin*, *vmax* if *vmin* > *vmax*. + + Returns + ------- + vmin, vmax : float + Endpoints, expanded and/or swapped if necessary. + If either input is inf or NaN, or if both inputs are 0 or very + close to zero, it returns -*expander*, *expander*. + """ + + if (not np.isfinite(vmin)) or (not np.isfinite(vmax)): + return -expander, expander + + swapped = False + if vmax < vmin: + vmin, vmax = vmax, vmin + swapped = True + + # Expand vmin, vmax to float: if they were integer types, they can wrap + # around in abs (abs(np.int8(-128)) == -128) and vmax - vmin can overflow. + vmin, vmax = map(float, [vmin, vmax]) + + maxabsvalue = max(abs(vmin), abs(vmax)) + if maxabsvalue < (1e6 / tiny) * np.finfo(float).tiny: + vmin = -expander + vmax = expander + + elif vmax - vmin <= maxabsvalue * tiny: + if vmax == 0 and vmin == 0: + vmin = -expander + vmax = expander + else: + vmin -= expander*abs(vmin) + vmax += expander*abs(vmax) + + if swapped and not increasing: + vmin, vmax = vmax, vmin + return vmin, vmax + + +def interval_contains(interval, val): + """ + Check, inclusively, whether an interval includes a given value. + + Parameters + ---------- + interval : (float, float) + The endpoints of the interval. + val : float + Value to check is within interval. + + Returns + ------- + bool + Whether *val* is within the *interval*. + """ + a, b = interval + if a > b: + a, b = b, a + return a <= val <= b + + +def _interval_contains_close(interval, val, rtol=1e-10): + """ + Check, inclusively, whether an interval includes a given value, with the + interval expanded by a small tolerance to admit floating point errors. + + Parameters + ---------- + interval : (float, float) + The endpoints of the interval. + val : float + Value to check is within interval. + rtol : float, default: 1e-10 + Relative tolerance slippage allowed outside of the interval. + For an interval ``[a, b]``, values + ``a - rtol * (b - a) <= val <= b + rtol * (b - a)`` are considered + inside the interval. + + Returns + ------- + bool + Whether *val* is within the *interval* (with tolerance). + """ + a, b = interval + if a > b: + a, b = b, a + rtol = (b - a) * rtol + return a - rtol <= val <= b + rtol + + +def interval_contains_open(interval, val): + """ + Check, excluding endpoints, whether an interval includes a given value. + + Parameters + ---------- + interval : (float, float) + The endpoints of the interval. + val : float + Value to check is within interval. + + Returns + ------- + bool + Whether *val* is within the *interval*. + """ + a, b = interval + return a < val < b or a > val > b + + +def offset_copy(trans, fig=None, x=0.0, y=0.0, units='inches'): + """ + Return a new transform with an added offset. + + Parameters + ---------- + trans : `Transform` subclass + Any transform, to which offset will be applied. + fig : `~matplotlib.figure.Figure`, default: None + Current figure. It can be None if *units* are 'dots'. + x, y : float, default: 0.0 + The offset to apply. + units : {'inches', 'points', 'dots'}, default: 'inches' + Units of the offset. + + Returns + ------- + `Transform` subclass + Transform with applied offset. + """ + _api.check_in_list(['dots', 'points', 'inches'], units=units) + if units == 'dots': + return trans + Affine2D().translate(x, y) + if fig is None: + raise ValueError('For units of inches or points a fig kwarg is needed') + if units == 'points': + x /= 72.0 + y /= 72.0 + # Default units are 'inches' + return trans + ScaledTranslation(x, y, fig.dpi_scale_trans) diff --git a/lib/python3.12/site-packages/sglang/srt/__pycache__/code_completion_parser.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/__pycache__/code_completion_parser.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a38f064531e14db7b88ad4a990eca341d8579e1 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/__pycache__/code_completion_parser.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/__pycache__/custom_op.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/__pycache__/custom_op.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f443a202f7f6e487a8745a752febd2d10dcd9e6 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/__pycache__/custom_op.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/__pycache__/model_parallel.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/__pycache__/model_parallel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c15b42c034df8e5ccc47b051d88f4dd42343259c Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/__pycache__/model_parallel.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/__pycache__/patch_torch.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/__pycache__/patch_torch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..66ffbc721ecfacceaddd8a7f6c6cbcc3e889d0a4 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/__pycache__/patch_torch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/__pycache__/reasoning_parser.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/__pycache__/reasoning_parser.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..feae1352663491b1d8fa946fc6f1e69801fd8bf2 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/__pycache__/reasoning_parser.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8df597f2448c268a98962f8763eeff9e2f8d985 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/activation.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/activation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0810cbace5ce8e588eb5feea3a0286167736a4c Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/activation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/dp_attention.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/dp_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c98504d1480ed09d963fdf99bcd26d06001d1856 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/dp_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/elementwise.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/elementwise.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe654265e337acbd8bbc16a955d88ccc671bcdde Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/elementwise.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/layernorm.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/layernorm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1afaf359bdc7bb138b6ccd84343af3b92b8193d0 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/layernorm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/linear.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/linear.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e828cdb59cfd2d3152529ba798b35531e50ed6e Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/linear.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/logits_processor.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/logits_processor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09033916ce4de24b1c8d66943328f0c215f606b6 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/logits_processor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/parameter.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/parameter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e6edae766e506046d6f232ddf486d2206ddc0929 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/parameter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/pooler.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/pooler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..19f21f25325490208194b01e5e4d0dffa4cf3543 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/pooler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/radix_attention.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/radix_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc9da8beb53aba8fab9dae9f808e911f8d331742 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/radix_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/rotary_embedding.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/rotary_embedding.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e181442f413cab30bd689991b792e3cf929caf96 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/rotary_embedding.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/sampler.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/sampler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d8dadbf0ff2988488c01dd679372dc6b0ee1d69 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/sampler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/torchao_utils.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/torchao_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d23c41c8af315f69da9a81722bc4be132240c11 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/torchao_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/vocab_parallel_embedding.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/vocab_parallel_embedding.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0bb6c558e6ad476a74ff5214c1c2b698e3166c54 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/__pycache__/vocab_parallel_embedding.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/base_attn_backend.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/base_attn_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..52bcd5fba5833f41e0ed1c5bc5fd08d72fedaf62 --- /dev/null +++ b/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() diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/cutlass_mla_backend.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/cutlass_mla_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..9ff5dfabf7546a636f275fb0ac08e9725b5d2db9 --- /dev/null +++ b/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) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/double_sparsity_backend.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/double_sparsity_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..47b867f61367d6dca5778f06584ccb9b5d50728e --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/flashattention_backend.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/flashattention_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..4e8543213a1ca398ec84e580aedeab8c13fb5e6a --- /dev/null +++ b/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, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/flashinfer_backend.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/flashinfer_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..8586006dc0b0f4b472d391afe21d78a1467b29c4 --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/flashinfer_mla_backend.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/flashinfer_mla_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..81afcb9dac5e66fb5c7b479513a9e0c9dea8c6df --- /dev/null +++ b/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, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/flashmla_backend.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/flashmla_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..1513c1c71265f158e03727da900e0d19ae52f701 --- /dev/null +++ b/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) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/torch_native_backend.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/torch_native_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..78ed042de85564b4554e27c51db1ba556d61f024 --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_backend.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..0aa3a695e4fcfd70a6681f14ac190891c87e0cf9 --- /dev/null +++ b/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) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/decode_attention.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/decode_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..b7dbdb16dafb9cc7974069fc06b7b0798dd1c7ae --- /dev/null +++ b/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, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/double_sparsity_attention.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/double_sparsity_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..6b6ede927bf7a8d5009601c788bc62399012efb0 --- /dev/null +++ b/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, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/extend_attention.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/extend_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..35e5c21c62fa15e3665f666ff7ba304269dec343 --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/prefill_attention.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/prefill_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..ac0fc72af1407d0b19028bcdc07ae9a87fcb70a7 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/prefill_attention.py @@ -0,0 +1,217 @@ +# 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 supporst page size = 1. +""" + +# Adapted from +# https://github.com/ModelTC/lightllm/blob/f2a54f0912293f683bf1d1695fd12c4098a5bf82/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py#L1 +import torch +import triton +import triton.language as tl + +from sglang.srt.utils import is_cuda, is_hip + +_is_cuda = is_cuda() +_is_hip = is_hip() + +if _is_cuda or _is_hip: + CUDA_CAPABILITY = torch.cuda.get_device_capability() + + +@triton.jit +def _fwd_kernel( + Q, + K, + V, + sm_scale, + B_Start_Loc, + B_Seqlen, + Out, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_obs, + stride_oh, + kv_group_num: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_CAUSAL: tl.constexpr, + Lk: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + start_m = tl.program_id(2) + + cur_kv_head = cur_head // kv_group_num + + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch) + + block_start_loc = BLOCK_M * start_m + + # initialize offsets + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + off_q = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs + + cur_head * stride_qh + + offs_d[None, :] + ) + off_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] + off_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] + + mask_d = offs_d < Lk + + q = tl.load( + Q + off_q, + mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :]), + other=0.0, + ) + + k_ptrs = K + off_k + v_ptrs = V + off_v + + # initialize pointer to m and l + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + + block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0) + + end_n = ( + cur_batch_seq_len + if not IS_CAUSAL + else tl.minimum((start_m + 1) * BLOCK_M, cur_batch_seq_len) + ) + for start_n in range(0, block_mask * end_n, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + k = tl.load( + k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs, + mask=((start_n + offs_n[None, :]) < cur_batch_seq_len) & (mask_d[:, None]), + other=0.0, + ) + # mask = tl.load(mask_ptrs + start_n, mask=start_n + offs_n < cur_batch_end_loc, other=0.0) + + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk += tl.dot(q, k) + qk *= sm_scale + + if IS_CAUSAL: + qk += tl.where( + (start_n + offs_n[None, :] < cur_batch_seq_len) + & (offs_m[:, None] >= (start_n + offs_n[None, :])), + 0, + float("-inf"), + ) + else: + qk += tl.where( + (start_n + offs_n[None, :]) < cur_batch_seq_len, 0, float("-inf") + ) + + # -- compute m_ij, p, l_ij + m_ij = tl.max(qk, 1) + p = tl.exp(qk - m_ij[:, None]) + l_ij = tl.sum(p, 1) + # -- update m_i and l_i + m_i_new = tl.maximum(m_i, m_ij) + alpha = tl.exp(m_i - m_i_new) + beta = tl.exp(m_ij - m_i_new) + l_i_new = alpha * l_i + beta * l_ij + # -- update output accumulator -- + # scale p + p_scale = beta / l_i_new + p = p * p_scale[:, None] + # scale acc + acc_scale = l_i / l_i_new * alpha + acc = acc * acc_scale[:, None] + # update acc + v = tl.load( + v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs, + mask=((start_n + offs_n[:, None]) < cur_batch_seq_len) & (mask_d[None, :]), + other=0.0, + ) + + p = p.to(v.dtype) + acc += tl.dot(p, v) + # update m_i and l_i + l_i = l_i_new + m_i = m_i_new + # initialize pointers to output + off_o = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + + cur_head * stride_oh + + offs_d[None, :] + ) + out_ptrs = Out + off_o + tl.store( + out_ptrs, acc, mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :]) + ) + + +def context_attention_fwd( + q, k, v, o, b_start_loc, b_seq_len, max_input_len, is_causal=True +): + """ + q, k, v: [b * s, head, head_dim] + b_start_loc: [b] + b_seq_len: [b] + out: [b * s, head, head_dim] + """ + if (_is_cuda or _is_hip) and CUDA_CAPABILITY[0] > 8: + BLOCK = 128 + else: + BLOCK = 64 + + Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] + + sm_scale = 1.0 / (Lq**0.5) + batch, head = b_seq_len.shape[0], q.shape[1] + kv_group_num = q.shape[1] // k.shape[1] + + grid = (batch, head, triton.cdiv(max_input_len, BLOCK)) + num_warps = 4 if Lk <= 64 else 8 + + _fwd_kernel[grid]( + q, + k, + v, + sm_scale, + b_start_loc, + b_seq_len, + o, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + o.stride(0), + o.stride(1), + kv_group_num=kv_group_num, + BLOCK_M=BLOCK, + BLOCK_DMODEL=triton.next_power_of_2(Lk), + BLOCK_N=BLOCK, + IS_CAUSAL=is_causal, + num_warps=num_warps, + num_stages=1, + Lk=Lk, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py new file mode 100644 index 0000000000000000000000000000000000000000..89b17ce2db1545661d26686581036e759ec1f6ea --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py @@ -0,0 +1,439 @@ +# 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 triton +import triton.language as tl + +from sglang.srt.layers.attention.triton_ops.decode_attention import ( + _decode_softmax_reducev_fwd, +) + + +def is_hip(): + return triton.runtime.driver.active.get_current_target().backend == "hip" + + +_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_grouped_kernel_stage1_rope( + Q, # Holds [Q_NOPE; Q_PE], b x h x (d+r) + K_Buffer, # Holds [KV; K_PE], b*s x (c+r) + V_buffer, # Holds [KV], b*s x (c) + cos_sin_cache, # max_seq_len x (rotary_dim * 2) + positions, # sequence positions + sm_scale, + kv_indptr, + kv_indices, + Att_Out, # b x h x NUM_KV_SPLITS x (kv_lora_rank + 1) + k_pe_t_out, + stride_qb, + stride_qh, + stride_buf_kbs, + stride_buf_vbs, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + stride_kpe_tokens_out_b, + stride_cos_sin_cache_s, + stride_positions_b, + rotary_dim: tl.constexpr, + kv_lora_rank: tl.constexpr, + qk_rope_head_dim: tl.constexpr, + kv_group_num: tl.constexpr, + q_head_num: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_R: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + NUM_KV_SPLITS: tl.constexpr, + logit_cap: tl.constexpr, + USE_ROPE: tl.constexpr, + IS_NEOX_STYLE: tl.constexpr, +): + + cur_batch = tl.program_id(0) + cur_head_id = tl.program_id(1) + 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_c = tl.arange(0, BLOCK_C) + offs_qk_r = tl.arange(kv_lora_rank, kv_lora_rank + BLOCK_R) # to get the k_pe + + off_q_pe = ( + cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_qk_r[None, :] + ) + offs_q = cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_c[None, :] + + mask_c = offs_c < kv_lora_rank + mask_qk_r = offs_qk_r < (kv_lora_rank + qk_rope_head_dim) + + 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 + + q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_c[None, :]), other=0.0) + q_pe = tl.load( + Q + off_q_pe, mask=(mask_h[:, None]) & (mask_qk_r[None, :]), other=0.0 + ) + + kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_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) + + # apply rotary embedding for q_pe, and k_pe (last token per batch of K_PE) + LAST_SPLIT = split_kv_end == cur_batch_seq_len + k_pe_last_token = tl.zeros([BLOCK_R], dtype=q.dtype) + + if USE_ROPE: + if IS_NEOX_STYLE: + # [BLOCK_ROTARY // 2, BLOCK_ROTARY // 2 + 1, BLOCK_ROTARY // 2 + 2, ..., 0, 1, 2, ..., BLOCK_ROTARY // 2 - 1, pass:] + offs_qk_rot_r = kv_lora_rank + ( + (tl.arange(0, BLOCK_R) + (rotary_dim // 2)) % rotary_dim + ) + # Which elements to flip + mask_rotate = tl.arange(0, BLOCK_R) < (rotary_dim // 2) + # [0 , 1, 2, ..., rotary_dim // 2 - 1, 0 , 1, 2, ..., rotary_dim // 2 - 1] + offs_rotary = tl.arange(0, BLOCK_R) % (rotary_dim // 2) + else: + # [1, 0, 3, 2, 5, 4, ..., BLOCK_R, BLOCK_R - 1] + offs_qk_rot_r = ( + kv_lora_rank + + (((tl.arange(0, BLOCK_R) + 1) % 2) * 2) + - 1 + + tl.arange(0, BLOCK_R) + ) + mask_rotate = tl.arange(0, BLOCK_R) % 2 < 1 + # [0, 0, 1, 1, ..., rotary_dim // 2 - 1, rotary_dim // 2 - 1] + offs_rotary = tl.arange(0, BLOCK_R) // 2 + + if qk_rope_head_dim > rotary_dim: + offs_qk_rot_r = tl.where( + tl.arange(0, BLOCK_R) < rotary_dim, offs_qk_rot_r, tl.arange(0, BLOCK_R) + ) + offs_rotary = tl.where( + tl.arange(0, BLOCK_R) < rotary_dim, offs_rotary, tl.arange(0, BLOCK_R) + ) + + mask_rotary = tl.arange(0, BLOCK_R) < rotary_dim + + pos = tl.load(positions + cur_batch * stride_positions_b) + cos = tl.load( + cos_sin_cache + pos * stride_cos_sin_cache_s + offs_rotary, + mask=mask_rotary, + other=1.0, + ) + sin = tl.load( + cos_sin_cache + + pos * stride_cos_sin_cache_s + + offs_rotary + + rotary_dim // 2, + mask_rotary, + other=0.0, + ) + + off_q_pe_rot = ( + cur_batch * stride_qb + + cur_head[:, None] * stride_qh + + offs_qk_rot_r[None, :] + ) + mask_qk_rot_r = offs_qk_rot_r < (kv_lora_rank + qk_rope_head_dim) + + # 0, 2, 4,.... 1, 3, 5... + q_pe_rot = tl.load( + Q + off_q_pe_rot, + mask=(mask_h[:, None]) & (mask_qk_rot_r[None, :]), + other=0.0, + ) + q_pe_rot = tl.where(mask_rotate[None, :], -q_pe_rot, q_pe_rot) + + q_pe = q_pe * cos + q_pe_rot * sin + + # we only apply to the last token in the K_PE + if LAST_SPLIT: + # debug assert + if (cur_batch == 0 and cur_head == 0) and split_kv_id < NUM_KV_SPLITS - 1: + tl.device_assert(False, "Only last split should compute k_pe") + + kv_loc = tl.load( + kv_indices + cur_batch_kv_start_idx + cur_batch_seq_len - 1 + ) + offs_buf_k_pe_last_token = kv_loc * stride_buf_kbs + offs_qk_r + offs_buf_k_pe_rot_last_token = kv_loc * stride_buf_kbs + offs_qk_rot_r + k_pe_last_token = tl.load(K_Buffer + offs_buf_k_pe_last_token) + + k_pe_rot_last_token = tl.load(K_Buffer + offs_buf_k_pe_rot_last_token) + k_pe_rot_last_token = tl.where( + mask_rotate, -k_pe_rot_last_token, k_pe_rot_last_token + ) + + k_pe_last_token = k_pe_last_token * cos + k_pe_rot_last_token * sin + + 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_C], dtype=tl.float32) + + if split_kv_end > split_kv_start: + 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_kv = kv_loc[None, :] * stride_buf_kbs + offs_c[:, None] + offs_buf_k_pe = kv_loc[None, :] * stride_buf_kbs + offs_qk_r[:, None] + + k_pe = tl.load( + K_Buffer + offs_buf_k_pe, + mask=(offs_n[None, :] < split_kv_end) & (mask_qk_r[:, None]), + other=0.0, + ) # positional embedding part of keys + + if (USE_ROPE and LAST_SPLIT) and start_n >= cur_batch_seq_len - BLOCK_N: + k_pe = tl.where( + offs_n[None, :] != (split_kv_end - 1), + k_pe, + k_pe_last_token[:, None], + ) + + # (16, 64) x (64, 32) + # dot product of rope parts + qk = tl.dot(q_pe, k_pe.to(q_pe.dtype)) + + kv = tl.load( + K_Buffer + offs_buf_kv, + mask=(offs_n[None, :] < split_kv_end) & (mask_c[:, None]), + other=0.0, + ) # the shared latent tensor for keys and values + + # (16, 512) x (512, 32) + # dot product of nope parts + qk += tl.dot(q, kv) + + 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 + offs_c[None, :] + v = tl.load( + V_buffer + offs_buf_v, + mask=(offs_n[:, None] < split_kv_end) & (mask_c[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] + # (16, 32) x (32, 512) + 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_c[None, :] + ) + + if USE_ROPE: + if LAST_SPLIT: + k_pe_last_token_ptrs = ( + k_pe_t_out + + cur_batch * stride_kpe_tokens_out_b + + tl.arange(0, BLOCK_R) + ) + tl.store(k_pe_last_token_ptrs, k_pe_last_token, mask=mask_qk_r) + + tl.store( + Att_Out + offs_mid_o, + acc / e_sum[:, None], + mask=(mask_h[:, None]) & (mask_c[None, :]), + ) + + offs_mid_o_1 = ( + cur_batch * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + + kv_lora_rank + ) + + tl.store( + Att_Out + offs_mid_o_1, + e_max + tl.log(e_sum), + mask=mask_h, + ) + + +# TODO rope offset +def _decode_grouped_att_m_fwd_rope( + q, + k_buffer, + v_buffer, + att_out, + k_pe_tokens_out, + kv_lora_rank, # c + cos_sin_cache, + positions, + rotary_dim, + kv_indptr, + kv_indices, + num_kv_splits, + sm_scale, + logit_cap, + use_rope, + is_neox_style=True, +): + if use_rope: + assert ( + k_pe_tokens_out is not None + ), "We must output the k_pe tokens with rope applied if rope fusion enabled." + + BLOCK = 32 + + # # [TODO] work around shmem limit on MI3xx + # if _is_hip and kv_lora_rank >= 576: + # BLOCK = 16 + + qk_rope_head_dim = k_buffer.shape[-1] - kv_lora_rank + batch, head_num = kv_indptr.shape[0] - 1, q.shape[1] + kv_group_num = q.shape[1] // k_buffer.shape[1] + + BLOCK_C = triton.next_power_of_2(kv_lora_rank) + BLOCK_R = triton.next_power_of_2(qk_rope_head_dim) + + BLOCK_H = 16 + NUM_KV_SPLITS = num_kv_splits + grid = ( + batch, + triton.cdiv(head_num, min(BLOCK_H, kv_group_num)), + NUM_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_rope[grid]( + q, + k_buffer, + v_buffer, + cos_sin_cache, + positions, + sm_scale, + kv_indptr, + kv_indices, + att_out, + k_pe_tokens_out, + q.stride(0), + q.stride(1), + k_buffer.stride(0), + v_buffer.stride(0), + att_out.stride(0), + att_out.stride(1), + att_out.stride(2), + k_pe_tokens_out.stride(0) if use_rope else 0, + cos_sin_cache.stride(0) if use_rope else 0, + positions.stride(0) if use_rope else 0, + rotary_dim, + kv_lora_rank, + qk_rope_head_dim, + kv_group_num=kv_group_num, + q_head_num=head_num, + BLOCK_C=BLOCK_C, + BLOCK_R=BLOCK_R, + BLOCK_N=BLOCK, + BLOCK_H=BLOCK_H, + NUM_KV_SPLITS=NUM_KV_SPLITS, + logit_cap=logit_cap, + USE_ROPE=use_rope, + IS_NEOX_STYLE=is_neox_style, + num_warps=4, + num_stages=num_stages, + **extra_kargs + ) + + +def decode_attention_fwd_grouped_rope( + q, + k_buffer, + v_buffer, + o, + kv_indptr, + kv_indices, + k_pe_tokens, + kv_lora_rank, + rotary_dim, + cos_sin_cache, + positions, + attn_logits, + num_kv_splits, + sm_scale, + logit_cap=0.0, + use_rope=False, + is_neox_style=False, +): + _decode_grouped_att_m_fwd_rope( + q, + k_buffer, + v_buffer, + attn_logits, + k_pe_tokens, + kv_lora_rank, + cos_sin_cache, + positions, + rotary_dim, + kv_indptr, + kv_indices, + num_kv_splits, + sm_scale, + logit_cap, + use_rope, + is_neox_style, + ) + _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, kv_indptr, num_kv_splits) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/utils.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c87aa45d7d8377be228c8aa6a70367c4d84cf254 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/attention/utils.py @@ -0,0 +1,92 @@ +import triton +import triton.language as tl + + +@triton.jit +def create_flashinfer_kv_indices_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices_ptr, + page_kernel_lens_ptr, + kv_indptr, + kv_start_idx, + kv_indices_ptr, + req_to_token_ptr_stride: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(axis=0) + + # find the req pool idx, this is for batch to token + req_pool_index = tl.load(req_pool_indices_ptr + pid) + kv_indices_offset = tl.load(kv_indptr + pid) + + kv_start = 0 + kv_end = 0 + if kv_start_idx: + kv_start = tl.load(kv_start_idx + pid).to(tl.int32) + kv_end = kv_start + kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32) + + num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < kv_end - kv_start + data = tl.load( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + kv_start + + offset, + mask=mask, + ) + tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask) + + +@triton.jit +def create_flashmla_kv_indices_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices_ptr, + page_kernel_lens_ptr, + kv_start_idx, + kv_indices_ptr, + req_to_token_ptr_stride: tl.constexpr, + kv_indices_ptr_stride: tl.constexpr, + PAGED_SIZE: tl.constexpr = 64, +): + BLOCK_SIZE: tl.constexpr = 4096 + NUM_PAGE_PER_BLOCK: tl.constexpr = 64 + pid = tl.program_id(axis=0) + + # find the req pool idx, this is for batch to token + req_pool_index = tl.load(req_pool_indices_ptr + pid) + + kv_start = 0 + kv_end = 0 + if kv_start_idx: + kv_start = tl.load(kv_start_idx + pid).to(tl.int32) + kv_end = kv_start + + kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32) + + num_paged = tl.cdiv(kv_end - kv_start, PAGED_SIZE) + num_pages_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) + + for i in range(num_pages_loop): + paged_offset = ( + tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK + ) * PAGED_SIZE + paged_offset_out = tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK + + mask = paged_offset <= num_paged * PAGED_SIZE + mask_out = paged_offset_out <= num_paged + + data = tl.load( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + kv_start + + paged_offset, + mask=mask, + ) + tl.store( + kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out, + data // PAGED_SIZE, + mask=mask_out, + ) diff --git a/lib/python3.12/site-packages/sglang/srt/layers/attention/vision.py b/lib/python3.12/site-packages/sglang/srt/layers/attention/vision.py new file mode 100644 index 0000000000000000000000000000000000000000..d65104beb587b320df118149f1e8f17c5e26f86e --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/attention/vision.py @@ -0,0 +1,368 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from sglang.srt.distributed import parallel_state +from sglang.srt.distributed import utils as dist_utils +from sglang.srt.layers.attention.triton_ops.prefill_attention import ( + context_attention_fwd, +) +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.quantization import QuantizationConfig +from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb, rotate_half +from sglang.srt.utils import add_prefix + + +class VisionAttention(nn.Module): + r""" + Multi-headed attention without any cache, mostly used for ViT. + + + Args: + use_qkv_parallel (bool, optional): If True, use QKV-parallel attention. + use_context_forward (bool, default to True): + if ``True``, a flash_attn style attention will be applied + Otherwise, a full-sequence attention will be applied. + softmax_in_single_precision (bool, default to False): + if ``True``, the softmax will be performed in single-precision + Otherwise, it will be performed in half-precision + + """ + + def __init__( + self, + embed_dim: int, + num_heads: int, + projection_size: int, + use_qkv_parallel: bool, + quant_config: Optional[QuantizationConfig] = None, + dropout: float = 0.0, + use_context_forward: bool = True, + softmax_in_single_precision: bool = False, + flatten_batch: bool = False, + prefix: str = "", + ): + super().__init__() + self.use_context_forward = use_context_forward + world_size = parallel_state.get_tensor_model_parallel_world_size() + self.dropout = dropout + self.head_size = embed_dim // num_heads + self.hidden_size_per_attention_head = dist_utils.divide( + projection_size, num_heads + ) + self.num_attention_heads_per_partition = dist_utils.divide( + num_heads, world_size + ) + + if self.use_context_forward: + self.qkv_backend = VisionTritonAttention() + else: + self.qkv_backend = VisionSdpaAttention( + head_size=self.head_size, + dropout=dropout, + flatten_batch=flatten_batch, + softmax_in_single_precision=softmax_in_single_precision, + ) + + self.use_qkv_parallel = use_qkv_parallel + if use_qkv_parallel: + self.qkv_proj = QKVParallelLinear( + hidden_size=embed_dim, + head_size=self.head_size, + total_num_heads=num_heads, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + else: + self.qkv_proj = ColumnParallelLinear( + input_size=embed_dim, + output_size=3 * projection_size, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.proj = RowParallelLinear( + input_size=embed_dim, + output_size=embed_dim, + quant_config=quant_config, + prefix=add_prefix("proj", prefix), + ) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + r""" + Args: + x: [b, s, embed_dim] + cu_seqlens: [b] + Returns: + [s, b, head * head_size] + """ + bsz, s, _ = x.shape + head = self.num_attention_heads_per_partition + if self.use_qkv_parallel: + # [b, s, embed_dim] --> [b, s, embed_dim] + qkv, _ = self.qkv_proj(x) + q, k, v = qkv.chunk(3, dim=-1) + + # [b, s, embed_dim] --> [b * s, head, head_size] + q, k, v = [x.reshape(bsz * s, head, -1).contiguous() for x in (q, k, v)] + else: + # [b, s, embed_dim] --> [s, b, embed_dim] + x = rearrange(x, "b s ... -> s b ...") + # [s, b, embed_dim] --> [s, b, head * 3 * head_size] + qkv, _ = self.qkv_proj(x) + # [s, b, head * 3 * head_size] --> [s, b, head, 3 * head_size] + new_x_shape = qkv.size()[:-1] + ( + head, + 3 * self.hidden_size_per_attention_head, + ) + qkv = qkv.view(*new_x_shape) + + # [s, b, head, 3 * head_size] --> 3 [s, b, head, head_size] + q, k, v = dist_utils.split_tensor_along_last_dim(qkv, 3) + + # [s, b, head, head_size] --> [b, s, head, head_size] + q, k, v = [ + rearrange(x, "s b ... -> b s ...").contiguous() for x in (q, k, v) + ] + + if position_embeddings is not None: + cos, sin = position_embeddings + original_shape = q.shape + # [total_tokens, head, head_size] + q = q.view(-1, head, self.head_size) + k = k.view(-1, head, self.head_size) + + q, k = apply_rotary_pos_emb(q, k, cos, sin) + + q = q.view(original_shape) + k = k.view(original_shape) + + if self.use_qkv_parallel: + pass + else: + # [b, s, head, head_size] --> [b * s, head, head_size] + q, k, v = [rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v]] + + output = self.qkv_backend.forward(q, k, v, bsz, cu_seqlens, attention_mask) + + if self.use_qkv_parallel: + # [b * s, h, head_size] --> [b, s, h * head_size] + output = rearrange(output, "(b s) ... h d -> b s ... (h d)", b=bsz) + + # [b, s, h * head_size] --> [b, s, h * head_size] + output, _ = self.proj(output) + else: + # [b * s, h, head_size] --> [s, b, h * head_size] + context_layer = rearrange( + output, "(b s) h d -> s b (h d)", b=bsz, s=s + ).contiguous() + + # [s, b, h * head_size] --> [s, b, h * head_size] + output, _ = self.proj(context_layer) + + # [s, b, h * head_size] --> [b, s, h * head_size] + output = output.view(bsz, s, -1) + + return output + + +class VisionSdpaAttention(nn.Module): + r""" + Scaled Dot Product Attention inner product + + """ + + def __init__( + self, + head_size: int, + dropout: float = 0.0, + flatten_batch: bool = False, + softmax_in_single_precision: bool = False, + ): + super().__init__() + self.head_size = head_size + self.flatten_batch = flatten_batch + self.softmax_in_single_precision = softmax_in_single_precision + self.dropout = dropout + + @staticmethod + @lru_cache(maxsize=128) + def _generate_mask_cache( + s: int, flatten_batch: bool, cu_seqlens: tuple + ) -> torch.BoolTensor: + """ + Generate a boolean attention mask with caching mechanism. + Args: + s: sequence length + flatten_batch: whether to flatten batch dimension + cu_seqlens: tuple of cumulative sequence lengths + Returns: + attention mask tensor + """ + if flatten_batch: + mask = torch.zeros([1, s, s], dtype=torch.bool) + for i in range(1, len(cu_seqlens)): + start = cu_seqlens[i - 1] + end = cu_seqlens[i] + mask[..., start:end, start:end] = True + else: + # [1, 1, 1, s] + row_indices = torch.arange(s).view(1, 1, 1, s) + # [1, 1, s, 1] + col_indices = torch.arange(s).view(1, 1, s, 1) + # [b, 1, 1, 1] + seq_lens = torch.tensor( + [end - start for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:])], + ).view(-1, 1, 1, 1) + + mask = (row_indices < seq_lens) & (col_indices < seq_lens) + + return mask + + def generate_patch_attention_mask( + self, + s: int, + cu_seqlens: Optional[torch.Tensor], + flatten_batch: bool = False, + ) -> Optional[torch.Tensor]: + r""" + Creates a non-causal 4D mask of shape `(b, 1, s, s)` or `(1, 1, s, s)`. + Args: + s: sequence length + cu_seqlens: cumulative sequence lengths tensor. If not, returns an empty mask + flatten_batch: whether to flatten batch dimension + Returns: + attention mask tensor or None + """ + if cu_seqlens is None: + return None + + cu_seqlens_tuple = tuple(cu_seqlens.cpu().tolist()) + + return self._generate_mask_cache(s, flatten_batch, cu_seqlens_tuple) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + bsz: int, + cu_seqlens: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + r""" + Args: + cu_seqlens: [b] + Returns: + [b * s, h, head_size] + """ + if self.flatten_batch: + assert bsz == 1, "flatten_batch is True, bsz must be 1" + + s = q.shape[0] // bsz + + # [b, 1, s, s] + if attention_mask is None: + attention_mask = self.generate_patch_attention_mask( + s, cu_seqlens, flatten_batch=self.flatten_batch + ) + + if attention_mask is None: + if self.softmax_in_single_precision: + raise RuntimeError("Empty attention mask") + else: + attention_mask = attention_mask.to(device=q.device) + + q, k, v = [rearrange(x, "(b s) h d -> b h s d", b=bsz) for x in [q, k, v]] + + if self.softmax_in_single_precision: + scale = self.head_size**-0.5 + k_transposed = rearrange(k, "b h s d -> b h d s") + attn_weights = torch.matmul(q, k_transposed) * scale + del k, k_transposed + attention_mask = (~attention_mask) * torch.finfo(q.dtype).min + attn_weights = attn_weights + attention_mask + del attention_mask + # full-precision + attn_weights = nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(q.dtype) + attn_weights = nn.functional.dropout( + attn_weights, p=self.dropout, training=False + ) + output = torch.matmul(attn_weights, v) + del attn_weights, v + else: + # SDPA + # [b, h, s, head_size] + output = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attention_mask, + dropout_p=self.dropout, + is_causal=False, + ) + + # [b, h, s, head_size] --> [b * s, h, head_size] + output = rearrange(output, "b h s d -> (b s) h d") + + return output + + +class VisionTritonAttention(nn.Module): + """ + Triton-implemented attention without a causal mask + """ + + def __init__( + self, + ): + super().__init__() + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + _bsz: int, + cu_seqlens: Optional[torch.Tensor], + ) -> torch.Tensor: + r""" + Args: + cu_seqlens: [b] + Returns: + [b * s, h, head_size] + """ + + # [b * s, head, head_size] + output = torch.empty_like(q) + seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] + max_seqlen = seq_lens.max().item() + context_attention_fwd( + q, + k, + v, + output, + cu_seqlens.cuda(), + seq_lens.cuda(), + max_seqlen, + is_causal=False, + ) + + return output diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63b4c7da46e1c21879812567f14ea1bdaa997e80 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/awq.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/awq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e453ed0c6692f1cee08f49a320f08892608035d Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/awq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/base_config.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/base_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d661782fc61b2ced92f9210b8731c6241c99662 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/base_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/blockwise_int8.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/blockwise_int8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dc5c277c2da91fcb9e84df27fca48d9d9cca88f8 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/blockwise_int8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/deep_gemm.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/deep_gemm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..69fd617c9260c244aa144b840bcf14323e5148ab Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/deep_gemm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/fp8.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/fp8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ebe7f126e038b06918ba7478c68231a60c24d026 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/fp8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/fp8_kernel.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/fp8_kernel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a6e8f39bd3b7d6ace2e0428eed3490613a5c96b6 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/fp8_kernel.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/fp8_utils.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/fp8_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e3bcd306a6eb1f284866a07c405c0c62b9e27a2 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/fp8_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/gptq.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/gptq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a709aa34bb03264e194ed4d270c5f4b390cd4b49 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/gptq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/int8_kernel.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/int8_kernel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08984180bb8e6b5d2703e6cfc19222602cbfafdf Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/int8_kernel.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/int8_utils.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/int8_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d79279b55ae3a3d378267cee2b22ad161dc4be1 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/int8_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/kv_cache.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/kv_cache.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45ffa9381ac79c56c4d4fef4a148e0b4708d87b5 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/kv_cache.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/modelopt_quant.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/modelopt_quant.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d061f2900b82226b791c0019a3ee421e351ad2bc Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/modelopt_quant.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/moe_wna16.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/moe_wna16.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4c11170fb4e14298d3e6bb39758e4191662d51b1 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/moe_wna16.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b42ce1f7a74c016727cbfd1e3354684760122dd3 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/w8a8_fp8.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/w8a8_fp8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f76d5749faabeb7645e1d26f6088592ca0b6039b Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/w8a8_fp8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/w8a8_int8.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/w8a8_int8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3068d06bc0d961d7b3de359f009aecf8ae8192b2 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/layers/quantization/__pycache__/w8a8_int8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..6ae89c7571a5dc1e866a9c1c698c8be5d79eea78 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..ae4c397ef37ff57f17e66484275f2880f738fbec --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..effb7bda707ac38654851aa0b05785a44e8fd488 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..eb18488946020293ac3cc04cb96d6ac2d4c6df00 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=1536,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..653bb997ee114a81b4cdf7f0760341e3a050087c --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..3618053b65831b95c4bb0f20ef3b9aa816b2d637 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..cb573ccdff98acf08726d763d81ec13dd892f0cc --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=1536,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..eb11c1696d113057c8d027dda35a136517b51a9f --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..035ec027fa56622196b24a03a5042ce010deaebf --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..dd66019ccf317803eef5066296a499248da0f606 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2048,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..5011fd572c49c863a2f33c2cd51c7b4ddcdf2b59 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..d1227c2157990216d2ca51c69ad0944017f53b6a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=2304,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..6f5adbb9361229bc24b19098e3711a3fd15e939b --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..7b10d6d781484473282d8fc81c0170c58a678708 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..77ba0d7477bdbcb036a43263e7aaa6b6913f8f4e --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..1c61451fb34e52deec827f8f63c80fb15830c202 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..66ca4ee165e350212f4d8233622fc1a34b9f2fd0 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=24576,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..d72288781e40ad7d51bc84d496406b81102325b0 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=256,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..6f5adbb9361229bc24b19098e3711a3fd15e939b --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..eccb86a76df0d7302b760ab6d83a8ceb9fa9d0d9 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=3072,K=1536,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..88af48431d8b8791af8df03429704606b670f1f7 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=3072,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..4225c78eb72cb55c7a65b72d27fab9b4f33cc87a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..30932a1ccab41b97e3cb409fcd3fcd6a61eaa1e0 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..65989f012523036746302e174071ab68550d07df --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=32768,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..4225c78eb72cb55c7a65b72d27fab9b4f33cc87a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4096,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..5e6789d00e0ae0a47ef2cc5aad489b24586f255f --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..03dba5ad15ba5f7f49100a5c78e8685e64334b2a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..96e1594a3eabbaedc792b84b07f05ae8752b7251 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..017dca003117d65aa4a82f57b6fbfb1deee58096 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=4608,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..49ac14d2a57613b59a61d9138bcf0378f75175ec --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..9a5ff48b8942957dde9b862aed848390dd267948 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..b57588f9bb6d90b8daadaa12a9db7c1fb08543de --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..eabc423949a24c2a1fb2368a73e5249caf8d07df --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=512,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..13e3c20b74828ebe9d54c533eb54e3c6c3cdf48d --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..ecbca9d8dfa5de3d5ce4ee89a1993b3be5583292 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=576,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..5e40bbe6471ace395efa41df05eaa2acee185bf6 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..40c01c0b92b4b26fe480879dda33f18c5eb59a6d --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..b96981dec4b2d9371bd731fe26041eb5142a3d63 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1024,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..e5c4a1d2c94e5c7864f462e083ea5f530b8efe3f --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=1152,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..d694716c3e2c9ca65a265acf7557ee2762635ec5 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..dfe5c1e43d687e22d066295483e3324a803feb87 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..dfe5c1e43d687e22d066295483e3324a803feb87 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..b2cab0831118e9a8df2d2ea482af3e865e267a80 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..9de2ba9c8fcd8969ad925ceecf7452050999d04d --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..0a1e14cffbb2a894a701352193947d272427db0d --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..c87cb4f957f6b7166d338dad001f8b2b9e606032 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..5bf208ce4b73085838afc23b69b5aeb9082c12a7 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=16384,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..b474e6265bc152f30f3120c5a04fd792b4dba3fd --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..a3f5371f1ce4ff0829bda4cc62cf1f315ad9e95a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=18432,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,26 @@ +{ + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..dfe5c1e43d687e22d066295483e3324a803feb87 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=2048,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..58cdd93e90b8c29bc7a211861711565dbeeb529a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..a5ce5e49b5923b8430e1ed359986a73ad2f588fa --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=2304,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..d8cc0f8967e9358b28e575e9c062d0b2d11422a7 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "128": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 0000000000000000000000000000000000000000..293adce387e066fce75b6e606d4b8b6a5aa10bdb --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + } +}