Kernels:

kernels-community
/

fla

Trusted publisher

fla / build /torch-cuda /utils /layers.py

Uploaded using `kernel-builder`.

e541093 verified about 2 months ago

12.5 kB

# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
#   https://github.com/fla-org/flash-linear-attention/graphs/contributors

# Code is adapted from flash-attn.bert_padding.py

# NOTE: This is a copy of `fla.layers.utils` but we do not sync the FLA layers in hf-kernels

import torch
from einops import rearrange, repeat

from ..ops.utils.index import (
    prepare_cu_seqlens_from_lens,
    prepare_cu_seqlens_from_mask,
    prepare_lens_from_mask,
)
from . import tensor_cache

_LAYER_IDX_REQUIRED_MSG = "{cls} requires `layer_idx` when `past_key_values` is provided."


class IndexFirstAxis(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x, indices):
        ctx.save_for_backward(indices)
        assert x.ndim >= 2
        ctx.first_axis_dim, other_shape = x.shape[0], x.shape[1:]
        second_dim = other_shape.numel()
        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
        # return x[indices]
        return torch.gather(
            rearrange(x, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim),
        ).reshape(-1, *other_shape)

    @staticmethod
    def backward(ctx, do):
        (indices,) = ctx.saved_tensors
        assert do.ndim >= 2
        other_shape = do.shape[1:]
        do = rearrange(do, "b ... -> b (...)")
        dx = torch.zeros(
            [ctx.first_axis_dim, do.shape[1]],
            device=do.device,
            dtype=do.dtype,
        )
        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
        # dx[indices] = do
        dx.scatter_(0, repeat(indices, "z -> z d", d=do.shape[1]), do)
        return dx.reshape(ctx.first_axis_dim, *other_shape), None


index_first_axis = IndexFirstAxis.apply


class IndexPutFirstAxis(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x, indices, first_axis_dim):
        ctx.save_for_backward(indices)
        assert indices.ndim == 1
        assert x.ndim >= 2
        y = torch.zeros(first_axis_dim, *x.shape[1:], device=x.device, dtype=x.dtype)
        # TODO [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
        y[indices] = x
        # y.scatter_(0, repeat(indices, 'z -> z d', d=x.shape[1]), x)
        return y

    @staticmethod
    def backward(ctx, do):
        (indices,) = ctx.saved_tensors
        # TODO [2022-03-04] For some reason torch.gather is a bit faster than indexing.
        dx = do[indices]
        # dx = torch.gather(do, 0, repeat(indices, 'z -> z d', d=do.shape[1]))
        return dx, None, None


index_put_first_axis = IndexPutFirstAxis.apply


