ejkernel 0.0.72__tar.gz → 0.0.75__tar.gz

{ejkernel-0.0.72 → ejkernel-0.0.75}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: ejkernel
-Version: 0.0.72
+Version: 0.0.75
 Summary: Accelerate, Optimize performance with streamlined training and serving options with JAX.
 Keywords: Deep Learning,Machine Learning,JAX,CUDA,XLA,Triton,Pallas
 Author: Erfan Zare Chavoshi

{ejkernel-0.0.72 → ejkernel-0.0.75}/ejkernel/__init__.py

@@ -43,7 +43,7 @@ import os as _os
 _os.environ.setdefault("TF_GPU_ALLOCATOR", "cuda_malloc_async")
 _os.environ.setdefault("CUTE_DSL_ENABLE_TVM_FFI", "1")
 
-__version__ = "0.0.72"
+__version__ = "0.0.75"
 
 from . import errors, kernels, modules, types, utils, xla_utils
 from .errors import EjkernelRuntimeError

{ejkernel-0.0.72 → ejkernel-0.0.75}/ejkernel/kernels/_pallas/tpu/__init__.py

@@ -21,7 +21,11 @@ High Bandwidth Memory (HBM) for efficient execution.
 
 Available Kernels:
     - flash_attention: Memory-efficient exact attention
+    - flash_mla: Multi-head latent attention forward pass
+    - deepseek_attn: DeepSeek Sparse Attention (MLA + Lightning Indexer)
+    - gated_delta_rule: Gated delta rule linear attention
     - multi_latent_ragged_page_attention: MLA ragged paged attention
+    - multi_latent_ragged_page_attention_v2: MLA ragged paged attention v2
     - blocksparse_attention: Block-sparse attention patterns
     - ring_attention: Distributed attention across devices
     - page_attention: Paged KV cache attention
@@ -33,12 +37,15 @@ Available Kernels:
 
 from .all_gather_matmul import all_gather_matmul
 from .blocksparse_attention import blocksparse_attention as blocksparse_attention
+from .deepseek_attn import deepseek_attn
 from .flash_attention import flash_attention
 from .flash_mla import flash_mla
 from .gated_delta_rule import gated_delta_rule
 from .grouped_matmul import grouped_matmul
 from .grouped_matmulv2 import grouped_matmulv2
+from .grouped_matmulv3 import grouped_matmulv3
 from .multi_latent_ragged_page_attention import multi_latent_ragged_page_attention
+from .multi_latent_ragged_page_attention_v2 import multi_latent_ragged_page_attention_v2
 from .page_attention import page_attention
 from .prefill_page_attention import prefill_page_attention
 from .quantized_matmul import quantized_matmul