@tensor_cache
def _get_unpad_indices_and_cu(
    attention_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    lens = prepare_lens_from_mask(attention_mask)
    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
    cu_seqlens = prepare_cu_seqlens_from_lens(lens)
    return indices, cu_seqlens


@tensor_cache
def _get_last_token_unpad_indices_and_cu(
    attention_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    return _get_unpad_indices_and_cu(attention_mask[:, -1:])


def get_unpad_indices_and_cu(
    attention_mask: torch.Tensor,
    q_len: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return unpadding indices and cumulative sequence lengths."""
    if q_len is None or attention_mask.shape[-1] == q_len:
        return _get_unpad_indices_and_cu(attention_mask)
    if q_len == 1:
        return _get_last_token_unpad_indices_and_cu(attention_mask)
    return _get_unpad_indices_and_cu(attention_mask[:, -q_len:])


@tensor_cache
def get_unpad_data(
    attention_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, int]:
    """
    Retrieves indexing data required to repad unpadded (ragged) tensors.
    Args:
        attention_mask (`torch.Tensor`):
            Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.
    Return:
        indices (`torch.Tensor`):
            The indices of non-masked tokens from the flattened input sequence.
        cu_seqlens (`torch.Tensor`):
            The cumulative sequence lengths, used to index into ragged (unpadded) tensors.
            `cu_seqlens` shape is [batch_size + 1].
        max_seqlen_in_batch (`int`):
            Maximum sequence length in batch.
    """
    lens = prepare_lens_from_mask(attention_mask)
    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
    max_seqlen_in_batch = lens.max().item()
    cu_seqlens = prepare_cu_seqlens_from_mask(attention_mask)
    return indices, cu_seqlens, max_seqlen_in_batch


def unpad_hidden_states(
    hidden_states: torch.Tensor,
    cu_seqlens: torch.LongTensor | None,
    attention_mask: torch.Tensor | None,
    q_len: int,
) -> tuple[torch.Tensor, torch.Tensor | None, torch.LongTensor | None]:
    """Unpad hidden states via attention_mask, unless explicit cu_seqlens are given.
    Explicit `cu_seqlens` take precedence: when they are provided, the input is
    returned unchanged and no unpadding happens. Only when `cu_seqlens` is None
    and `attention_mask` is present are the hidden states unpadded via the mask.
    Arguments:
        hidden_states (`torch.Tensor`):
            Input hidden states with padding. Shape: [batch_size, q_len, ...].
        cu_seqlens (`torch.LongTensor | None`):
            Explicit cumulative sequence lengths provided by the caller.
        attention_mask (`torch.Tensor | None`):
            Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.
        q_len (`int`):
            Target length.
    Return:
        hidden_states (`torch.Tensor`):
            The unpadded hidden states of shape [1, total_tokens, ...],
            or the input unchanged when unpadding is skipped.
        indices (`torch.Tensor | None`):
            The indices of non-masked tokens, or None when unpadding is skipped.
        cu_seqlens (`torch.LongTensor | None`):
            The cumulative sequence lengths derived from the mask,
            or the explicit value passed in when unpadding is skipped.
    """
    if cu_seqlens is None and attention_mask is not None:
        indices, cu_seqlens = get_unpad_indices_and_cu(attention_mask, q_len)
        hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
        return hidden_states, indices, cu_seqlens
    return hidden_states, None, cu_seqlens


def repad_hidden_states(
    hidden_states: torch.Tensor,
    indices: torch.Tensor | None,
    batch_size: int,
    q_len: int,
) -> torch.Tensor:
    """Inverse of `unpad_hidden_states`; a no-op when `indices` is None.
    Arguments:
        hidden_states (`torch.Tensor`):
            Hidden states of shape [1, total_tokens, ...].
        indices (`torch.Tensor | None`):
            The indices returned by `unpad_hidden_states`, or None.
        batch_size (`int`):
            Batch size for the padded sequence.
        q_len (`int`):
            Target length.
    Return:
        hidden_states of shape [batch_size, q_len, ...],
        or the input unchanged when `indices` is None.
    """
    if indices is None:
        return hidden_states
    return pad_input(hidden_states.squeeze(0), indices, batch_size, q_len)


def unpad_input(
    q: torch.Tensor | tuple[torch.Tensor, ...],
    states: tuple[torch.Tensor],
    attention_mask: torch.Tensor,
    q_len: int,
    keepdim: bool = False,
):
    """
    Unpads query, key, and values tensors, using a single dimension for all tokens
    even though they belong to different batches.
    Arguments:
        q (`torch.Tensor` or `Tuple[torch.Tensor]`):
            Query state with padding. Shape: [batch_size, q_len, ...].
            When it is a tuple, do unpadding for each tensor in the tuple.
        states (`Tuple[torch.Tensor]`):
            Attention state with padding. Shape: [batch_size, seq_len, ...].
        attention_mask (`torch.Tensor`):
            Boolean or int tensor of shape [batch_size, sequence_length], 1 means valid and 0 means not valid.
        q_len (`int`):
            Target length.
        keepdim (`bool`):
            Whether to keep the batch dimension. Default: `False`.
    Return:
        q (`torch.Tensor` or `Tuple[torch.Tensor]`):
            Query state without padding.
            Shape: [1, total_target_length, ...] if `keepdim=True` else [total_target_length, ...].
            When the `q` passed in is a tuple, return a tuple of such unpadded tensors.
        states (`Tuple[torch.Tensor]`):
            Attention state without padding.
            Shape: [1, total_source_length, ...] if `keepdim=True` else [total_source_length, ...].
        indices_q (`torch.Tensor`):
            The indices of non-masked tokens from the flattened input target sequence.
        (cu_seqlens_q, cu_seqlens_k) (`Tuple[torch.LongTensor, torch.LongTensor]`):
            The cumulative sequence lengths for the target (query) and source (key, value),
            used to index into ragged (unpadded) tensors.
            `cu_seqlens` shape is [batch_size + 1].
        (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`Tuple[int, int]`):
            Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence
            i.e. query, `max_seqlen_in_batch_k` for the source sequence i.e. key/value).
    """
    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = get_unpad_data(attention_mask)
    batch_size, seq_len, *_ = states[0].shape

    state = tuple(
        index_first_axis(rearrange(s, "b s ... -> (b s) ..."), indices_k)
        for s in states
    )
    if isinstance(q, torch.Tensor):
        q = (q,)
        cast_tuple = True
    else:
        cast_tuple = False

    if q_len == seq_len:
        q = tuple(index_first_axis(rearrange(q_, "b s ... -> (b s) ..."), indices_k) for q_ in q)
        cu_seqlens_q = cu_seqlens_k
        max_seqlen_in_batch_q = max_seqlen_in_batch_k
        indices_q = indices_k
    elif q_len == 1:
        max_seqlen_in_batch_q = 1
        cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=q[0].device)
        indices_q = cu_seqlens_q[:-1]
        q = tuple(q_.squeeze(1) for q_ in q)
    else:
        raise NotImplementedError("We only support either q_len == k_len (prefilling) or q_len == 1 (decoding)")

    if keepdim:
        q = tuple(q_.unsqueeze(0) for q_ in q)
        state = tuple(s.unsqueeze(0) for s in state)
    if cast_tuple:
        q = q[0]

    return (
        q,
        state,
        indices_q,
        (cu_seqlens_q, cu_seqlens_k),
        (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
    )


def pad_input(
    hidden_states: torch.Tensor,
    indices: torch.LongTensor,
    batch_size: int,
    seq_len: int,
) -> torch.Tensor:
    """
    Args:
        hidden_states ([total_tokens, ...]):
            where total_tokens denotes the number of tokens in selected in attention_mask.
        indices ([total_tokens]):
            the indices that represent the non-masked tokens of the original padded input sequence.
        batch_size (int):
            batch_size size for the padded sequence.
        seq_len (int):
            maximum sequence length for the padded sequence.
    Return:
        hidden_states of shape [batch_size, seq_len, ...]
    """
    output = index_put_first_axis(hidden_states, indices, batch_size * seq_len)
    return rearrange(output, "(b s) ... -> b s ...", b=batch_size)


def require_cache_layer_idx(module, past_key_values):
    layer_idx = getattr(module, "layer_idx", None)
    if past_key_values is not None and layer_idx is None:
        raise ValueError(_LAYER_IDX_REQUIRED_MSG.format(cls=module.__class__.__name__))
    return layer_idx


def get_layer_cache(module, past_key_values):
    layer_idx = require_cache_layer_idx(module, past_key_values)
    if past_key_values is not None and len(past_key_values) > layer_idx:
        return past_key_values[layer_idx]
    return None


def update_layer_cache(module, past_key_values, **kwargs):
    layer_idx = require_cache_layer_idx(module, past_key_values)
    if past_key_values is not None:
        return past_key_values.update(layer_idx=layer_idx, **kwargs)
    return None

	# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
	#
	# This source code is licensed under the MIT license found in the
	# LICENSE file in the root directory of this source tree.
	# For a list of all contributors, visit:
	# https://github.com/fla-org/flash-linear-attention/graphs/contributors

	# Code is adapted from flash-attn.bert_padding.py

	# NOTE: This is a copy of `fla.layers.utils` but we do not sync the FLA layers in hf-kernels

	import torch
	from einops import rearrange, repeat

	from ..ops.utils.index import (
	prepare_cu_seqlens_from_lens,
	prepare_cu_seqlens_from_mask,
	prepare_lens_from_mask,
	)
	from . import tensor_cache

	_LAYER_IDX_REQUIRED_MSG = "{cls} requires `layer_idx` when `past_key_values` is provided."


	class IndexFirstAxis(torch.autograd.Function):

	@staticmethod
	def forward(ctx, x, indices):
	ctx.save_for_backward(indices)
	assert x.ndim >= 2
	ctx.first_axis_dim, other_shape = x.shape[0], x.shape[1:]
	second_dim = other_shape.numel()
	# TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
	# return x[indices]
	return torch.gather(
	rearrange(x, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim),
	).reshape(-1, *other_shape)

	@staticmethod
	def backward(ctx, do):
	(indices,) = ctx.saved_tensors
	assert do.ndim >= 2
	other_shape = do.shape[1:]
	do = rearrange(do, "b ... -> b (...)")
	dx = torch.zeros(
	[ctx.first_axis_dim, do.shape[1]],
	device=do.device,
	dtype=do.dtype,
	)
	# TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
	# dx[indices] = do
	dx.scatter_(0, repeat(indices, "z -> z d", d=do.shape[1]), do)
	return dx.reshape(ctx.first_axis_dim, *other_shape), None


	index_first_axis = IndexFirstAxis.apply


	class IndexPutFirstAxis(torch.autograd.Function):

	@staticmethod
	def forward(ctx, x, indices, first_axis_dim):
	ctx.save_for_backward(indices)
	assert indices.ndim == 1
	assert x.ndim >= 2
	y = torch.zeros(first_axis_dim, *x.shape[1:], device=x.device, dtype=x.dtype)
	# TODO [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
	y[indices] = x
	# y.scatter_(0, repeat(indices, 'z -> z d', d=x.shape[1]), x)
	return y

	@staticmethod
	def backward(ctx, do):
	(indices,) = ctx.saved_tensors
	# TODO [2022-03-04] For some reason torch.gather is a bit faster than indexing.
	dx = do[indices]
	# dx = torch.gather(do, 0, repeat(indices, 'z -> z d', d=do.shape[1]))
	return dx, None, None


	index_put_first_axis = IndexPutFirstAxis.apply


	@tensor_cache
	def _get_unpad_indices_and_cu(
	attention_mask: torch.Tensor,
	) -> tuple[torch.Tensor, torch.Tensor]:
	lens = prepare_lens_from_mask(attention_mask)
	indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
	cu_seqlens = prepare_cu_seqlens_from_lens(lens)
	return indices, cu_seqlens


	@tensor_cache
	def _get_last_token_unpad_indices_and_cu(
	attention_mask: torch.Tensor,
	) -> tuple[torch.Tensor, torch.Tensor]:
	return _get_unpad_indices_and_cu(attention_mask[:, -1:])


	def get_unpad_indices_and_cu(
	attention_mask: torch.Tensor,
	q_len: int \| None = None,
	) -> tuple[torch.Tensor, torch.Tensor]:
	"""Return unpadding indices and cumulative sequence lengths."""
	if q_len is None or attention_mask.shape[-1] == q_len:
	return _get_unpad_indices_and_cu(attention_mask)
	if q_len == 1:
	return _get_last_token_unpad_indices_and_cu(attention_mask)
	return _get_unpad_indices_and_cu(attention_mask[:, -q_len:])


	@tensor_cache
	def get_unpad_data(
	attention_mask: torch.Tensor,
	) -> tuple[torch.Tensor, torch.Tensor, int]:
	"""
	Retrieves indexing data required to repad unpadded (ragged) tensors.

	Args:
	attention_mask (`torch.Tensor`):
	Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.

	Return:
	indices (`torch.Tensor`):
	The indices of non-masked tokens from the flattened input sequence.
	cu_seqlens (`torch.Tensor`):
	The cumulative sequence lengths, used to index into ragged (unpadded) tensors.
	`cu_seqlens` shape is [batch_size + 1].
	max_seqlen_in_batch (`int`):
	Maximum sequence length in batch.
	"""
	lens = prepare_lens_from_mask(attention_mask)
	indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
	max_seqlen_in_batch = lens.max().item()
	cu_seqlens = prepare_cu_seqlens_from_mask(attention_mask)
	return indices, cu_seqlens, max_seqlen_in_batch


	def unpad_hidden_states(
	hidden_states: torch.Tensor,
	cu_seqlens: torch.LongTensor \| None,
	attention_mask: torch.Tensor \| None,
	q_len: int,
	) -> tuple[torch.Tensor, torch.Tensor \| None, torch.LongTensor \| None]:
	"""Unpad hidden states via attention_mask, unless explicit cu_seqlens are given.

	Explicit `cu_seqlens` take precedence: when they are provided, the input is
	returned unchanged and no unpadding happens. Only when `cu_seqlens` is None
	and `attention_mask` is present are the hidden states unpadded via the mask.

	Arguments:
	hidden_states (`torch.Tensor`):
	Input hidden states with padding. Shape: [batch_size, q_len, ...].
	cu_seqlens (`torch.LongTensor \| None`):
	Explicit cumulative sequence lengths provided by the caller.
	attention_mask (`torch.Tensor \| None`):
	Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.
	q_len (`int`):
	Target length.

	Return:
	hidden_states (`torch.Tensor`):
	The unpadded hidden states of shape [1, total_tokens, ...],
	or the input unchanged when unpadding is skipped.
	indices (`torch.Tensor \| None`):
	The indices of non-masked tokens, or None when unpadding is skipped.
	cu_seqlens (`torch.LongTensor \| None`):
	The cumulative sequence lengths derived from the mask,
	or the explicit value passed in when unpadding is skipped.
	"""
	if cu_seqlens is None and attention_mask is not None:
	indices, cu_seqlens = get_unpad_indices_and_cu(attention_mask, q_len)
	hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
	return hidden_states, indices, cu_seqlens
	return hidden_states, None, cu_seqlens


	def repad_hidden_states(
	hidden_states: torch.Tensor,
	indices: torch.Tensor \| None,
	batch_size: int,
	q_len: int,
	) -> torch.Tensor:
	"""Inverse of `unpad_hidden_states`; a no-op when `indices` is None.

	Arguments:
	hidden_states (`torch.Tensor`):
	Hidden states of shape [1, total_tokens, ...].
	indices (`torch.Tensor \| None`):
	The indices returned by `unpad_hidden_states`, or None.
	batch_size (`int`):
	Batch size for the padded sequence.
	q_len (`int`):
	Target length.

	Return:
	hidden_states of shape [batch_size, q_len, ...],
	or the input unchanged when `indices` is None.
	"""
	if indices is None:
	return hidden_states
	return pad_input(hidden_states.squeeze(0), indices, batch_size, q_len)


	def unpad_input(
	q: torch.Tensor \| tuple[torch.Tensor, ...],
	states: tuple[torch.Tensor],
	attention_mask: torch.Tensor,
	q_len: int,
	keepdim: bool = False,
	):
	"""
	Unpads query, key, and values tensors, using a single dimension for all tokens
	even though they belong to different batches.


	Arguments:
	q (`torch.Tensor` or `Tuple[torch.Tensor]`):
	Query state with padding. Shape: [batch_size, q_len, ...].
	When it is a tuple, do unpadding for each tensor in the tuple.
	states (`Tuple[torch.Tensor]`):
	Attention state with padding. Shape: [batch_size, seq_len, ...].
	attention_mask (`torch.Tensor`):
	Boolean or int tensor of shape [batch_size, sequence_length], 1 means valid and 0 means not valid.
	q_len (`int`):
	Target length.
	keepdim (`bool`):
	Whether to keep the batch dimension. Default: `False`.

	Return:
	q (`torch.Tensor` or `Tuple[torch.Tensor]`):
	Query state without padding.
	Shape: [1, total_target_length, ...] if `keepdim=True` else [total_target_length, ...].
	When the `q` passed in is a tuple, return a tuple of such unpadded tensors.
	states (`Tuple[torch.Tensor]`):
	Attention state without padding.
	Shape: [1, total_source_length, ...] if `keepdim=True` else [total_source_length, ...].
	indices_q (`torch.Tensor`):
	The indices of non-masked tokens from the flattened input target sequence.
	(cu_seqlens_q, cu_seqlens_k) (`Tuple[torch.LongTensor, torch.LongTensor]`):
	The cumulative sequence lengths for the target (query) and source (key, value),
	used to index into ragged (unpadded) tensors.
	`cu_seqlens` shape is [batch_size + 1].
	(max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`Tuple[int, int]`):
	Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence
	i.e. query, `max_seqlen_in_batch_k` for the source sequence i.e. key/value).
	"""
	indices_k, cu_seqlens_k, max_seqlen_in_batch_k = get_unpad_data(attention_mask)
	batch_size, seq_len, *_ = states[0].shape

	state = tuple(
	index_first_axis(rearrange(s, "b s ... -> (b s) ..."), indices_k)
	for s in states
	)
	if isinstance(q, torch.Tensor):
	q = (q,)
	cast_tuple = True
	else:
	cast_tuple = False

	if q_len == seq_len:
	q = tuple(index_first_axis(rearrange(q_, "b s ... -> (b s) ..."), indices_k) for q_ in q)
	cu_seqlens_q = cu_seqlens_k
	max_seqlen_in_batch_q = max_seqlen_in_batch_k
	indices_q = indices_k
	elif q_len == 1:
	max_seqlen_in_batch_q = 1
	cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=q[0].device)
	indices_q = cu_seqlens_q[:-1]
	q = tuple(q_.squeeze(1) for q_ in q)
	else:
	raise NotImplementedError("We only support either q_len == k_len (prefilling) or q_len == 1 (decoding)")

	if keepdim:
	q = tuple(q_.unsqueeze(0) for q_ in q)
	state = tuple(s.unsqueeze(0) for s in state)
	if cast_tuple:
	q = q[0]

	return (
	q,
	state,
	indices_q,
	(cu_seqlens_q, cu_seqlens_k),
	(max_seqlen_in_batch_q, max_seqlen_in_batch_k),
	)


	def pad_input(
	hidden_states: torch.Tensor,
	indices: torch.LongTensor,
	batch_size: int,
	seq_len: int,
	) -> torch.Tensor:
	"""
	Args:
	hidden_states ([total_tokens, ...]):
	where total_tokens denotes the number of tokens in selected in attention_mask.
	indices ([total_tokens]):
	the indices that represent the non-masked tokens of the original padded input sequence.
	batch_size (int):
	batch_size size for the padded sequence.
	seq_len (int):
	maximum sequence length for the padded sequence.

	Return:
	hidden_states of shape [batch_size, seq_len, ...]
	"""
	output = index_put_first_axis(hidden_states, indices, batch_size * seq_len)
	return rearrange(output, "(b s) ... -> b s ...", b=batch_size)


	def require_cache_layer_idx(module, past_key_values):
	layer_idx = getattr(module, "layer_idx", None)
	if past_key_values is not None and layer_idx is None:
	raise ValueError(_LAYER_IDX_REQUIRED_MSG.format(cls=module.__class__.__name__))
	return layer_idx


	def get_layer_cache(module, past_key_values):
	layer_idx = require_cache_layer_idx(module, past_key_values)
	if past_key_values is not None and len(past_key_values) > layer_idx:
	return past_key_values[layer_idx]
	return None


	def update_layer_cache(module, past_key_values, **kwargs):
	layer_idx = require_cache_layer_idx(module, past_key_values)
	if past_key_values is not None:
	return past_key_values.update(layer_idx=layer_idx, **kwargs)
	return None