@@ -51,12 +58,15 @@ from .ring_attention import ring_attention
 __all__ = (
     "all_gather_matmul",
     "blocksparse_attention",
+    "deepseek_attn",
     "flash_attention",
     "flash_mla",
     "gated_delta_rule",
     "grouped_matmul",
     "grouped_matmulv2",
+    "grouped_matmulv3",
     "multi_latent_ragged_page_attention",
+    "multi_latent_ragged_page_attention_v2",
     "page_attention",
     "prefill_page_attention",
     "quantized_matmul",

ejkernel-0.0.75/ejkernel/kernels/_pallas/tpu/deepseek_attn/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2025 The EasyDeL/ejKernel Author @erfanzar (Erfan Zare Chavoshi).
+#
+# 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
+#
+#     https://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.
+
+
+"""Pallas TPU backend for DeepSeek Sparse Attention (DSA).
+
+Combines MLA (compressed KV latent + on-the-fly projection) with a
+Lightning Indexer for dynamic top-k token selection on TPU.
+
+Reference:
+    DeepSeek-V3.2: https://arxiv.org/abs/2512.02556
+"""
+
+from ._interface import deepseek_attn
+
+__all__ = ("deepseek_attn",)

ejkernel-0.0.75/ejkernel/kernels/_pallas/tpu/deepseek_attn/_interface.py

@@ -0,0 +1,159 @@
+# Copyright 2025 The EasyDeL/ejKernel Author @erfanzar (Erfan Zare Chavoshi).
+#
+# 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
+#
+#     https://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.
+
+"""DeepSeek Sparse Attention (DSA) — Pallas TPU interface.
+
+Two-phase implementation:
+  Phase 1 (XLA): Lightning Indexer selects top-k tokens.
+  Phase 2 (Pallas): Tiled MLA attention with sparse mask.
+
+Reference:
+    DeepSeek-V3.2: https://arxiv.org/abs/2512.02556
+"""
+
+from __future__ import annotations
+
+import math
+
+import jax.numpy as jnp
+import jaxtyping
+from beartype import beartype
+from jaxtyping import Array, Float
+
+from ...._registry import Backend, Platform, kernel_registry
+from ...._xla.deepseek_attn._xla_impl_fwd import _lightning_indexer
+from ._pallas_impl_fwd import (
+    ROPE_DECOUPLED,
+    ROPE_FUSED,
+    ROPE_NONE,
+    _build_sparse_mask,
+    deepseek_attn_pallas_impl,
+)
+
+
+@kernel_registry.register("deepseek_attn", Platform.PALLAS, Backend.TPU)
+@jaxtyping.jaxtyped(typechecker=beartype)
+def deepseek_attn(
+    query: Float[Array, "batch seq_len q_heads q_head_dim"],
+    key_value: Float[Array, "batch seq_len kv_lora_rank"],
+    w_kc: Float[Array, "kv_lora_rank kv_heads qk_nope_head_dim"],
+    w_vc: Float[Array, "kv_lora_rank kv_heads v_head_dim"],
+    query_index: Float[Array, "batch seq_len index_heads index_head_dim"],
+    key_index: Float[Array, "batch seq_len index_head_dim"],
+    index_weights: Float[Array, "batch seq_len index_heads"],
+    index_topk: int = 2048,
+    softmax_scale: float | None = None,
+    index_softmax_scale: float | None = None,
+    b_q: Float[Array, "batch seq_len qk_rope_head_dim"] | None = None,
+    b_k: Float[Array, "batch seq_len qk_rope_head_dim"] | None = None,
+    causal: bool = True,
+) -> Float[Array, "batch seq_len q_heads v_head_dim"]:
+    """DeepSeek Sparse Attention on TPU using Pallas.
+
+    Phase 1 runs the Lightning Indexer (XLA) to select top-k tokens.
+    Phase 2 runs a tiled Pallas MLA kernel with the sparse mask applied.
+
+    Args:
+        query: Query tensor [batch, seq_len, q_heads, q_head_dim].
+        key_value: Compressed KV latent [batch, seq_len, kv_lora_rank].
+        w_kc: Key projection [kv_lora_rank, kv_heads, qk_nope_head_dim].
+        w_vc: Value projection [kv_lora_rank, kv_heads, v_head_dim].
+        query_index: Indexer queries [batch, seq_len, index_heads, index_head_dim].
+        key_index: Indexer keys [batch, seq_len, index_head_dim].
+        index_weights: Per-head weights [batch, seq_len, index_heads].
+        index_topk: Number of tokens to select (default: 2048).
+        softmax_scale: Attention scale (default: effective_head_dim^-0.5).
+        index_softmax_scale: Indexer scale (default: index_head_dim^-0.5).
+        b_q: Optional query RoPE [batch, seq_len, qk_rope_head_dim].
+        b_k: Optional key RoPE [batch, seq_len, qk_rope_head_dim].
+        causal: Whether to apply causal masking (default: True).
+
+    Returns:
+        Attention output [batch, seq_len, q_heads, v_head_dim].
+    """
+    _batch, seq_q, q_heads, q_head_dim = query.shape
+    kv_len = key_value.shape[1]
+    kv_heads = int(w_kc.shape[1])
+    d_nope = int(w_kc.shape[2])
+
+    if q_heads % kv_heads != 0:
+        raise ValueError(f"q_heads ({q_heads}) must be divisible by kv_heads ({kv_heads}).")
+
+    if b_k is None:
+        rope_mode = ROPE_NONE
+        if q_head_dim != d_nope:
+            raise ValueError(f"No RoPE: query head_dim ({q_head_dim}) must equal d_nope ({d_nope}).")
+        effective_dim = d_nope
+    elif b_q is None:
+        rope_mode = ROPE_FUSED
+        rope_dim = int(b_k.shape[2])
+        if q_head_dim != d_nope + rope_dim:
+            raise ValueError(
+                f"Fused RoPE: query head_dim ({q_head_dim}) must equal d_nope+rope_dim ({d_nope + rope_dim})."
+            )
+        effective_dim = d_nope + rope_dim
+    else:
+        rope_mode = ROPE_DECOUPLED
+        rope_dim = int(b_k.shape[2])
+        if b_q.shape[2] != rope_dim:
+            raise ValueError(f"b_q/b_k rope_dim mismatch: {b_q.shape[2]} vs {rope_dim}.")
+        effective_dim = d_nope + rope_dim
+
+    if softmax_scale is None:
+        softmax_scale = 1.0 / math.sqrt(effective_dim)
+
+    if index_softmax_scale is None:
+        index_softmax_scale = float(query_index.shape[-1] ** -0.5)
+
+    topk_indices = _lightning_indexer(
+        query_index=query_index,
+        key_index=key_index,
+        index_weights=index_weights,
+        index_topk=index_topk,
+        softmax_scale=index_softmax_scale,
+        causal=causal,
+    )
+
+    sparse_mask = _build_sparse_mask(topk_indices, kv_len)
+
+    block_q = min(128, seq_q)
+    block_k = min(128, kv_len)
+    block_b = 1
+
+    while seq_q % block_q != 0 and block_q > 64:
+        block_q //= 2
+    while kv_len % block_k != 0 and block_k > 64:
+        block_k //= 2
+
+    q_t = query.transpose(0, 2, 1, 3)
+    w_kc_t = jnp.transpose(w_kc, (1, 0, 2))
+    w_vc_t = jnp.transpose(w_vc, (1, 0, 2))
+
+    o = deepseek_attn_pallas_impl(
+        q_t,
+        key_value,
+        w_kc_t,
+        w_vc_t,
+        b_q,
+        b_k,
+        sparse_mask,
+        rope_mode,
+        causal,
+        softmax_scale,
+        block_b,
+        block_q,
+        block_k,
+    )
+
+    return o.transpose(0, 2, 1, 3)

ejkernel-0.0.75/ejkernel/kernels/_pallas/tpu/deepseek_attn/_pallas_impl_bwd.py

@@ -0,0 +1,97 @@
+# Copyright 2025 The EasyDeL/ejKernel Author @erfanzar (Erfan Zare Chavoshi).
+#
+# 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
+#
+#     https://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.
+
+"""Backward utilities for DeepSeek Sparse Attention on TPU Pallas."""
+
+from __future__ import annotations
+
+import jax.numpy as jnp
+
+from ..flash_mla._pallas_impl_bwd import _flash_mla_bwd_impl
+from ..flash_mla._utils import DEFAULT_MASK_VALUE
+
+
+def _sparse_mask_to_bias(sparse_mask):
+    """Convert a float keep-mask into an additive bias for Flash MLA backward.
+
+    Same logic as the forward ``_sparse_mask_to_bias`` — kept as a
+    separate copy to avoid a circular import between fwd and bwd modules.
+
+    Args:
+        sparse_mask: Float32 keep-mask [batch, seq_q, kv_len].
+
+    Returns:
+        Additive bias [batch, 1, seq_q, kv_len].
+    """
+    mask_value = jnp.asarray(DEFAULT_MASK_VALUE, dtype=jnp.float32)
+    return jnp.where(sparse_mask[:, None, :, :] > 0.5, 0.0, mask_value)
+
+
+def _deepseek_attn_bwd_impl(
+    rope_mode,
+    causal,
+    softmax_scale,
+    block_q,
+    block_k,
+    residuals,
+    do,
+):
+    """Compute DeepSeek TPU gradients using the Flash-MLA backward kernels.
+
+    Converts the keep-mask to an additive bias and delegates to
+    ``_flash_mla_bwd_impl``.  The bias gradient from Flash-MLA is
+    discarded and replaced with ``zeros_like(sparse_mask)`` because
+    the top-k selection is non-differentiable.
+
+    Args:
+        rope_mode: RoPE mode constant (ROPE_NONE/FUSED/DECOUPLED).
+        causal: Whether causal masking was applied.
+        softmax_scale: Attention scaling factor.
+        block_q: Query block size for the backward kernel.
+        block_k: Key block size for the backward kernel.
+        residuals: Tuple ``(q, kv_latent, w_kc, w_vc, b_q, b_k,
+            sparse_mask, o, l, m)`` saved during forward.
+        do: Upstream gradient of the attention output.
+
+    Returns:
+        Seven-element tuple of gradients: ``(dq, dkv_latent, dw_kc, dw_vc,
+        db_q, db_k, d_sparse_mask)`` where ``d_sparse_mask`` is always zeros.
+    """
+    q, kv_latent, w_kc, w_vc, b_q, b_k, sparse_mask, o, l, m = residuals
+    bias = _sparse_mask_to_bias(sparse_mask)
+
+    dq, dkv_latent, dw_kc, dw_vc, db_q, db_k, _dbias = _flash_mla_bwd_impl(
+        rope_mode,
+        causal,
+        softmax_scale,
+        None,
+        None,
+        block_q,
+        block_k,
+        (q, kv_latent, w_kc, w_vc, b_q, b_k, bias, o, l, m),
+        do,
+    )
+
+    return (
+        dq,
+        dkv_latent,
+        dw_kc,
+        dw_vc,
+        db_q,
+        db_k,
+        jnp.zeros_like(sparse_mask),
+    )
+
+
+__all__ = ("_deepseek_attn_bwd_impl",)

ejkernel-0.0.75/ejkernel/kernels/_pallas/tpu/deepseek_attn/_pallas_impl_fwd.py

@@ -0,0 +1,259 @@
+# Copyright 2025 The EasyDeL/ejKernel Author @erfanzar (Erfan Zare Chavoshi).
+#
+# 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
+#
+#     https://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.
+
+"""DeepSeek Sparse Attention forward pass and custom VJP wiring for TPU Pallas."""
+
+from __future__ import annotations
+
+import functools
+
+import jax
+import jax.numpy as jnp
+
+from ..flash_mla._pallas_impl_fwd import _flash_mla_pallas_call
+from ..flash_mla._utils import DEFAULT_MASK_VALUE
+from ._pallas_impl_bwd import _deepseek_attn_bwd_impl
+
+ROPE_NONE = 0
+ROPE_FUSED = 1
+ROPE_DECOUPLED = 2
+
+
+def _build_sparse_mask(
+    topk_indices: jnp.ndarray,
+    seq_len: int,
+) -> jnp.ndarray:
+    """Build a dense float32 keep-mask from top-k indices.
+
+    Sets 1.0 at each selected position and guarantees the diagonal
+    (self-attend) is always 1.0 regardless of the indexer output.
+
+    Args:
+        topk_indices: Token indices selected by the Lightning Indexer
+            [batch, seq_q, topk].
+        seq_len: Full KV sequence length (width of the output mask).
+
+    Returns:
+        Float32 mask [batch, seq_q, seq_len] with 1.0 at selected positions.
+    """
+    batch, seq_q, _topk = topk_indices.shape
+    mask = jnp.zeros((batch, seq_q, seq_len), dtype=jnp.float32)
+    mask = mask.at[
+        jnp.arange(batch)[:, None, None],
+        jnp.arange(seq_q)[None, :, None],
+        topk_indices,
+    ].set(1.0)
+    self_mask = jnp.eye(seq_len, dtype=jnp.float32)[None, :, :]
+    return jnp.maximum(mask, self_mask)
+
+
+def _sparse_mask_to_bias(sparse_mask: jnp.ndarray) -> jnp.ndarray:
+    """Convert a float keep-mask into the additive bias Flash MLA expects.
+
+    Positions with mask > 0.5 get bias 0.0 (attend); positions with
+    mask <= 0.5 get ``DEFAULT_MASK_VALUE`` (masked out before softmax).
+
+    Args:
+        sparse_mask: Float32 keep-mask [batch, seq_q, kv_len].
+
+    Returns:
+        Additive bias [batch, 1, seq_q, kv_len] broadcastable over heads.
+    """
+    mask_value = jnp.asarray(DEFAULT_MASK_VALUE, dtype=jnp.float32)
+    return jnp.where(sparse_mask[:, None, :, :] > 0.5, 0.0, mask_value)
+
+
+def _deepseek_attn_pallas_call(
+    q,
+    kv_latent,
+    w_kc,
+    w_vc,
+    b_q,
+    b_k,
+    sparse_mask,
+    *,
+    rope_mode,
+    causal,
+    softmax_scale,
+    block_b,
+    block_q,
+    block_k,
+    save_residuals: bool = False,
+):
+    """Run the DeepSeek TPU forward pass via the raw Flash-MLA Pallas call.
+
+    Converts the float keep-mask into an additive bias (with
+    ``stop_gradient``) and delegates to ``_flash_mla_pallas_call``.
+
+    Args:
+        q: Query tensor [batch, q_heads, seq_q, q_head_dim].
+        kv_latent: Compressed KV latent [batch, kv_len, kv_lora_rank].
+        w_kc: Key projection [kv_heads, kv_lora_rank, d_nope].
+        w_vc: Value projection [kv_heads, kv_lora_rank, v_head_dim].
+        b_q: Optional query RoPE [batch, seq_q, rope_dim].
+        b_k: Optional key RoPE [batch, kv_len, rope_dim].
+        sparse_mask: Float32 keep-mask [batch, seq_q, kv_len].
+        rope_mode: RoPE mode constant (ROPE_NONE/FUSED/DECOUPLED).
+        causal: Whether to apply causal masking.
+        softmax_scale: Attention scaling factor.
+        block_b: Batch block size for Pallas grid.
+        block_q: Query block size for Pallas grid.
+        block_k: Key block size for Pallas grid.
+        save_residuals: If True, also return (logsumexp, max) for backward.
+
+    Returns:
+        Attention output, or (output, logsumexp, max) when ``save_residuals=True``.
+    """
+    bias = _sparse_mask_to_bias(jax.lax.stop_gradient(sparse_mask))
+    return _flash_mla_pallas_call(
+        q,
+        kv_latent,
+        w_kc,
+        w_vc,
+        b_q,
+        b_k,
+        bias,
+        rope_mode=rope_mode,
+        causal=causal,
+        softmax_scale=softmax_scale,
+        sliding_window=None,
+        logits_soft_cap=None,
+        block_b=block_b,
+        block_q=block_q,
+        block_k=block_k,
+        save_residuals=save_residuals,
+    )
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=range(7, 13))
+def deepseek_attn_pallas_impl(
+    q,  # [batch, q_heads, seq_q, q_head_dim]
+    kv_latent,  # [batch, kv_len, kv_lora_rank]
+    w_kc,  # [kv_heads, kv_lora_rank, d_nope]
+    w_vc,  # [kv_heads, kv_lora_rank, v_head_dim]
+    b_q,  # [batch, seq_q, rope_dim] or None
+    b_k,  # [batch, kv_len, rope_dim] or None
+    sparse_mask,  # [batch, seq_q, kv_len] — float32, 1.0=selected
+    rope_mode,
+    causal,
+    softmax_scale,
+    block_b,
+    block_q,
+    block_k,
+):
+    """DeepSeek Sparse Attention forward with explicit TPU ``custom_vjp``.
+
+    Wraps ``_deepseek_attn_pallas_call`` so that the backward pass uses
+    the Flash-MLA Pallas backward kernels with the sparse bias.
+
+    Non-diff args (indices 7-12): ``rope_mode``, ``causal``,
+    ``softmax_scale``, ``block_b``, ``block_q``, ``block_k``.
+    """
+    return _deepseek_attn_pallas_call(
+        q,
+        kv_latent,
+        w_kc,
+        w_vc,
+        b_q,
+        b_k,
+        sparse_mask,
+        rope_mode=rope_mode,
+        causal=causal,
+        softmax_scale=softmax_scale,
+        block_b=block_b,
+        block_q=block_q,
+        block_k=block_k,
+    )
+
+
+def _deepseek_attn_fwd(
+    q,
+    kv_latent,
+    w_kc,
+    w_vc,
+    b_q,
+    b_k,
+    sparse_mask,
+    rope_mode,
+    causal,
+    softmax_scale,
+    block_b,
+    block_q,
+    block_k,
+):
+    """Forward rule for ``custom_vjp``: run raw forward and save residuals.
+
+    Saves ``(q, kv_latent, w_kc, w_vc, b_q, b_k, sparse_mask, o, l, m)``
+    where ``l`` and ``m`` are the per-row logsumexp and max from Flash-MLA.
+    """
+    o, l, m = _deepseek_attn_pallas_call(
+        q,
+        kv_latent,
+        w_kc,
+        w_vc,
+        b_q,
+        b_k,
+        sparse_mask,
+        rope_mode=rope_mode,
+        causal=causal,
+        softmax_scale=softmax_scale,
+        block_b=block_b,
+        block_q=block_q,
+        block_k=block_k,
+        save_residuals=True,
+    )
+    residuals = (q, kv_latent, w_kc, w_vc, b_q, b_k, sparse_mask, o, l, m)
+    return o, residuals
+
+
+def _deepseek_attn_bwd(
+    rope_mode,
+    causal,
+    softmax_scale,
+    block_b,
+    block_q,
+    block_k,
+    residuals,
+    do,
+):
+    """Backward rule for ``custom_vjp``.
+
+    Delegates to ``_deepseek_attn_bwd_impl`` which converts the sparse mask
+    back to an additive bias and calls the Flash-MLA backward kernels.
+    Returns gradients for all seven differentiable inputs; the sparse-mask
+    gradient is always zeros.
+    """
+    del block_b
+    return _deepseek_attn_bwd_impl(
+        rope_mode,
+        causal,
+        softmax_scale,
+        block_q,
+        block_k,
+        residuals,
+        do,
+    )
+
+
+deepseek_attn_pallas_impl.defvjp(_deepseek_attn_fwd, _deepseek_attn_bwd)
+
+
+__all__ = (
+    "ROPE_DECOUPLED",
+    "ROPE_FUSED",
+    "ROPE_NONE",
+    "_build_sparse_mask",
+    "_sparse_mask_to_bias",
+    "deepseek_attn_pallas_impl",
+)

{ejkernel-0.0.72 → ejkernel-0.0.75}/ejkernel/kernels/_pallas/tpu/gated_delta_rule/_interface.py

@@ -36,7 +36,7 @@ def gated_delta_rule(
     beta: Float[Array, "batch seq_len num_heads"],
     decay: Float[Array, "batch seq_len num_heads"] | None = None,
     *,
-    chunk_size: int = 64,
+    chunk_size: int = 256,
     initial_state: Float[Array, "batch num_heads qk_head_dim v_head_dim"] | None = None,
     use_qk_l2norm: bool = True,
     use_chunked: bool = True,