Upload 6 files

README.md

@@ -0,0 +1,69 @@
+---
+language:
+- ja
+pipeline_tag: automatic-speech-recognition
+base_model:
+  - efwkjn/cohere-asr-ja
+---
+Experimental single step block diffusion model for speculative decoding similar to [Orthrus](https://arxiv.org/abs/2605.12825)/[DFlash](https://arxiv.org/abs/2602.06036). Like Orthrus it conditions on the AR KV cache but like DFlash only the embedding layer is shared. Training and inference code is included based on pure pytorch/flex attention with torch compile.
+
+|bsz=1  |transformers           |nano-cohere-transcribe |diffusion=False        |diffusion=True         |avg accept |
+|-------|----------------------:|----------------------:|----------------------:|----------------------:|----------:|
+|fleurs |                  1.00x|                  2.35x|                  2.97x|                  5.28x|       8.92|
+|jsut   |                  1.00x|                  2.02x|                  2.72x|                  4.24x|       9.88|
+|reazon |                  1.00x|                  1.89x|                  2.59x|                  3.82x|       8.45|
+
+Notes
+* Diffusion for fun, not very practical (increase batching instead)
+* torch.compile involves more warmup time
+* Pad inputs to static shapes for compile efficiency
+* diffusion=False ~10% faster than nano if encoder isn't compiled
+* Benchmark sequences short, average of jsut/reazon is below block size
+
+```python
+from huggingface_hub import snapshot_download
+model_dir = 'diff'
+snapshot_download('efwkjn/cohere-asr-ja-diffusion', cache_dir=model_dir, local_dir=model_dir)
+```
+```python
+import importlib
+import subprocess
+
+import numpy as np
+import torch
+from tokenizers import Tokenizer
+
+cohere_asr = importlib.import_module(model_dir + '.cohere_asr')
+file = 'audio.wav'
+device = 'cuda'
+
+cmd = [
+    'ffmpeg',
+    '-threads', '1',
+    '-nostdin',
+    '-i', file,
+    '-f', 's16le',
+    '-c:a', 'pcm_s16le',
+    '-ar', '16000',
+    '-ac', '1',
+    '-'
+]
+with subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) as p:
+    audio = np.frombuffer(p.stdout.read(), dtype=np.int16).astype(np.float32) / 0x8000
+
+fe = cohere_asr.CohereAsrFeatureExtractor(model_dir)
+model = cohere_asr.CohereAsr.from_pretrained(model_dir, device)
+tokenizer: Tokenizer = Tokenizer.from_file(model_dir + '/tokenizer.json')
+prompt = tokenizer.encode(
+    '<|startofcontext|><|startoftranscript|><|emo:undefined|>'
+    '<|ja|><|ja|><|pnc|><|noitn|><|notimestamp|><|nodiarize|>',
+    add_special_tokens=False
+).ids
+
+features, lengths = fe([audio[:480000]])
+input_ids = torch.tensor(prompt, device=device)[None, :]
+features = features.to(device, torch.bfloat16)
+lengths = lengths.to(device)
+ids = model.generate(input_ids, features, lengths, compile=False, diffusion=True)
+print(tokenizer.decode_batch(ids.tolist())[0])
+```

cohere_asr.py

@@ -0,0 +1,1769 @@
+import json
+import math
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from safetensors import safe_open
+from torch import nn
+from torch.nn.attention.flex_attention import _DEFAULT_SPARSE_BLOCK_SIZE, BlockMask, flex_attention
+
+
+class FilterbankFeatures(nn.Module):
+    '''Filterbank features extraction module.
+    Follows NeMo's FilterbankFeatures implementation.
+    '''
+
+    window: torch.Tensor
+    fb: torch.Tensor
+
+    def __init__(
+        self,
+        pretrained_model_or_path=None,
+        sample_rate=16000,
+        n_window_size=400,
+        n_window_stride=160,
+        window='hann',
+        normalize='per_feature',
+        n_fft=512,
+        preemph=0.97,
+        nfilt=128,
+        lowfreq=0,
+        highfreq=None,
+        log=True,
+        log_zero_guard_type='add',
+        log_zero_guard_value=2**-24,
+        dither=1e-5,
+        pad_to=0,
+        max_duration=35,
+        frame_splicing=1,
+        exact_pad=False,
+        pad_value=0,
+        mag_power=2.0,
+        use_grads=False,
+        nb_augmentation_prob=0.0,
+        nb_max_freq=4000,
+        mel_norm='slaney',
+        device='cpu',
+    ):
+        super().__init__()
+        if exact_pad and n_window_stride % 2 == 1:
+            raise NotImplementedError(f'{self} received exact_pad=True with odd hop length ({n_window_stride}).')
+
+        if (
+            n_window_size is None
+            or n_window_stride is None
+            or not isinstance(n_window_size, int)
+            or not isinstance(n_window_stride, int)
+            or n_window_size <= 0
+            or n_window_stride <= 0
+        ):
+            raise ValueError('n_window_size and n_window_stride must be positive ints.')
+
+        self.log_zero_guard_value = log_zero_guard_value
+        self.sample_rate = sample_rate
+        self.win_length = n_window_size
+        self.hop_length = n_window_stride
+        self.n_fft = n_fft or 2 ** math.ceil(math.log2(self.win_length))
+        self.stft_pad_amount = (self.n_fft - self.hop_length) // 2 if exact_pad else None
+        self.exact_pad = exact_pad
+        self.max_duration = max_duration
+
+        if pretrained_model_or_path:
+            safetensor_path = Path(pretrained_model_or_path) / 'model.safetensors'
+            if safetensor_path.exists():
+                try:
+                    with safe_open(safetensor_path.as_posix(), 'pt') as file:
+                        self.register_buffer('fb', file.get_tensor('preprocessor.featurizer.fb'))
+                        self.register_buffer('window', file.get_tensor('preprocessor.featurizer.window'))
+                except:
+                    pass
+
+        if not hasattr(self, 'window'):
+            torch_windows = {
+                'hann': torch.hann_window,
+                'hamming': torch.hamming_window,
+                'blackman': torch.blackman_window,
+                'bartlett': torch.bartlett_window,
+                'none': None,
+            }
+            window_fn = torch_windows.get(window)
+            window_tensor = window_fn(self.win_length, periodic=False) if window_fn else None
+            self.register_buffer('window', window_tensor)
+
+        self.normalize = normalize
+        self.log = log
+        self.dither = dither
+        self.frame_splicing = frame_splicing
+        self.nfilt = nfilt
+        self.preemph = preemph
+        self.pad_to = pad_to
+        highfreq = highfreq or sample_rate / 2
+        self.pad_min_duration = 0.0
+        self.pad_direction = 'both'
+        self.pad_value = pad_value
+        self.mag_power = mag_power
+        self.nb_augmentation_prob = nb_augmentation_prob
+
+        if not hasattr(self, 'fb'):
+            self.register_buffer('fb', self._get_mel_filters(lowfreq, highfreq, mel_norm))
+
+        max_length = self.get_seq_len(torch.tensor(max_duration * sample_rate, dtype=torch.float))
+        max_pad = pad_to - (max_length % pad_to) if pad_to > 0 else 0
+        self.max_length = max_length + max_pad
+
+        if log_zero_guard_type not in ['add', 'clamp']:
+            raise ValueError('log_zero_guard_type must be `add` or `clamp`.')
+        self.log_zero_guard_type = log_zero_guard_type
+
+        self.use_grads = use_grads
+        if not use_grads:
+            self.forward = torch.no_grad()(self.forward)
+
+        if self.nb_augmentation_prob > 0.0:
+            if nb_max_freq >= sample_rate / 2:
+                self.nb_augmentation_prob = 0.0
+            else:
+                self._nb_max_fft_bin = int((nb_max_freq / sample_rate) * self.n_fft)
+
+        if self.window is None:
+            raise RuntimeError('Expected a window tensor for STFT feature extraction.')
+        if self.fb is None:
+            raise RuntimeError('Expected mel filterbank weights for feature extraction.')
+        self.window = self.window.to(dtype=torch.float32)
+        self.fb = self.fb.to(dtype=torch.float32)
+        self.generator = torch.Generator(device=device)
+        self.generator.manual_seed(0)
+
+    def _get_mel_filters(self, fmin: float, fmax: float, mel_norm: str):
+        assert mel_norm == 'slaney'
+        min_log_hertz = 1000.0
+        min_log_mel = 15.0
+        logstep = math.log(6.4) / 27.0
+        fftfreqs = torch.fft.rfftfreq(n=self.n_fft, d=1.0 / self.sample_rate, dtype=torch.float64)
+
+        def hz_to_mel(freq: float) -> float:
+            if freq < min_log_hertz:
+                return 3 * freq / 200.0
+            return min_log_mel + math.log(freq / min_log_hertz) * (1 / logstep)
+        mels = torch.linspace(hz_to_mel(fmin), hz_to_mel(fmax), self.nfilt + 2, dtype=torch.float64)
+        freqs = 200.0 / 3 * mels
+        log_t = mels >= min_log_mel
+        freqs[log_t] = min_log_hertz * torch.exp(logstep * (mels[log_t] - min_log_mel))
+
+        fdiff = torch.diff(freqs)
+        ramps = freqs[:, None] - fftfreqs[None, :]
+        lower = (-ramps[:-2, :] / fdiff[:-1, None]).float()
+        upper = (ramps[2:, :] / fdiff[1:, None]).float()
+        weights = torch.maximum(torch.minimum(lower, upper), torch.zeros(1))
+        enorm = 2.0 / (freqs[2: self.nfilt + 2] - freqs[:self.nfilt])
+        weights *= enorm[:, None]
+        return weights
+
+    @torch._dynamo.disable
+    def _apply_dither(self, x: torch.Tensor, seq_len_time: torch.Tensor) -> torch.Tensor:
+        '''Apply deterministic per-sample dither outside torch.compile.
+        Each sample is seeded by its valid waveform length so that dither noise
+        is batch-composition invariant (a sample's features depend only on its
+        own content, not on what else is in the batch).
+        '''
+        if self.dither <= 0:
+            return x
+        for i in range(x.shape[0]):
+            valid_samples = min(int(seq_len_time[i].item()), x.shape[1])
+            if valid_samples <= 0:
+                continue
+            self.generator.manual_seed(valid_samples)
+            noise = torch.randn(
+                (valid_samples,),
+                dtype=x.dtype,
+                device=x.device,
+                generator=self.generator,
+            )
+            x[i, :valid_samples] += self.dither * noise
+        return x
+
+    def log_zero_guard_value_fn(self, x: torch.Tensor) -> float:
+        if isinstance(self.log_zero_guard_value, str):
+            if self.log_zero_guard_value == 'tiny':
+                return torch.finfo(x.dtype).tiny
+            if self.log_zero_guard_value == 'eps':
+                return torch.finfo(x.dtype).eps
+            raise ValueError('log_zero_guard_value must be number, `tiny`, or `eps` when str.')
+        return self.log_zero_guard_value
+
+    def get_seq_len(self, seq_len: torch.Tensor) -> torch.Tensor:
+        pad_amount = self.stft_pad_amount * 2 if self.stft_pad_amount is not None else self.n_fft // 2 * 2
+        seq_len = torch.floor_divide((seq_len + pad_amount - self.n_fft), self.hop_length)
+        return seq_len.to(dtype=torch.long)
+
+    def splice_frames(self, x: torch.Tensor, frame_splicing: int) -> torch.Tensor:
+        seq = [x]
+        for n in range(1, frame_splicing):
+            seq.append(torch.cat([x[:, :, :n], x[:, :, n:]], dim=2))
+        return torch.cat(seq, dim=1)
+
+    def normalize_batch(self, x: torch.Tensor, seq_len: torch.Tensor, normalize_type: str) -> torch.Tensor:
+        if normalize_type != 'per_feature':
+            raise ValueError('Only per_feature normalization is supported.')
+        batch_size = x.shape[0]
+        max_time = x.shape[2]
+        time_steps = torch.arange(max_time, device=x.device).unsqueeze(0).expand(batch_size, max_time)
+        valid_mask = time_steps < seq_len.unsqueeze(1)
+        x_mean_num = torch.where(valid_mask.unsqueeze(1), x, 0.0).sum(dim=2)
+        x_mean_den = valid_mask.sum(dim=1)
+        x_mean = x_mean_num / x_mean_den.unsqueeze(1)
+        x_std = torch.sqrt(
+            torch.sum(
+                torch.where(valid_mask.unsqueeze(1), x - x_mean.unsqueeze(2), 0.0) ** 2,
+                dim=2,
+            )
+            / (x_mean_den.unsqueeze(1) - 1.0)
+        )
+        x_std = x_std.masked_fill(x_std.isnan(), 0.0)
+        x_std += 1e-5
+        return (x - x_mean.unsqueeze(2)) / x_std.unsqueeze(2), x_mean, x_std
+
+    def forward(self, x: torch.Tensor, seq_len: torch.Tensor, linear_spec: bool = False) -> tuple[torch.Tensor, torch.Tensor]:
+        if x.shape[1] < self.sample_rate * self.pad_min_duration:
+            pad_amount = int(self.sample_rate * self.pad_min_duration) - x.shape[1]
+            if self.pad_direction == 'right':
+                x = F.pad(x, (0, pad_amount), value=self.pad_value)
+            elif self.pad_direction == 'left':
+                x = F.pad(x, (pad_amount, 0), value=self.pad_value)
+            elif self.pad_direction == 'both':
+                left_pad = pad_amount // 2
+                right_pad = pad_amount - left_pad
+                x = F.pad(x, (left_pad, right_pad), value=self.pad_value)
+            else:
+                raise ValueError(f'Invalid pad_direction: {self.pad_direction}')
+            seq_len = torch.tensor([x.shape[1]], dtype=torch.float, device=x.device)
+
+        seq_len_time = seq_len
+        seq_len_unfixed = self.get_seq_len(seq_len)
+        seq_len = torch.where(seq_len == 0, torch.zeros_like(seq_len_unfixed), seq_len_unfixed)
+
+        if self.stft_pad_amount is not None:
+            x = F.pad(x.unsqueeze(1), (self.stft_pad_amount, self.stft_pad_amount), 'constant').squeeze(1)
+
+        x = self._apply_dither(x, seq_len_time)
+
+        if self.preemph is not None:
+            timemask = torch.arange(x.shape[1], device=x.device).unsqueeze(0) < seq_len_time.unsqueeze(1)
+            x = torch.cat((x[:, 0].unsqueeze(1), x[:, 1:] - self.preemph * x[:, :-1]), dim=1)
+            x = x.masked_fill(~timemask, 0.0)
+
+        x = torch.stft(
+            x,
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            center=not self.exact_pad,
+            window=self.window,
+            return_complex=True,
+            pad_mode='constant',
+        )
+        x = x[:, :, :seq_len.max()]
+        guard = 0 if not self.use_grads else 1e-5
+        x = torch.sqrt(torch.view_as_real(x).pow(2).sum(-1) + guard)
+
+        if self.mag_power != 1.0:
+            x = x.pow(self.mag_power)
+        if linear_spec:
+            return x, seq_len
+
+        x = torch.matmul(self.fb, x)
+
+        if self.log:
+            if self.log_zero_guard_type == 'add':
+                x = torch.log(x + self.log_zero_guard_value_fn(x))
+            elif self.log_zero_guard_type == 'clamp':
+                x = torch.log(torch.clamp(x, min=self.log_zero_guard_value_fn(x)))
+            else:
+                raise ValueError('log_zero_guard_type was not understood')
+
+        if self.frame_splicing > 1:
+            x = self.splice_frames(x, self.frame_splicing)
+        if self.normalize:
+            x, _, _ = self.normalize_batch(x, seq_len, normalize_type=self.normalize)
+
+        max_len = x.size(-1)
+        mask = torch.arange(max_len, device=x.device)
+        mask = mask.repeat(x.size(0), 1) >= seq_len.unsqueeze(1)
+        x = x.masked_fill(mask.unsqueeze(1).to(device=x.device), self.pad_value)
+        del mask
+
+        if self.pad_to == 'max':
+            x = F.pad(x, (0, self.max_length - x.size(-1)), value=self.pad_value)
+        elif self.pad_to > 0:
+            pad_amt = x.size(-1) % self.pad_to
+            if pad_amt != 0:
+                x = F.pad(x, (0, self.pad_to - pad_amt), value=self.pad_value)
+        return x, seq_len
+
+
+class CohereAsrFeatureExtractor(nn.Module):
+    def __init__(
+        self,
+        pretrained_model_or_path=None,
+        feature_size=128,
+        sampling_rate=16000,
+        padding_value=0.0,
+        max_duration=35,
+        n_window_size=400,
+        n_window_stride=160,
+        window='hann',
+        normalize='per_feature',
+        n_fft=512,
+        preemph=0.97,
+        lowfreq=0,
+        highfreq=None,
+        log=True,
+        log_zero_guard_type='add',
+        log_zero_guard_value=2**-24,
+        dither=1e-5,
+        pad_to=0,
+        frame_splicing=1,
+        exact_pad=False,
+        mag_power=2.0,
+        nb_augmentation_prob=0.0,
+        nb_max_freq=4000,
+        mel_norm='slaney',
+        device='cpu',
+    ):
+        super().__init__()
+        self.max_duration = max_duration
+        self.hop_length = n_window_stride
+        self._device = str(device)
+        self.filterbank = FilterbankFeatures(
+            pretrained_model_or_path=pretrained_model_or_path,
+            sample_rate=sampling_rate,
+            n_window_size=n_window_size,
+            n_window_stride=n_window_stride,
+            window=window,
+            normalize=normalize,
+            n_fft=n_fft,
+            preemph=preemph,
+            nfilt=feature_size,
+            lowfreq=lowfreq,
+            highfreq=highfreq,
+            log=log,
+            log_zero_guard_type=log_zero_guard_type,
+            log_zero_guard_value=log_zero_guard_value,
+            dither=dither,
+            pad_to=pad_to,
+            max_duration=max_duration,
+            frame_splicing=frame_splicing,
+            exact_pad=exact_pad,
+            pad_value=padding_value,
+            mag_power=mag_power,
+            nb_augmentation_prob=nb_augmentation_prob,
+            nb_max_freq=nb_max_freq,
+            mel_norm=mel_norm,
+            device=device,
+        )
+        self.filterbank.eval()
+        self.filterbank = self.filterbank.to(self._device)
+
+    def get_seq_len(self, seq_len: torch.Tensor) -> torch.Tensor:
+        return self.filterbank.get_seq_len(seq_len)
+
+    def __call__(self, raw_speech: list[np.ndarray], sampling_rate: int = None) -> tuple[torch.Tensor, torch.Tensor]:
+        '''Extract mel features from raw waveform input.'''
+        if sampling_rate is not None and int(sampling_rate) != int(self.filterbank.sampling_rate):
+            raise ValueError(f'Expected sampling_rate={self.filterbank.sampling_rate}, got {sampling_rate}')
+
+        if isinstance(raw_speech, np.ndarray):
+            if raw_speech.ndim == 1:
+                raw_speech = [raw_speech]
+            else:
+                raw_speech = [s for s in raw_speech]
+        elif isinstance(raw_speech, torch.Tensor):
+            if raw_speech.ndim == 1:
+                raw_speech = [raw_speech.detach().cpu().numpy()]
+            else:
+                raw_speech = [s.detach().cpu().numpy() for s in raw_speech]
+        elif not isinstance(raw_speech, (list, tuple)):
+            raise TypeError('raw_speech must be an array/tensor or list of arrays.')
+
+        normalized = []
+        for sample in raw_speech:
+            arr = np.asarray(sample, dtype=np.float32)
+            if arr.ndim != 1:
+                raise ValueError('Each audio sample must be 1D waveform.')
+            normalized.append(arr)
+
+        seq_len = torch.tensor([s.shape[0] for s in normalized], dtype=torch.long)
+        max_len = max(s.shape[0] for s in normalized)
+        padded = np.zeros((len(normalized), max_len), dtype=np.float32)
+        for i, s in enumerate(normalized):
+            padded[i, : s.shape[0]] = s
+
+        audio_tensor = torch.from_numpy(padded).to(self._device)
+        seq_len = seq_len.to(self._device)
+        with torch.no_grad():
+            input_features, length = self.filterbank(audio_tensor, seq_len)
+
+        return input_features, length
+
+
+class MaskedConvSequential(nn.Sequential):
+    def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # x: (batch, channels, time, features)
+        current_lengths = lengths.clone().float()
+        mask = self._create_mask(x, current_lengths.long())
+        for layer in self:
+            x = self.apply_channel_mask(x, mask)
+            x = layer(x)
+            if hasattr(layer, 'stride') and layer.stride != (1, 1):
+                current_lengths = self.calculate_conv_output_size(
+                    current_lengths, layer.kernel_size[0], layer.stride[0], layer.padding
+                )
+                mask = self._create_mask(x, current_lengths.long())
+        x = self.apply_channel_mask(x, mask)
+        return x, current_lengths.long()
+
+    def _create_mask(self, tensor: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
+        batch_size, _, time, features = tensor.shape
+        time_mask = torch.arange(time, device=tensor.device).expand(batch_size, time) < lengths.unsqueeze(1)
+        return time_mask.unsqueeze(-1).expand(batch_size, time, features).to(tensor.dtype)
+
+    def apply_channel_mask(self, tensor: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        batch_size, channels, time, features = tensor.shape
+        expanded_mask = mask.unsqueeze(1).expand(batch_size, channels, time, features)
+        return tensor * expanded_mask
+
+    def calculate_conv_output_size(
+        self,
+        input_size: torch.Tensor,
+        kernel_size: int,
+        stride: int,
+        padding: tuple[int, int],
+    ) -> torch.Tensor:
+        return (input_size + padding[0] + padding[1] - kernel_size) // stride + 1
+
+
+class ConvSubsampling(nn.Module):
+    def __init__(self, config: dict):
+        super().__init__()
+        feat_in = int(config['feat_in'])
+        conv_channels = int(config['subsampling_conv_channels'])
+        self._conv_channels = conv_channels
+        feat_out = int(config['feat_out'])
+        if feat_out <= 0:
+            feat_out = int(config['d_model'])
+        subsampling_factor = int(config['subsampling_factor'])
+
+        self.conv = MaskedConvSequential(
+            nn.Conv2d(1, conv_channels, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(conv_channels, conv_channels, kernel_size=3, stride=2, padding=1, groups=conv_channels),
+            nn.Conv2d(conv_channels, conv_channels, kernel_size=1),
+            nn.ReLU(),
+            nn.Conv2d(conv_channels, conv_channels, kernel_size=3, stride=2, padding=1, groups=conv_channels),
+            nn.Conv2d(conv_channels, conv_channels, kernel_size=1),
+            nn.ReLU(),
+        )
+        self.out = nn.Linear(conv_channels * (feat_in // subsampling_factor), feat_out)
+
+    def forward(self, x: torch.Tensor, lengths: torch.Tensor):
+        # x: (B, feat_in, T) -> (B, 1, T, feat_in)
+        x = x.transpose(1, 2).unsqueeze(1)
+        x, lengths = self.conv(x, lengths)
+        b, c, t, f = x.size()
+        x = x.transpose(1, 2).reshape(b, t, -1)
+        x = self.out(x)
+        return x, lengths
+
+
+class RelPositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, max_len: int = 5000):
+        super().__init__()
+        self.d_model = d_model
+        self.max_len = max_len
+        positions = torch.arange(
+            max_len - 1, -max_len, -1, dtype=torch.float32
+        ).unsqueeze(1)
+        pe = self._create_pe(positions=positions, dtype=torch.float32)
+        self.register_buffer('pe', pe, persistent=False)
+
+    def _create_pe(self, positions: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+        pos_length = positions.size(0)
+        pe = torch.zeros(pos_length, self.d_model, device=positions.device)
+        div_term = torch.exp(
+            torch.arange(0, self.d_model, 2, dtype=torch.float32, device=positions.device)
+            * -(math.log(10000.0) / self.d_model)
+        )
+        pe[:, 0::2] = torch.sin(positions * div_term)
+        pe[:, 1::2] = torch.cos(positions * div_term)
+        return pe.unsqueeze(0).to(dtype)
+
+    @torch._dynamo.disable
+    def _materialize_pe(self, length: int, device: torch.device, dtype: torch.dtype):
+        needed_size = 2 * length - 1
+        if hasattr(self, 'pe') and self.pe.device.type != 'meta' and self.pe.size(1) >= needed_size:
+            if self.pe.device != device:
+                self.pe = self.pe.to(device=device)
+            if self.pe.dtype != dtype:
+                self.pe = self.pe.to(dtype=dtype)
+            return
+        effective_length = max(length, self.max_len)
+        positions = torch.arange(
+            effective_length - 1, -effective_length, -1, dtype=torch.float32, device=device
+        ).unsqueeze(1)
+        pe = self._create_pe(positions=positions, dtype=dtype)
+        if hasattr(self, 'pe'):
+            self.pe = pe
+        else:
+            self.register_buffer('pe', pe, persistent=False)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # center_pos would be the index of position 0
+        # negative positions would be used for right and
+        # positive for left tokens
+        # for input of length L, 2*L-1 positions are needed,
+        # positions from (L-1) to -(L-1)
+        input_len = x.size(1)
+        center_pos = self.pe.size(1) // 2 + 1
+        start_pos = center_pos - input_len
+        end_pos = center_pos + input_len - 1
+        pos_emb = self.pe[:, start_pos:end_pos]
+
+        return x, pos_emb
+
+
+class ConformerFeedForward(nn.Module):
+    def __init__(self, d_model: int, d_ff: int, dropout: float):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, d_ff)
+        self.activation = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(d_ff, d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.linear1(x)
+        x = self.activation(x)
+        x = self.dropout(x)
+        x = self.linear2(x)
+        return x
+
+
+class ConformerConvolution(nn.Module):
+    def __init__(self, d_model: int, kernel_size: int):
+        super().__init__()
+        self.pointwise_conv1 = nn.Conv1d(d_model, d_model * 2, kernel_size=1)
+        self.depthwise_conv = nn.Conv1d(
+            d_model, d_model, kernel_size=kernel_size, groups=d_model, padding=(kernel_size - 1) // 2
+        )
+        self.batch_norm = nn.BatchNorm1d(d_model)
+        self.activation = nn.SiLU()
+        self.pointwise_conv2 = nn.Conv1d(d_model, d_model, kernel_size=1)
+
+    def forward(self, x: torch.Tensor, pad_mask: torch.Tensor = None) -> torch.Tensor:
+        x = x.transpose(1, 2)
+        x = self.pointwise_conv1(x)
+        x = F.glu(x, dim=1)
+        if pad_mask is not None:
+            x = x.masked_fill(pad_mask.unsqueeze(1), 0.0)
+        x = self.depthwise_conv(x)
+        x = self.batch_norm(x)
+        x = self.activation(x)
+        x = self.pointwise_conv2(x)
+        return x.transpose(1, 2)
+
+
+class RelPositionMultiHeadAttention(nn.Module):
+    def __init__(self, n_head: int, n_feat: int, dropout_rate: float):
+        super().__init__()
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.scaling = self.d_k**-0.5
+        self.pos_bias_u = nn.Parameter(torch.zeros(self.h, self.d_k))
+        self.pos_bias_v = nn.Parameter(torch.zeros(self.h, self.d_k))
+
+    def rel_shift(self, x: torch.Tensor) -> torch.Tensor:
+        '''Compute relative positional encoding.
+        Args:
+            x (torch.Tensor): (batch, nheads, time, 2*time-1)
+        '''
+        b, h, qlen, pos_len = x.size()  # (b, h, t1, t2)
+        # need to add a column of zeros on the left side of
+        # last dimension to perform the relative shifting
+        x = F.pad(x, pad=(1, 0))  # (b, h, t1, t2+1)
+        x = x.view(b, h, -1, qlen)  # (b, h, t2+1, t1)
+        # need to drop the first row
+        x = x[:, :, 1:].view(b, h, qlen, pos_len)  # (b, h, t1, t2)
+        return x
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        pos_emb: torch.Tensor,
+        mask: BlockMask | torch.Tensor = None,
+    ) -> torch.Tensor:
+        batch_size = x.size(0)
+        q = self.linear_q(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
+        k = self.linear_k(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
+        v = self.linear_v(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
+
+        # pos_emb might be shared across batch
+        if pos_emb.size(0) == 1 and batch_size > 1:
+            pos_emb = pos_emb.expand(batch_size, -1, -1)
+        p = self.linear_pos(pos_emb).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
+
+        q_with_u = q + self.pos_bias_u.unsqueeze(0).unsqueeze(2)
+        q_with_v = q + self.pos_bias_v.unsqueeze(0).unsqueeze(2)
+        matrix_bd = torch.matmul(q_with_v, p.transpose(-1, -2))
+
+        if isinstance(mask, BlockMask):
+            def score_mod(score, b, h, q_idx, kv_idx):
+                i = kv_idx - q_idx + q.size(2) - 1
+                return score + matrix_bd[b, h, q_idx, i] * self.scaling
+
+            x = torch.compile(flex_attention)(
+                q_with_u, k, v,
+                score_mod=score_mod,
+                block_mask=mask,
+                scale=self.scaling,
+            )
+
+        else:
+            matrix_bd = self.rel_shift(matrix_bd)
+            matrix_bd = matrix_bd[:, :, :, : k.size(-2)] * self.scaling
+
+            if mask is not None:
+                matrix_bd.masked_fill_(mask.unsqueeze(1), -1e9)
+
+            x = F.scaled_dot_product_attention(
+                q_with_u, k, v,
+                attn_mask=matrix_bd,
+                dropout_p=self.dropout.p if self.training else 0,
+                scale=self.scaling,
+            )
+
+        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)
+        return self.linear_out(x)
+
+
+class ConformerLayer(nn.Module):
+    def __init__(self, d_model: int, d_ff: int, n_heads: int, conv_kernel_size: int, dropout: float):
+        super().__init__()
+        self.norm_feed_forward1 = nn.LayerNorm(d_model)
+        self.feed_forward1 = ConformerFeedForward(d_model, d_ff, dropout)
+        self.norm_self_att = nn.LayerNorm(d_model)
+        self.self_attn = RelPositionMultiHeadAttention(n_heads, d_model, dropout)
+        self.norm_conv = nn.LayerNorm(d_model)
+        self.conv = ConformerConvolution(d_model, conv_kernel_size)
+        self.norm_feed_forward2 = nn.LayerNorm(d_model)
+        self.feed_forward2 = ConformerFeedForward(d_model, d_ff, dropout)
+        self.norm_out = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, pos_emb: torch.Tensor, mask: torch.Tensor = None, pad_mask: torch.Tensor = None) -> torch.Tensor:
+        residual = x
+        x = self.norm_feed_forward1(x)
+        x = residual + 0.5 * self.dropout(self.feed_forward1(x))
+
+        residual = x
+        x = self.norm_self_att(x)
+        x = residual + self.dropout(self.self_attn(x, pos_emb, mask))
+
+        residual = x
+        x = self.norm_conv(x)
+        x = residual + self.dropout(self.conv(x, pad_mask=pad_mask))
+
+        residual = x
+        x = self.norm_feed_forward2(x)
+        x = residual + 0.5 * self.dropout(self.feed_forward2(x))
+
+        return self.norm_out(x)
+
+
+class ConformerEncoder(nn.Module):
+    '''
+    Fast Conformer encoder.
+
+    Follows [Fast Conformer with Linearly Scalable Attention for Efficient Speech
+    Recognition](https://arxiv.org/abs/2305.05084).
+    '''
+
+    main_input_name = 'input_features'
+
+    def __init__(self, config: dict | str):
+        super().__init__()
+        if not isinstance(config, dict):
+            with open(config, 'r', encoding='utf-8') as file:
+                config = json.load(file)
+        enc_config = config['encoder']
+        self.d_model = enc_config['d_model']
+        d_ff = self.d_model * enc_config['ff_expansion_factor']
+        n_heads = enc_config['n_heads']
+        conv_kernel_size = enc_config['conv_kernel_size']
+        dropout = enc_config['dropout']
+        n_layers = enc_config['n_layers']
+        pos_emb_max_len = enc_config['pos_emb_max_len']
+
+        self.pre_encode = ConvSubsampling(enc_config)
+        self.pos_enc = RelPositionalEncoding(self.d_model, pos_emb_max_len)
+
+        self.layers = nn.ModuleList(
+            [ConformerLayer(self.d_model, d_ff, n_heads, conv_kernel_size, dropout) for _ in range(n_layers)]
+        )
+        self.encoder_decoder_proj = nn.Linear(self.d_model, config['transf_decoder']['config_dict']['hidden_size'])
+
+    @torch.no_grad()
+    def load(self, model: str):
+        with safe_open(model, 'pt') as file:
+            for name, param in self.named_parameters():
+                if name.startswith('encoder_decoder_proj.'):
+                    key = name
+                else:
+                    key = 'encoder.' + name
+                param.copy_(file.get_tensor(key), non_blocking=True)
+        return self
+
+    def _create_masks(
+        self,
+        padding_length: torch.Tensor,
+        max_audio_length: int,
+        device: torch.device,
+        block_size: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+    ) -> tuple[torch.Tensor, BlockMask]:
+        padding_length = padding_length.int()
+        bsz = padding_length.shape[0]
+        pad_mask = torch.arange(0, max_audio_length, device=device).expand(bsz, -1) >= padding_length[:, None]
+
+        def mask_mod(b, h, q_idx, kv_idx):
+            return (q_idx < padding_length[b]) & (kv_idx < padding_length[b])
+
+        bsz = padding_length.size(0)
+        num_blocks = -(-max_audio_length // block_size)
+        arange = torch.arange(num_blocks, dtype=torch.int32, device=device)
+
+        full_blocks = (padding_length // block_size)[:, None]
+        full = arange < full_blocks
+        full_kv_num_blocks = torch.where(full, full_blocks, 0).view(bsz, 1, num_blocks)
+        full_kv_indices = arange.expand(bsz, 1, num_blocks, num_blocks).contiguous()
+
+        partial = (padding_length % block_size != 0)[:, None]
+        partial_blocks = -(-padding_length // block_size)[:, None]
+        mask = partial & (full_blocks == arange)
+        extra_block = full & partial
+        kv_num_blocks = torch.where(mask, partial_blocks, extra_block).view(bsz, 1, num_blocks)
+        kv_indices = torch.where(
+            extra_block.view(bsz, 1, num_blocks, 1),
+            partial_blocks.view(bsz, 1, 1, 1) - 1,
+            full_kv_indices,
+        )
+
+        att_mask = BlockMask.from_kv_blocks(
+            kv_num_blocks=kv_num_blocks,
+            kv_indices=kv_indices,
+            full_kv_num_blocks=full_kv_num_blocks,
+            full_kv_indices=full_kv_indices,
+            BLOCK_SIZE=(block_size, block_size),
+            mask_mod=mask_mod,
+            seq_lengths=(max_audio_length, max_audio_length),
+            compute_q_blocks=False,
+        )
+        return pad_mask, att_mask
+
+    def forward(self, input_features: torch.Tensor = None, length: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
+        if input_features is None:
+            raise ValueError('Expected `input_features` for encoder forward.')
+        if length is None:
+            length = torch.full(
+                (input_features.shape[0],),
+                input_features.shape[-1],
+                device=input_features.device,
+                dtype=torch.long,
+            )
+        x, length = self.pre_encode(input_features, length)
+        max_audio_length = x.size(1)
+        x, pos_emb = self.pos_enc(x)
+        pad_mask, att_mask = self._create_masks(
+            padding_length=length,
+            max_audio_length=max_audio_length,
+            device=x.device,
+        )
+        for layer in self.layers:
+            x = layer(x, pos_emb, mask=att_mask, pad_mask=pad_mask)
+        x = self.encoder_decoder_proj(x)
+        return x, length
+
+
+class FixedPositionalEncoding(nn.Module):
+    def __init__(self, hidden_size: int, max_sequence_length: int = 512):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.max_sequence_length = max_sequence_length
+
+        pos_enc = torch.zeros(max_sequence_length, hidden_size)
+        position = torch.arange(0.0, max_sequence_length).unsqueeze(1)
+        coef = -math.log(10000.0) / hidden_size
+        div_term = torch.exp(coef * torch.arange(0.0, hidden_size, 2))
+        pos_enc[:, 0::2] = torch.sin(position * div_term)
+        pos_enc[:, 1::2] = torch.cos(position * div_term)
+        pos_enc.div_(math.sqrt(hidden_size))
+        self.register_buffer('pos_enc', pos_enc)
+
+    def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
+        return torch.index_select(self.pos_enc, 0, position_ids.reshape(-1)).reshape(*position_ids.shape, -1)
+
+
+class Cache(nn.Module):
+    def __init__(
+        self,
+        batch_size: int,
+        max_cache_len: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        num_key_value_heads: int,
+        head_dim: int,
+        num_hidden_layers: int,
+    ):
+        super().__init__()
+        self.batch_size = batch_size
+        self.max_cache_len = max_cache_len
+        self.cur_cache_len = max_cache_len
+        self.device = device
+        self.dtype = dtype
+        self.layers = num_hidden_layers
+
+        cache_shape = batch_size, num_key_value_heads, max_cache_len, head_dim
+        self.key_cache: list[torch.Tensor] = []
+        self.value_cache: list[torch.Tensor] = []
+        for i in range(num_hidden_layers):
+            self.register_buffer(f'key_cache_{i}', torch.zeros(cache_shape, dtype=dtype, device=device), False)
+            self.register_buffer(f'value_cache_{i}', torch.zeros(cache_shape, dtype=dtype, device=device), False)
+            k = getattr(self, f'key_cache_{i}')
+            v = getattr(self, f'value_cache_{i}')
+            torch._dynamo.mark_static_address(k)
+            torch._dynamo.mark_static_address(v)
+            self.key_cache.append(k)
+            self.value_cache.append(v)
+
+    def update(
+        self,
+        layer_idx: int,
+        bsz: int,
+        key: torch.Tensor = None,
+        value: torch.Tensor = None,
+        positions: torch.Tensor = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        k = self.key_cache[layer_idx][:bsz, :, :self.cur_cache_len, :]
+        v = self.value_cache[layer_idx][:bsz, :, :self.cur_cache_len, :]
+        if key is None:  # cross decode
+            pass
+        elif positions is None:  # prefill
+            l = key.shape[2]
+            k = k[:, :, :l, :].copy_(key)
+            v = v[:, :, :l, :].copy_(value)
+        else:  # decode
+            k.index_copy_(2, positions, key)
+            v.index_copy_(2, positions, value)
+        return k, v
+
+    def reset(self):
+        for k, v in zip(self.key_cache, self.value_cache):
+            k.zero_()
+            v.zero_()
+
+    def reorder(self, indices: torch.Tensor, finished: torch.Tensor, l: int):
+        indices = indices[~finished]
+        for i in range(self.layers):
+            self.key_cache[i][~finished, :, :l, :] = self.key_cache[i][indices, :, :l, :]
+            self.value_cache[i][~finished, :, :l, :] = self.value_cache[i][indices, :, :l, :]
+
+
+class EncoderDecoderCache(nn.Module):
+    def __init__(self, self_attention_cache: Cache, cross_attention_cache: Cache):
+        super().__init__()
+        self.self_attention_cache = self_attention_cache
+        self.cross_attention_cache = cross_attention_cache
+        self.is_updated = {}
+        for layer_idx in range(len(cross_attention_cache.key_cache)):
+            self.is_updated[layer_idx] = False
+
+    def reset(self):
+        self.self_attention_cache.reset()
+        self.cross_attention_cache.reset()
+        for layer_idx in self.is_updated:
+            self.is_updated[layer_idx] = False
+
+
+class DecoderAttention(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, layer_idx: int, has_kv: bool = True):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.layer_idx = layer_idx
+        self.head_dim = hidden_size // num_heads
+        self.scale = self.head_dim**-0.5
+        self.query_net = nn.Linear(hidden_size, hidden_size)
+        if has_kv:
+            self.key_net = nn.Linear(hidden_size, hidden_size)
+            self.value_net = nn.Linear(hidden_size, hidden_size)
+        self.out_projection = nn.Linear(hidden_size, hidden_size)
+
+    def _reshape(self, x: torch.Tensor) -> torch.Tensor:
+        b, t, _ = x.shape
+        return x.view(b, t, self.num_heads, self.head_dim).transpose(1, 2)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        context_states: torch.Tensor = None,
+        attention_mask: BlockMask | torch.Tensor = None,
+        past_key_values: EncoderDecoderCache = None,
+        positions: torch.Tensor = None,
+        diffusion: bool = False,
+    ) -> torch.Tensor:
+        bsz, tgt_len, _ = hidden_states.size()
+        query = self._reshape(self.query_net(hidden_states))
+        _flex_attention = torch.compile(flex_attention)
+
+        if context_states is not None:
+            enc_bsz = context_states.shape[0]
+            beam_size = bsz // enc_bsz
+            if past_key_values is not None and past_key_values.is_updated.get(self.layer_idx):
+                key, value = past_key_values.cross_attention_cache.update(self.layer_idx, enc_bsz)
+            else:
+                key = self._reshape(self.key_net(context_states))
+                value = self._reshape(self.value_net(context_states))
+                if past_key_values is not None:
+                    key, value = past_key_values.cross_attention_cache.update(self.layer_idx, enc_bsz, key, value)
+                    past_key_values.is_updated[self.layer_idx] = True
+
+            if beam_size != 1:  # use gqa
+                query = query.transpose(1, 2)
+                query = query.view(enc_bsz, beam_size, tgt_len, self.num_heads, self.head_dim)
+                query = query.permute(0, 3, 1, 2, 4)
+                query = query.reshape(enc_bsz, self.num_heads * beam_size, tgt_len, self.head_dim)
+
+                kernel_options = {'BLOCK_M': 1} if tgt_len == 1 else None  # incase decoding kernel not used
+                attn_output = _flex_attention(query, key, value, block_mask=attention_mask, scale=self.scale, enable_gqa=True, kernel_options=kernel_options)
+                attn_output = attn_output.view(enc_bsz, self.num_heads, beam_size, tgt_len, self.head_dim)
+                attn_output = attn_output.permute(0, 2, 3, 1, 4).reshape(bsz, tgt_len, self.hidden_size)
+            elif isinstance(attention_mask, torch.Tensor):
+                attn_output = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, scale=self.scale)
+                attn_output = attn_output.transpose(1, 2).reshape(bsz, tgt_len, self.hidden_size)
+            else:
+                attn_output = _flex_attention(query, key, value, block_mask=attention_mask, scale=self.scale)
+                attn_output = attn_output.transpose(1, 2).reshape(bsz, tgt_len, self.hidden_size)
+
+        else:
+            key = self._reshape(self.key_net(hidden_states))
+            value = self._reshape(self.value_net(hidden_states))
+            if past_key_values is not None:
+                key, value = past_key_values.self_attention_cache.update(self.layer_idx, bsz, key, value, positions)
+
+            if isinstance(attention_mask, BlockMask):
+                attn_output = _flex_attention(query, key, value, block_mask=attention_mask, scale=self.scale)
+            else:
+                attn_output = F.scaled_dot_product_attention(
+                    query=query,
+                    key=key,
+                    value=value,
+                    attn_mask=attention_mask,
+                    scale=self.scale,
+                    is_causal=attention_mask is None and not diffusion,
+                )
+            attn_output = attn_output.transpose(1, 2).reshape(bsz, tgt_len, self.hidden_size)
+
+        return self.out_projection(attn_output)
+
+
+class DecoderFeedForward(nn.Module):
+    def __init__(self, hidden_size: int, inner_size: int, hidden_act: str = 'relu'):
+        super().__init__()
+        self.dense_in = nn.Linear(hidden_size, inner_size)
+        assert hidden_act == 'relu'
+        self.activation = nn.ReLU()
+        self.dense_out = nn.Linear(inner_size, hidden_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.dense_out(self.activation(self.dense_in(x)))
+
+
+class TransformerDecoderLayer(nn.Module):
+    def __init__(self, hidden_size: int, inner_size: int, num_heads: int, layer_idx: int, hidden_act: str, diffusion: bool):
+        super().__init__()
+        self.layer_norm_1 = nn.LayerNorm(hidden_size)
+        self.first_sub_layer = DecoderAttention(hidden_size, num_heads, layer_idx)
+        self.layer_norm_2 = nn.LayerNorm(hidden_size)
+        self.second_sub_layer = DecoderAttention(hidden_size, num_heads, layer_idx, has_kv=not diffusion)
+        self.layer_norm_3 = nn.LayerNorm(hidden_size)
+        self.third_sub_layer = DecoderFeedForward(hidden_size, inner_size, hidden_act=hidden_act)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        self_attention_mask: BlockMask | torch.Tensor = None,
+        cross_attention_mask: BlockMask | torch.Tensor = None,
+        past_key_values: EncoderDecoderCache = None,
+        positions: torch.Tensor = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.layer_norm_1(hidden_states)
+        self_out = self.first_sub_layer(
+            hidden_states,
+            context_states=None,
+            attention_mask=self_attention_mask,
+            past_key_values=past_key_values,
+            positions=positions,
+        )
+        hidden_states = residual + self_out
+
+        residual = hidden_states
+        hidden_states = self.layer_norm_2(hidden_states)
+        cross_out = self.second_sub_layer(
+            hidden_states,
+            context_states=encoder_hidden_states,
+            attention_mask=cross_attention_mask,
+            past_key_values=past_key_values,
+        )
+        hidden_states = residual + cross_out
+
+        residual = hidden_states
+        hidden_states = self.layer_norm_3(hidden_states)
+        hidden_states = residual + self.third_sub_layer(hidden_states)
+        return hidden_states
+
+
+class TransformerDecoderEmbedding(nn.Module):
+    def __init__(self, vocab_size: int, hidden_size: int, max_sequence_length: int, padding_idx: int = 2):
+        super().__init__()
+        self.token_embedding = nn.Embedding(vocab_size, hidden_size, padding_idx)
+        self.position_embedding = FixedPositionalEncoding(hidden_size, max_sequence_length)
+        self.layer_norm = nn.LayerNorm(hidden_size)
+
+    def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        token_embeds = self.token_embedding(input_ids)
+        if positions is None:
+            pos_embeds = self.position_embedding.pos_enc[:input_ids.shape[-1]]
+        else:
+            pos_embeds = self.position_embedding(positions)
+        return self.layer_norm(token_embeds + pos_embeds)
+
+
+class TransformerDecoderCore(nn.Module):
+    def __init__(self, hidden_size: int, inner_size: int, num_heads: int, num_layers: int, hidden_act: str, diffusion: bool):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                TransformerDecoderLayer(hidden_size, inner_size, num_heads, i, hidden_act, diffusion)
+                for i in range(num_layers)
+            ]
+        )
+        self.final_layer_norm = nn.LayerNorm(hidden_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        self_attention_mask: BlockMask | torch.Tensor = None,
+        cross_attention_mask: BlockMask | torch.Tensor = None,
+        past_key_values: EncoderDecoderCache = None,
+        positions: torch.Tensor = None,
+    ) -> torch.Tensor:
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                self_attention_mask=self_attention_mask,
+                cross_attention_mask=cross_attention_mask,
+                past_key_values=past_key_values,
+                positions=positions,
+            )
+        return self.final_layer_norm(hidden_states)
+
+
+class TransformerDecoderWrapper(nn.Module):
+    def __init__(self, config: str | dict):
+        super().__init__()
+        if not isinstance(config, dict):
+            with open(config, 'r', encoding='utf-8') as file:
+                config = json.load(file)
+        dec_config = config['transf_decoder']['config_dict']
+        hidden_size = dec_config['hidden_size']
+        self._embedding = TransformerDecoderEmbedding(
+            vocab_size=config['vocab_size'],
+            hidden_size=hidden_size,
+            max_sequence_length=dec_config['max_sequence_length'],
+            padding_idx=2,
+        )
+        self._decoder = TransformerDecoderCore(
+            hidden_size=hidden_size,
+            inner_size=dec_config['inner_size'],
+            num_heads=dec_config['num_attention_heads'],
+            num_layers=dec_config['num_layers'],
+            hidden_act=dec_config.get('hidden_act', 'relu'),
+            diffusion=False,
+        )
+        self._lm_head = nn.Linear(hidden_size, config['vocab_size'])
+        self._lm_head.weight = self._embedding.token_embedding.weight
+
+    @torch.no_grad()
+    def load(self, model: str):
+        with safe_open(model, 'pt') as file:
+            for name, param in self.named_parameters():
+                if name.startswith('_lm_head.'):
+                    key = 'log_softmax.mlp.layer0.' + name.removeprefix('_lm_head.')
+                else:
+                    key = 'transf_decoder.' + name
+                param.copy_(file.get_tensor(key), non_blocking=True)
+        return self
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor = None,
+        encoder_hidden_states: BlockMask | torch.Tensor = None,
+        self_attention_mask: BlockMask | torch.Tensor = None,
+        cross_attention_mask: BlockMask | torch.Tensor = None,
+        past_key_values: EncoderDecoderCache = None,
+    ) -> torch.Tensor:
+        hidden_states = self._embedding(input_ids, positions)
+        hidden_states = self._decoder(
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            self_attention_mask=self_attention_mask,
+            cross_attention_mask=cross_attention_mask,
+            past_key_values=past_key_values,
+            positions=positions,
+        )
+        return self._lm_head(hidden_states)
+
+
+class CohereAsr(nn.Module):
+    def __init__(self, config: dict | str, pad_id: int = 2, eot_id: int = 3):
+        super().__init__()
+        if not isinstance(config, dict):
+            with open(config, 'r', encoding='utf-8') as file:
+                config = json.load(file)
+        self.encoder = ConformerEncoder(config)
+        self.transf_decoder = TransformerDecoderWrapper(config)
+        dec_config = config['transf_decoder']['config_dict']
+        self.decoder_hidden_size = dec_config['hidden_size']
+        self.diffusion = dec_config.get('diffusion', 0)
+        if self.diffusion > 0:
+            self.diff_decoder = TransformerDecoderCore(
+                hidden_size=self.decoder_hidden_size,
+                inner_size=dec_config['inner_size'],
+                num_heads=dec_config['num_attention_heads'],
+                num_layers=dec_config['num_layers'],
+                hidden_act=dec_config.get('hidden_act', 'relu'),
+                diffusion=True,
+            )
+
+        self.pad_id = pad_id
+        self.eot_id = eot_id
+
+    @torch.no_grad()
+    def load(self, model: str):
+        with safe_open(model, 'pt') as file:
+            for name, param in self.named_parameters():
+                if name.startswith('encoder.encoder_decoder_proj.'):
+                    key = name.removeprefix('encoder.')
+                elif name.startswith('transf_decoder._lm_head.'):
+                    key = 'log_softmax.mlp.layer0.' + name.removeprefix('transf_decoder._lm_head.')
+                else:
+                    key = name
+                param.copy_(file.get_tensor(key), non_blocking=True)
+        return self
+
+    @classmethod
+    def from_pretrained(cls, path: Path, device: torch.device = 'cpu', dtype: torch.dtype = None):
+        path = Path(path)
+        with torch.device('meta'):
+            model = cls(path/'config.json')
+        state = {}
+        with safe_open(path/'model.safetensors', 'pt', device) as file:
+            for key in file.keys():
+                if key.startswith('preprocessor'):
+                    continue
+                elif key.startswith('encoder_decoder_proj.'):
+                    name = 'encoder.' + key
+                elif key.startswith('log_softmax.mlp.layer0.'):
+                    name = 'transf_decoder._lm_head.' + key.removeprefix('log_softmax.mlp.layer0.')
+                else:
+                    name = key
+                state[name] = file.get_tensor(key).to(dtype)
+        model.load_state_dict(state, assign=True)
+        model.encoder.pos_enc._materialize_pe(model.encoder.pos_enc.max_len, device, state['encoder.pre_encode.out.weight'].dtype)
+        model.transf_decoder._lm_head.weight = model.transf_decoder._embedding.token_embedding.weight
+        return model.eval()
+
+    def create_block_masks(
+        self,
+        encoder_lengths: torch.Tensor,
+        bsz: int,
+        q_seq_len: int,
+        kv_seq_len: int,
+        device: torch.device,
+        block_size: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+    ) -> tuple[BlockMask, BlockMask]:
+        num_blocks = -(-q_seq_len // block_size)
+        assert num_blocks == 1
+
+        def causal_mask(b, h, q_idx, kv_idx):
+            return kv_idx <= q_idx
+        kv_num_blocks = torch.ones((1, 1, 1), dtype=torch.int32, device=device)
+        kv_indices = torch.zeros((1, 1, 1, 1), dtype=torch.int32, device=device)
+        self_attn_mask = BlockMask.from_kv_blocks(
+            kv_num_blocks=kv_num_blocks,
+            kv_indices=kv_indices,
+            full_kv_num_blocks=None,
+            full_kv_indices=None,
+            BLOCK_SIZE=(block_size, block_size),
+            mask_mod=causal_mask,
+            seq_lengths=(q_seq_len, q_seq_len),
+            compute_q_blocks=False,
+        )
+        encoder_lengths = encoder_lengths.int()
+
+        def cross_mask(b, h, q_idx, kv_idx):
+            return kv_idx < encoder_lengths[b]
+        num_kv_blocks = -(-kv_seq_len // block_size)
+        kv_num_blocks = (encoder_lengths % block_size != 0).int().view(bsz, 1, 1)
+        full_kv_num_blocks = (encoder_lengths // block_size).view(bsz, 1, 1)
+        kv_indices = full_kv_num_blocks.view(bsz, 1, 1, 1).expand(bsz, 1, 1, num_kv_blocks).contiguous()
+        full_kv_indices = torch.arange(num_kv_blocks, dtype=torch.int32, device=device).expand(bsz, 1, 1, num_kv_blocks).contiguous()
+        cross_attn_mask = BlockMask.from_kv_blocks(
+            kv_num_blocks=kv_num_blocks,
+            kv_indices=kv_indices,
+            full_kv_num_blocks=full_kv_num_blocks,
+            full_kv_indices=full_kv_indices,
+            BLOCK_SIZE=(block_size, block_size),
+            mask_mod=cross_mask,
+            seq_lengths=(q_seq_len, kv_seq_len),
+            compute_q_blocks=False,
+        )
+        return self_attn_mask, cross_attn_mask
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor = None,
+        input_features: torch.Tensor = None,
+        encoder_hidden_states: torch.Tensor = None,
+        lengths: torch.Tensor = None,
+        self_attention_mask: BlockMask | torch.Tensor = None,
+        cross_attention_mask: BlockMask | torch.Tensor = None,
+        past_key_values: EncoderDecoderCache = None,
+        diffusion: bool = False,
+    ) -> torch.Tensor:
+        if encoder_hidden_states is None:
+            encoder_hidden_states, _ = self.encoder(input_features, lengths)
+
+        logits = self.transf_decoder(
+            input_ids=input_ids,
+            positions=positions,
+            encoder_hidden_states=encoder_hidden_states,
+            self_attention_mask=self_attention_mask,
+            cross_attention_mask=cross_attention_mask,
+            past_key_values=past_key_values,
+            diffusion=diffusion,
+        )
+        return logits
+
+    def prepare_cache(
+        self,
+        bsz: int,
+        max_len: int,
+        kv_len: int,
+        device: torch.device,
+        beam_size: int = 1,
+    ) -> EncoderDecoderCache:
+        past_key_values = None
+        if hasattr(self, '_cache'):
+            past_key_values = self._cache
+            if (past_key_values.self_attention_cache.batch_size >= bsz
+                and past_key_values.self_attention_cache.max_cache_len >= max_len
+                and past_key_values.cross_attention_cache.batch_size >= bsz
+                    and past_key_values.cross_attention_cache.max_cache_len >= kv_len):
+                past_key_values.reset()
+                past_key_values.self_attention_cache.cur_cache_len = max_len
+                past_key_values.cross_attention_cache.cur_cache_len = kv_len
+            else:
+                past_key_values = None
+
+        if past_key_values is None:
+            attn: DecoderAttention = self.transf_decoder._decoder.layers[0].first_sub_layer
+            cache_kwargs = {
+                'batch_size': bsz * beam_size,
+                'max_cache_len': max_len,
+                'device': device,
+                'dtype': attn.query_net.weight.dtype,
+                'num_key_value_heads': attn.num_heads,
+                'head_dim': attn.head_dim,
+                'num_hidden_layers': len(self.transf_decoder._decoder.layers),
+            }
+            cross_kwargs = cache_kwargs | {'batch_size': bsz, 'max_cache_len': kv_len}
+            past_key_values = EncoderDecoderCache(Cache(**cache_kwargs), Cache(**cross_kwargs))
+            self._cache = past_key_values
+        return past_key_values
+
+    def _sample(
+        self,
+        input_ids: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_lengths: torch.Tensor,
+        sample: bool = False,
+        topk: int = 0,
+        temp: float = 1.0,
+        max_len: int = 256,
+        flex_attn: bool = True,
+        block_size: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+    ) -> torch.Tensor:
+        device = input_ids.device
+        bsz, prompt_len = input_ids.shape
+
+        kv_len = encoder_hidden_states.shape[1]
+        input_ids = F.pad(input_ids, (0, max_len - prompt_len), 'constant', self.pad_id)
+        kv_seq_len = torch.zeros((), dtype=torch.int64, device=device)
+        finished = input_ids[:, prompt_len-1] == self.pad_id
+        done = torch.zeros((), dtype=torch.bool, pin_memory=True)
+        event = torch.Event(device)
+        past_key_values = self.prepare_cache(bsz, max_len, kv_len, device)
+
+        if flex_attn:
+            self_attention_mask, cross_attention_mask = self.create_block_masks(
+                encoder_lengths=encoder_lengths,
+                bsz=bsz,
+                q_seq_len=prompt_len,
+                kv_seq_len=kv_len,
+                device=device,
+                block_size=block_size,
+            )
+            assert prompt_len <= block_size
+        else:
+            self_attention_mask = None
+            mask = torch.arange(kv_len, device=device) < encoder_lengths[:, None]
+            cross_attention_mask = mask[:, None, None, :]
+
+        for i in range(prompt_len, max_len):
+            next_ids = input_ids[:, :prompt_len] if i == prompt_len else input_ids[:, i-1:i]
+            logits: torch.Tensor = self.transf_decoder(
+                input_ids=next_ids,
+                positions=kv_seq_len if next_ids.shape[1] == 1 else None,
+                encoder_hidden_states=encoder_hidden_states,
+                self_attention_mask=self_attention_mask,
+                cross_attention_mask=cross_attention_mask,
+                past_key_values=past_key_values,
+            )
+            if sample:
+                if temp != 1.0:
+                    logits.div_(temp)
+                if topk != 0:
+                    indices_to_remove = logits < torch.topk(logits, topk)[0][..., -1, None]
+                    logits.masked_fill_(indices_to_remove, float('-inf'))
+                probs = F.softmax(logits[:, -1, :], dim=1, dtype=torch.float32)
+                new_tokens = torch.multinomial(probs, num_samples=1)[:, 0]
+            else:
+                new_tokens = torch.argmax(logits[:, -1, :], dim=1)
+            new_tokens[finished] = self.pad_id
+
+            if i == prompt_len and flex_attn:
+                def decode_mask(b, h, q_idx, kv_idx):
+                    return kv_idx <= kv_seq_len
+                num_blocks = -(-max_len // block_size)
+                kv_num_blocks = torch.ones((bsz, 1, 1), dtype=torch.int32, device=device)
+                kv_indices = torch.zeros((bsz, 1, 1, num_blocks), dtype=torch.int32, device=device)
+                full_kv_num_blocks = torch.zeros((bsz, 1, 1), dtype=torch.int32, device=device)
+                full_kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device).expand(bsz, 1, 1, num_blocks).contiguous()
+                self_attention_mask = BlockMask.from_kv_blocks(
+                    kv_num_blocks=kv_num_blocks,
+                    kv_indices=kv_indices,
+                    full_kv_num_blocks=full_kv_num_blocks,
+                    full_kv_indices=full_kv_indices,
+                    BLOCK_SIZE=(1, block_size),
+                    mask_mod=decode_mask,
+                    seq_lengths=(1, max_len),
+                    compute_q_blocks=False,
+                )
+                cross_attention_mask.seq_lengths = 1, kv_len
+                cross_attention_mask.BLOCK_SIZE = 1, cross_attention_mask.BLOCK_SIZE[1]
+            elif i == prompt_len:
+                mask = torch.arange(max_len, device=device) <= prompt_len
+                self_attention_mask = mask[None, None, None, :]
+            else:
+                event.synchronize()
+                if done:
+                    break
+
+            finished.logical_or_(new_tokens == self.eot_id)
+            done.copy_(finished.all(), non_blocking=True)
+            event.record()
+            kv_seq_len += next_ids.shape[1]
+            input_ids.index_copy_(1, kv_seq_len, new_tokens[:, None])
+
+            if flex_attn:
+                if (i - 1) % block_size == 0:
+                    self_attention_mask.kv_indices += 1
+                    self_attention_mask.full_kv_num_blocks += 1
+                self_attention_mask.kv_num_blocks[finished] = 0
+                self_attention_mask.full_kv_num_blocks[finished] = 0
+                cross_attention_mask.kv_num_blocks[finished] = 0
+                cross_attention_mask.full_kv_num_blocks[finished] = 0
+            else:
+                self_attention_mask[:, :, :, kv_seq_len] = True
+
+        return input_ids
+
+    def _beam_search(
+        self,
+        input_ids: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_lengths: torch.Tensor,
+        beam_size: int = 5,
+        max_len: int = 256,
+        flex_attn: bool = True,
+        block_size: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+    ) -> torch.Tensor:
+        device = input_ids.device
+        bsz, prompt_len = input_ids.shape
+
+        kv_len = encoder_hidden_states.shape[1]
+        flat_size = bsz * beam_size
+        kv_seq_len = torch.zeros((), dtype=torch.int64, device=device)
+        done = torch.zeros((), dtype=torch.bool, pin_memory=True)
+        event = torch.Event(device)
+        past_key_values = self.prepare_cache(bsz, max_len, kv_len, device, beam_size)
+        finish_ids = torch.tensor([self.pad_id, self.eot_id], device=device)
+
+        if flex_attn:
+            self_attention_mask, cross_attention_mask = self.create_block_masks(
+                encoder_lengths=encoder_lengths,
+                bsz=bsz,
+                q_seq_len=prompt_len,
+                kv_seq_len=kv_len,
+                device=device,
+                block_size=block_size,
+            )
+            assert prompt_len <= block_size
+        else:
+            self_attention_mask = None
+            mask = torch.arange(kv_len, device=device) < encoder_lengths[:, None]
+            cross_attention_mask = mask[:, None, None, :]
+
+        for i in range(prompt_len, max_len):
+            next_ids = input_ids if i == prompt_len else input_ids[:, -1:]
+            logits: torch.Tensor = self.transf_decoder(
+                input_ids=next_ids.contiguous(),
+                positions=kv_seq_len if next_ids.shape[1] == 1 else None,
+                encoder_hidden_states=encoder_hidden_states,
+                self_attention_mask=self_attention_mask,
+                cross_attention_mask=cross_attention_mask,
+                past_key_values=past_key_values,
+            )
+            logprobs = F.log_softmax(logits[:, -1, :], dim=1, dtype=torch.float32)
+
+            if i == prompt_len:
+                if flex_attn:
+                    def decode_mask(b, h, q_idx, kv_idx):
+                        return kv_idx <= kv_seq_len
+                    num_blocks = -(-max_len // block_size)
+                    kv_num_blocks = torch.ones((flat_size, 1, 1), dtype=torch.int32, device=device)
+                    kv_indices = torch.zeros((flat_size, 1, 1, num_blocks), dtype=torch.int32, device=device)
+                    full_kv_num_blocks = torch.zeros((flat_size, 1, 1), dtype=torch.int32, device=device)
+                    full_kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device).expand(flat_size, 1, 1, num_blocks).contiguous()
+                    self_attention_mask = BlockMask.from_kv_blocks(
+                        kv_num_blocks=kv_num_blocks,
+                        kv_indices=kv_indices,
+                        full_kv_num_blocks=full_kv_num_blocks,
+                        full_kv_indices=full_kv_indices,
+                        BLOCK_SIZE=(1, block_size),
+                        mask_mod=decode_mask,
+                        seq_lengths=(1, max_len),
+                        compute_q_blocks=False,
+                    )
+                    cross_attention_mask.seq_lengths = 1, kv_len
+                    cross_attention_mask.BLOCK_SIZE = 1, cross_attention_mask.BLOCK_SIZE[1]
+                else:
+                    mask = torch.arange(max_len, device=device) <= prompt_len
+                    self_attention_mask = mask[None, None, None, :]
+
+                scores, new_tokens = torch.topk(logprobs, k=beam_size, dim=1)
+                new_indices = torch.arange(bsz, device=device).repeat_interleave(beam_size, dim=0)
+
+            else:
+                event.synchronize()
+                if done:
+                    break
+                vocab_size = logprobs.shape[1]
+                logprobs[finished] = 0
+                batch_logprobs = (logprobs + scores.view(-1, 1)).view(bsz, -1)
+                scores, topk_indices = torch.topk(batch_logprobs, k=beam_size, dim=1)
+                new_indices = topk_indices // vocab_size
+                new_indices += torch.arange(0, flat_size, beam_size, device=device)[:, None]
+                new_indices = new_indices.flatten()
+                new_tokens = topk_indices % vocab_size
+
+            input_ids = input_ids[new_indices]
+            finished = torch.isin(input_ids[:, -1].view(bsz, beam_size), finish_ids)
+            finished[finished[:, 0]] = True
+            new_tokens[finished] = self.pad_id
+            input_ids = torch.cat((input_ids, new_tokens.view(-1, 1)), dim=1)
+            finished = torch.isin(input_ids[:, -1], finish_ids)
+            done.copy_(finished.all(), non_blocking=True)
+            event.record()
+            past_key_values.self_attention_cache.reorder(new_indices, finished, i)
+            kv_seq_len += next_ids.shape[1]
+
+            if flex_attn:
+                if (i - 1) % block_size == 0:
+                    self_attention_mask.kv_indices += 1
+                    self_attention_mask.full_kv_num_blocks += 1
+                self_attention_mask.kv_num_blocks[finished] = 0
+                self_attention_mask.full_kv_num_blocks[finished] = 0
+                cross_finished = finished.view(bsz, beam_size).all(dim=1)
+                cross_attention_mask.kv_num_blocks[cross_finished] = 0
+                cross_attention_mask.full_kv_num_blocks[cross_finished] = 0
+            else:
+                self_attention_mask[:, :, :, kv_seq_len] = True
+
+        return input_ids.view(bsz, beam_size, -1)[:, 0, :]
+
+    def _spec_forward(
+        self,
+        input_ids: torch.LongTensor,  # [b,1]
+        encoder_hidden_states: torch.Tensor,  # [b,t,d]
+        cross_attention_mask: BlockMask,
+        diff_attention_mask: BlockMask,
+        self_attention_mask: BlockMask,
+        past_key_values: EncoderDecoderCache,
+        finished: torch.BoolTensor,  # [b]
+        kv_seq_len: torch.IntTensor,  # []
+        diff_len: int,
+    ) -> tuple[torch.LongTensor, torch.IntTensor]:
+        device = input_ids.device
+        kv_len = encoder_hidden_states.shape[1]
+        cross_attention_mask.seq_lengths = diff_len, kv_len
+        cross_attention_mask.BLOCK_SIZE = diff_len, cross_attention_mask.BLOCK_SIZE[1]
+
+        positions = torch.arange(diff_len, device=device) + kv_seq_len
+        diff_states = self.transf_decoder._embedding(
+            input_ids=F.pad(input_ids, (0, diff_len - 1), 'constant', self.pad_id),
+            positions=positions,
+        )
+        diff_states = self.diff_decoder(
+            hidden_states=diff_states,
+            encoder_hidden_states=encoder_hidden_states,
+            self_attention_mask=diff_attention_mask,
+            cross_attention_mask=cross_attention_mask,
+            past_key_values=past_key_values,
+            positions=positions,
+        )
+        diff_logits = self.transf_decoder._lm_head(diff_states)
+        diff_tokens = torch.argmax(diff_logits, dim=2)
+        diff_eot = (diff_tokens == self.eot_id).int()
+        diff_pad = diff_eot.cumsum(dim=1) - diff_eot > 0
+        diff_tokens[finished[:, None] | diff_pad] = self.pad_id
+
+        cross_attention_mask.seq_lengths = diff_len + 1, kv_len
+        cross_attention_mask.BLOCK_SIZE = diff_len + 1, cross_attention_mask.BLOCK_SIZE[1]
+        ar_logits = self.transf_decoder(
+            input_ids=torch.cat((input_ids, diff_tokens), dim=1),
+            positions=torch.arange(diff_len + 1, device=device) + kv_seq_len,
+            encoder_hidden_states=encoder_hidden_states,
+            self_attention_mask=self_attention_mask,
+            cross_attention_mask=cross_attention_mask,
+            past_key_values=past_key_values,
+        )
+        ar_tokens = torch.argmax(ar_logits, dim=2)
+        ar_eot = (ar_tokens == self.eot_id).int()
+        ar_pad = ar_eot.cumsum(dim=1) - ar_eot > 0
+        ar_tokens[finished[:, None] | ar_pad] = self.pad_id
+
+        matches = (diff_tokens == ar_tokens[:, :-1]).int().cumprod(dim=1)
+        match_len = matches.sum(dim=1).min() + 1
+        return ar_tokens, match_len
+
+    def _speculate(
+        self,
+        input_ids: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_lengths: torch.Tensor,
+        max_len: int = 256,
+        block_size: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+        compile: bool = False,
+        compile_options: dict = {},
+    ) -> torch.Tensor:
+        device = input_ids.device
+        bsz, prompt_len = input_ids.shape
+        _spec_forward = (torch.compile(self._spec_forward, **compile_options)
+                         if compile else self._spec_forward)
+
+        kv_len = encoder_hidden_states.shape[1]
+        input_ids = F.pad(input_ids, (0, max_len - prompt_len), 'constant', self.pad_id)
+        finished = input_ids[:, prompt_len-1] == self.pad_id
+        diff_len = self.diffusion
+        cache_len = max_len + diff_len
+        past_key_values = self.prepare_cache(bsz, cache_len, kv_len, device)
+
+        self_attention_mask, cross_attention_mask = self.create_block_masks(
+            encoder_lengths=encoder_lengths,
+            bsz=bsz,
+            q_seq_len=prompt_len,
+            kv_seq_len=kv_len,
+            device=device,
+            block_size=block_size,
+        )
+        assert prompt_len + diff_len <= block_size
+
+        # prefill
+        logits: torch.Tensor = self.transf_decoder(
+            input_ids=input_ids[:, :prompt_len],
+            positions=None,
+            encoder_hidden_states=encoder_hidden_states,
+            self_attention_mask=self_attention_mask,
+            cross_attention_mask=cross_attention_mask,
+            past_key_values=past_key_values,
+        )
+        new_tokens = torch.argmax(logits[:, -1, :], dim=1)
+        finished.logical_or_(new_tokens == self.eot_id)
+        kv_seq_len = torch.full((), prompt_len, dtype=torch.int64, device=device)
+        input_ids[:, prompt_len] = new_tokens
+
+        # decode
+        num_blocks = -(-cache_len // block_size)
+
+        def diff_mask(b, h, q_idx, kv_idx):
+            return kv_idx <= kv_seq_len + diff_len
+        kv_num_blocks = torch.ones((bsz, 1, 1), dtype=torch.int32, device=device)
+        kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device).expand(bsz, 1, 1, num_blocks).contiguous()
+        full_kv_num_blocks = torch.zeros((bsz, 1, 1), dtype=torch.int32, device=device)
+        full_kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device).expand(bsz, 1, 1, num_blocks).contiguous()
+        diff_attention_mask = BlockMask.from_kv_blocks(
+            kv_num_blocks=kv_num_blocks,
+            kv_indices=kv_indices,
+            full_kv_num_blocks=full_kv_num_blocks,
+            full_kv_indices=full_kv_indices,
+            BLOCK_SIZE=(diff_len, block_size),
+            mask_mod=diff_mask,
+            seq_lengths=(diff_len, cache_len),
+            compute_q_blocks=False,
+        )
+
+        def decode_mask(b, h, q_idx, kv_idx):
+            return kv_idx <= kv_seq_len + q_idx
+        kv_num_blocks = torch.ones((bsz, 1, 1), dtype=torch.int32, device=device)
+        kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device).expand(bsz, 1, 1, num_blocks).contiguous()
+        self_attention_mask = BlockMask.from_kv_blocks(
+            kv_num_blocks=kv_num_blocks,
+            kv_indices=kv_indices,
+            full_kv_num_blocks=None,
+            full_kv_indices=None,
+            BLOCK_SIZE=(diff_len + 1, block_size),
+            mask_mod=decode_mask,
+            seq_lengths=(diff_len + 1, cache_len),
+            compute_q_blocks=False,
+        )
+
+        i = prompt_len + 1
+        while i < max_len:
+            ar_tokens, match_len = _spec_forward(
+                input_ids=input_ids[:, i-1:i],
+                encoder_hidden_states=encoder_hidden_states,
+                cross_attention_mask=cross_attention_mask,
+                diff_attention_mask=diff_attention_mask,
+                self_attention_mask=self_attention_mask,
+                past_key_values=past_key_values,
+                finished=finished,
+                kv_seq_len=kv_seq_len,
+                diff_len=diff_len,
+            )
+            new_len = min(match_len.item(), max_len - i)
+            new_tokens = ar_tokens[:, :new_len]
+            input_ids[:, i:i+new_len] = new_tokens
+            finished.logical_or_(torch.any(new_tokens == self.eot_id, dim=1))
+            if finished.all():
+                break
+
+            diff_i = i + diff_len
+            if diff_i // block_size < (diff_i + new_len) // block_size:
+                diff_attention_mask.kv_indices += 1
+                diff_attention_mask.full_kv_num_blocks += 1
+                self_attention_mask.kv_num_blocks += 1
+            diff_attention_mask.kv_num_blocks[finished] = 0
+            diff_attention_mask.full_kv_num_blocks[finished] = 0
+            self_attention_mask.kv_num_blocks[finished] = 0
+            cross_attention_mask.kv_num_blocks[finished] = 0
+            cross_attention_mask.full_kv_num_blocks[finished] = 0
+            kv_seq_len += match_len
+            i += new_len
+
+        return input_ids
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        input_features: torch.Tensor = None,
+        lengths: torch.Tensor = None,
+        encoder_hidden_states: torch.Tensor = None,
+        sample: bool = False,
+        topk: int = 0,
+        temp: float = 1.0,
+        beam_size: int = 1,
+        max_len: int = 256,
+        flex_attn: bool = True,
+        block_size: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+        compile: bool = False,
+        compile_options: dict = {'mode': 'reduce-overhead', 'fullgraph': True, 'dynamic': False},
+        diffusion: bool = False,
+    ) -> torch.Tensor:
+        if compile:
+            self.encoder.compile(**compile_options)
+            self.transf_decoder.compile(**compile_options)
+        if encoder_hidden_states is None:
+            encoder_hidden_states, encoder_lengths = self.encoder(input_features, lengths)
+            encoder_hidden_states = encoder_hidden_states.clone()
+        else:
+            encoder_lengths = lengths
+
+        if diffusion:
+            if sample or beam_size != 1 or not flex_attn or not self.diffusion:
+                raise NotImplementedError()
+            return self._speculate(
+                input_ids=input_ids,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_lengths=encoder_lengths,
+                max_len=max_len,
+                block_size=block_size,
+                compile=compile,
+                compile_options=compile_options,
+            )
+        if beam_size != 1:
+            if sample:
+                raise NotImplementedError()
+            return self._beam_search(
+                input_ids=input_ids,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_lengths=encoder_lengths,
+                beam_size=beam_size,
+                max_len=max_len,
+                flex_attn=flex_attn,
+                block_size=block_size,
+            )
+        return self._sample(
+            input_ids=input_ids,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_lengths=encoder_lengths,
+            sample=sample,
+            topk=topk,
+            temp=temp,
+            max_len=max_len,
+            flex_attn=flex_attn,
+            block_size=block_size,
+        )

config.json

@@ -0,0 +1,173 @@
+{
+  "architectures": [
+    "CohereAsrForConditionalGeneration"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_cohere_asr.CohereAsrConfig",
+    "AutoFeatureExtractor": "processing_cohere_asr.CohereAsrFeatureExtractor",
+    "AutoModel": "modeling_cohere_asr.CohereAsrModel",
+    "AutoModelForSpeechSeq2Seq": "modeling_cohere_asr.CohereAsrForConditionalGeneration",
+    "AutoProcessor": "processing_cohere_asr.CohereAsrProcessor",
+    "AutoTokenizer": "tokenization_cohere_asr.CohereAsrTokenizer"
+  },
+  "batch_size": 128,
+  "decoding": {
+    "beam": {
+      "beam_size": 1,
+      "len_pen": 0.0,
+      "max_generation_delta": 50
+    },
+    "return_best_hypothesis": true,
+    "strategy": "beam"
+  },
+  "encoder": {
+    "att_context_size": [
+      -1,
+      -1
+    ],
+    "causal_downsampling": false,
+    "conv_context_size": null,
+    "conv_kernel_size": 9,
+    "conv_norm_type": "batch_norm",
+    "d_model": 1280,
+    "dropout": 0,
+    "dropout_att": 0,
+    "dropout_emb": 0,
+    "dropout_pre_encoder": 0,
+    "feat_in": 128,
+    "feat_out": -1,
+    "ff_expansion_factor": 4,
+    "n_heads": 8,
+    "n_layers": 48,
+    "pos_emb_max_len": 5000,
+    "reduction": null,
+    "reduction_factor": 1,
+    "reduction_position": null,
+    "self_attention_model": "rel_pos",
+    "subsampling": "dw_striding",
+    "subsampling_conv_channels": 256,
+    "subsampling_factor": 8,
+    "untie_biases": true,
+    "xscaling": false
+  },
+  "head": {
+    "activation": "relu",
+    "dropout": 0,
+    "hidden_size": 1024,
+    "log_softmax": true,
+    "num_classes": 16384,
+    "num_layers": 1,
+    "use_transformer_init": true
+  },
+  "is_encoder_decoder": true,
+  "log_batch_stats": false,
+  "log_prediction": true,
+  "max_audio_clip_s": 35,
+  "max_seq_len": 1024,
+  "model_defaults": {
+    "asr_enc_hidden": 1280,
+    "lm_dec_hidden": 1024,
+    "lm_enc_hidden": 1024
+  },
+  "model_type": "cohere_asr",
+  "multitask_metrics_cfg": {
+    "log_predictions": true,
+    "metrics": {
+      "wer": {
+        "constraint": ".source_lang==.target_lang"
+      }
+    }
+  },
+  "overlap_chunk_second": 5,
+  "preprocessor": {
+    "dither": 1e-05,
+    "features": 128,
+    "frame_splicing": 1,
+    "log": true,
+    "n_fft": 512,
+    "normalize": "per_feature",
+    "pad_to": 0,
+    "pad_value": 0.0,
+    "sample_rate": 16000,
+    "window": "hann",
+    "window_size": 0.025,
+    "window_stride": 0.01
+  },
+  "prompt_defaults": [
+    {
+      "role": "user",
+      "slots": {
+        "decodercontext": "",
+        "diarize": "<|nodiarize|>",
+        "emotion": "<|emo:undefined|>",
+        "itn": "<|noitn|>",
+        "pnc": "<|pnc|>",
+        "source_lang": "<|en|>",
+        "target_lang": "<|en|>",
+        "timestamp": "<|notimestamp|>"
+      }
+    },
+    {
+      "role": "user_partial",
+      "slots": {
+        "decodercontext": ""
+      }
+    }
+  ],
+  "prompt_format": "cohere_asr",
+  "sample_rate": 16000,
+  "supported_languages": [
+    "en",
+    "fr",
+    "de",
+    "es",
+    "it",
+    "pt",
+    "nl",
+    "pl",
+    "el",
+    "ar",
+    "ja",
+    "zh",
+    "vi",
+    "ko"
+  ],
+  "transf_decoder": {
+    "config_dict": {
+      "attn_layer_dropout": 0,
+      "attn_score_dropout": 0,
+      "diffusion": 32,
+      "embedding_dropout": 0,
+      "ffn_dropout": 0,
+      "hidden_act": "relu",
+      "hidden_size": 1024,
+      "inner_size": 4096,
+      "learn_positional_encodings": false,
+      "lm_dec_hidden": 1280,
+      "max_sequence_length": 1024,
+      "num_attention_heads": 8,
+      "num_layers": 8,
+      "num_token_types": 0,
+      "pre_ln": true,
+      "vocab_size": "None"
+    },
+    "encoder": null,
+    "model_name": null,
+    "pre_ln_final_layer_norm": true,
+    "pretrained": false
+  },
+  "transf_encoder": {
+    "attn_layer_dropout": 0,
+    "attn_score_dropout": 0,
+    "ffn_dropout": 0,
+    "hidden_size": 1024,
+    "inner_size": 4096,
+    "mask_future": false,
+    "num_attention_heads": 8,
+    "num_layers": 0,
+    "pre_ln": true,
+    "pre_ln_final_layer_norm": true
+  },
+  "use_loss_mask_for_prompt": false,
+  "vocab_size": 16384
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e683346a1a0c860342513dfb775add615b6e83375a5d73e2b4fa197e14d185c
+size 4367048880

train.py

@@ -0,0 +1,520 @@
+import json
+import logging
+import math
+import os
+import time
+from functools import partial
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from huggingface_hub import HfApi
+from safetensors import safe_open
+from torch import nn
+from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention
+from torch.utils.tensorboard import SummaryWriter
+from tqdm.auto import tqdm
+
+LOGGER = logging.getLogger(__name__)
+HF_TOKEN = os.environ['HF_TOKEN']
+API = HfApi(token=HF_TOKEN)
+
+
+class FixedPositionalEncoding(nn.Module):
+    def __init__(self, hidden_size, max_sequence_length=512):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.max_sequence_length = max_sequence_length
+
+        pos_enc = torch.zeros(max_sequence_length, hidden_size)
+        position = torch.arange(0.0, max_sequence_length).unsqueeze(1)
+        coef = -math.log(10000.0) / hidden_size
+        div_term = torch.exp(coef * torch.arange(0.0, hidden_size, 2))
+        pos_enc[:, 0::2] = torch.sin(position * div_term)
+        pos_enc[:, 1::2] = torch.cos(position * div_term)
+        pos_enc.div_(math.sqrt(hidden_size))
+        self.register_buffer('pos_enc', pos_enc)
+
+    def forward(self, position_ids):
+        return torch.index_select(self.pos_enc, 0, position_ids.reshape(-1)).reshape(*position_ids.shape, -1)
+
+
+class DecoderAttention(nn.Module):
+    def __init__(self, hidden_size, num_heads, layer_idx, kv=True):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.layer_idx = layer_idx
+        self.head_dim = hidden_size // num_heads
+        self.scale = self.head_dim**-0.5
+        self.query_net = nn.Linear(hidden_size, hidden_size)
+        if kv:
+            self.key_net = nn.Linear(hidden_size, hidden_size)
+            self.value_net = nn.Linear(hidden_size, hidden_size)
+        self.out_projection = nn.Linear(hidden_size, hidden_size)
+
+    def _reshape(self, x):
+        b, t, _ = x.shape
+        return x.view(b, t, self.num_heads, self.head_dim).transpose(1, 2)
+
+    def forward(
+        self,
+        hidden_states,
+        context_states=None,
+        attention_mask=None,
+        past_key_values=None,
+        diffusion=False,
+    ):
+        self_attn = context_states is None
+        name = 'self' if self_attn else 'cross'
+        if self_attn:
+            context_states = hidden_states
+        bsz, tgt_len, _ = hidden_states.size()
+        query = self._reshape(self.query_net(hidden_states))
+
+        if diffusion:
+            ar_key = past_key_values[f'{self.layer_idx}.{name}.key']
+            ar_value = past_key_values[f'{self.layer_idx}.{name}.value']
+            if self_attn:
+                diff_key = self._reshape(self.key_net(context_states))
+                diff_value = self._reshape(self.value_net(context_states))
+                key = torch.cat((ar_key, diff_key), dim=2)
+                value = torch.cat((ar_value, diff_value), dim=2)
+            else:
+                key = ar_key
+                value = ar_value
+        else:
+            key = self._reshape(self.key_net(context_states))
+            value = self._reshape(self.value_net(context_states))
+            past_key_values[f'{self.layer_idx}.{name}.key'] = key
+            past_key_values[f'{self.layer_idx}.{name}.value'] = value
+
+        attn_output = flex_attention(query, key, value, block_mask=attention_mask, scale=self.scale)
+        attn_output = attn_output.transpose(1, 2).reshape(bsz, tgt_len, self.hidden_size)
+        return self.out_projection(attn_output)
+
+
+class DecoderFeedForward(nn.Module):
+    def __init__(self, hidden_size, inner_size, hidden_act='relu'):
+        super().__init__()
+        self.dense_in = nn.Linear(hidden_size, inner_size)
+        assert hidden_act == 'relu'
+        self.activation = nn.ReLU()
+        self.dense_out = nn.Linear(inner_size, hidden_size)
+
+    def forward(self, x):
+        return self.dense_out(self.activation(self.dense_in(x)))
+
+
+class TransformerDecoderLayer(nn.Module):
+    def __init__(self, hidden_size, inner_size, num_heads, diffusion, layer_idx, hidden_act='relu'):
+        super().__init__()
+        self.layer_norm_1 = nn.LayerNorm(hidden_size)
+        self.first_sub_layer = DecoderAttention(hidden_size, num_heads, layer_idx=layer_idx)
+        self.layer_norm_2 = nn.LayerNorm(hidden_size)
+        self.second_sub_layer = DecoderAttention(hidden_size, num_heads, layer_idx=layer_idx, kv=not diffusion)
+        self.layer_norm_3 = nn.LayerNorm(hidden_size)
+        self.third_sub_layer = DecoderFeedForward(hidden_size, inner_size, hidden_act=hidden_act)
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        self_attention_mask=None,
+        cross_attention_mask=None,
+        past_key_values=None,
+        diffusion=False,
+    ):
+        residual = hidden_states
+        hidden_states = self.layer_norm_1(hidden_states)
+        self_out = self.first_sub_layer(
+            hidden_states,
+            context_states=None,
+            attention_mask=self_attention_mask,
+            past_key_values=past_key_values,
+            diffusion=diffusion,
+        )
+        hidden_states = residual + self_out
+
+        residual = hidden_states
+        hidden_states = self.layer_norm_2(hidden_states)
+        cross_out = self.second_sub_layer(
+            hidden_states,
+            context_states=encoder_hidden_states,
+            attention_mask=cross_attention_mask,
+            past_key_values=past_key_values,
+            diffusion=diffusion,
+        )
+        hidden_states = residual + cross_out
+
+        residual = hidden_states
+        hidden_states = self.layer_norm_3(hidden_states)
+        hidden_states = residual + self.third_sub_layer(hidden_states)
+        return hidden_states
+
+
+class TransformerDecoderEmbedding(nn.Module):
+    def __init__(self, vocab_size, hidden_size, max_sequence_length, padding_idx=2):
+        super().__init__()
+        self.token_embedding = nn.Embedding(vocab_size, hidden_size, padding_idx)
+        self.position_embedding = FixedPositionalEncoding(hidden_size, max_sequence_length)
+        self.layer_norm = nn.LayerNorm(hidden_size)
+
+    def forward(self, input_ids, positions):
+        token_embeds = self.token_embedding(input_ids)
+        if positions is None:
+            pos_embeds = self.position_embedding.pos_enc[:input_ids.shape[-1]]
+        else:
+            pos_embeds = self.position_embedding(positions)
+        return self.layer_norm(token_embeds + pos_embeds)
+
+
+class TransformerDecoderCore(nn.Module):
+    def __init__(self, hidden_size, inner_size, num_heads, num_layers, diffusion, hidden_act='relu'):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                TransformerDecoderLayer(hidden_size, inner_size, num_heads, diffusion, layer_idx=i, hidden_act=hidden_act)
+                for i in range(num_layers)
+            ]
+        )
+        self.final_layer_norm = nn.LayerNorm(hidden_size)
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        self_attention_mask=None,
+        cross_attention_mask=None,
+        past_key_values=None,
+        diffusion=False
+    ):
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                self_attention_mask=self_attention_mask,
+                cross_attention_mask=cross_attention_mask,
+                past_key_values=past_key_values,
+                diffusion=diffusion,
+            )
+        return self.final_layer_norm(hidden_states)
+
+
+class TransformerDecoderWrapper(nn.Module):
+    def __init__(self, config, diffusion=False):
+        super().__init__()
+        if not isinstance(config, dict):
+            with open(config, 'r', encoding='utf-8') as file:
+                config = json.load(file)
+        dec_config = config['transf_decoder']['config_dict']
+        hidden_size = dec_config['hidden_size']
+        self._embedding = TransformerDecoderEmbedding(
+            vocab_size=config['vocab_size'],
+            hidden_size=hidden_size,
+            max_sequence_length=dec_config['max_sequence_length'],
+            padding_idx=2,
+        )
+        self._decoder = TransformerDecoderCore(
+            hidden_size=hidden_size,
+            inner_size=dec_config['inner_size'],
+            num_heads=dec_config['num_attention_heads'],
+            num_layers=dec_config['num_layers'],
+            diffusion=diffusion,
+            hidden_act=dec_config.get('hidden_act', 'relu'),
+        )
+        self.diffusion = diffusion
+        self._lm_head = nn.Linear(hidden_size, config['vocab_size'])
+        self._lm_head.weight = self._embedding.token_embedding.weight
+
+    def load(self, model: str):
+        state = {}
+        with safe_open(model, 'pt') as file:
+            for key in file.keys():
+                if self.diffusion and ('.second_sub_layer.key_net.' in key or '.second_sub_layer.value_net.' in key):
+                    continue
+                elif key.startswith('transf_decoder.'):
+                    state[key.removeprefix('transf_decoder.')] = file.get_tensor(key)
+                elif key == 'log_softmax.mlp.layer0.bias':
+                    state['_lm_head.bias'] = file.get_tensor(key)
+                elif key == 'log_softmax.mlp.layer0.weight':
+                    state['_lm_head.weight'] = file.get_tensor(key)
+        self.load_state_dict(state)
+        return self
+
+    def forward(
+        self,
+        input_ids,
+        positions=None,
+        encoder_hidden_states=None,
+        self_attention_mask=None,
+        cross_attention_mask=None,
+        past_key_values=None,
+        diffusion=False
+    ):
+        hidden_states = self._embedding(input_ids, positions)
+        hidden_states = self._decoder(
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            self_attention_mask=self_attention_mask,
+            cross_attention_mask=cross_attention_mask,
+            past_key_values=past_key_values,
+            diffusion=diffusion,
+        )
+        return self._lm_head(hidden_states)
+
+
+def cosine_schedule(step: int, *, warmup_steps: int, max_steps: int) -> float:
+    if step < warmup_steps:
+        return step / warmup_steps
+    progress = (step - warmup_steps) / (max_steps - warmup_steps)
+    return 0.5 * (1 + math.cos(math.pi * progress))
+
+
+@torch.no_grad()
+def load(model_id: str, hub_id: str, warmup_steps: int, max_steps: int, output_dir: Path) -> tuple[
+    int,
+    TransformerDecoderWrapper,
+    TransformerDecoderCore,
+    torch.optim.AdamW,
+    torch.optim.lr_scheduler.LambdaLR,
+]:
+    config = API.hf_hub_download(model_id, 'config.json')
+    teacher = TransformerDecoderWrapper(config).to(device='cuda', dtype=torch.bfloat16).eval()
+    student = TransformerDecoderWrapper(config, diffusion=True).to(device='cuda', dtype=torch.bfloat16)
+    model = API.hf_hub_download(model_id, 'model.safetensors')
+    teacher.load(model)
+    student.load(model)
+    student = student._decoder
+    teacher.requires_grad_(False)
+    cosine_lr = partial(cosine_schedule, warmup_steps=warmup_steps, max_steps=max_steps)
+
+    decay_params = []
+    no_decay_params = []
+    for param in student.parameters():
+        if param.ndim < 2:
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+    optimizer = torch.optim.AdamW([
+        {'params': decay_params, 'weight_decay': 0.01},
+        {'params': no_decay_params, 'weight_decay': 0.0},
+    ], lr=4e-3, betas=(0.9, 0.98), eps=1e-8)
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, cosine_lr)
+
+    cur_step = 0
+    checkpoints = [x.path for x in API.list_bucket_tree(hub_id) if x.path.endswith('.pt')]
+    checkpoints.sort(key=lambda x: int(x.split('.')[0]))
+    if checkpoints:
+        checkpoint = checkpoints[-1]
+        LOGGER.info('Checkpoint found %s', checkpoint)
+        API.download_bucket_files(hub_id, files=[(checkpoint, output_dir/checkpoint)])
+        state = torch.load(output_dir/checkpoint, weights_only=True)
+        cur_step = state['step']
+        student.load_state_dict(state['model'])
+        optimizer.load_state_dict(state['optimizer'])
+        scheduler.load_state_dict(state['scheduler'])
+    else:
+        LOGGER.info('Checkpoint not detected')
+
+    return cur_step, teacher, student, optimizer, scheduler
+
+
+@torch.compile(mode='default', fullgraph=True, dynamic=False)
+def forward(
+    teacher: TransformerDecoderWrapper,
+    student: TransformerDecoderCore,
+    encoder_states: torch.Tensor,  # [b,t,d]
+    input_ids: torch.LongTensor,  # [b,s]
+    indices: torch.LongTensor,  # [b,u]
+    causal_attn_mask: BlockMask,
+    self_attn_mask: BlockMask,
+    cross_attn_mask: BlockMask,
+    bsz: int,
+    ar_len: int,
+    diff_len: int,
+    kv_len: int,
+    diff_block: int,
+    diff_blocks: int,
+    pad_id: int,
+    eot_id: int,
+) -> torch.Tensor:
+    past_key_values = {}
+    with torch.no_grad():
+        teacher_attn_mask = cross_attn_mask._adjust(ar_len, kv_len)
+        teacher_attn_mask.seq_lengths = ar_len, kv_len
+        ar_logits = teacher(
+            input_ids=input_ids,
+            positions=None,
+            encoder_hidden_states=encoder_states,
+            self_attention_mask=causal_attn_mask,
+            cross_attention_mask=teacher_attn_mask,
+            past_key_values=past_key_values,
+            diffusion=False,
+        )
+        ar_logprobs = torch.log_softmax(ar_logits, dim=2, dtype=torch.float32)
+
+        positions = (indices[:, :, None] + torch.arange(0, diff_block, device=input_ids.device)).view(bsz, diff_len)
+        valid = positions < ar_len
+        pad_mask = valid & (input_ids.gather(1, torch.where(valid, positions, 0)) > eot_id)
+        target = ar_logprobs.take_along_dim(torch.clamp_max(positions, ar_len-1)[:, :, None], dim=1)
+        del ar_logits, ar_logprobs
+        train_ids = input_ids.gather(1, indices)[:, :, None]
+        train_ids = F.pad(train_ids, (0, diff_block - 1), 'constant', pad_id).reshape(bsz, diff_len)
+        diff_states = teacher._embedding(train_ids, positions)
+
+    diff_states = student(
+        diff_states,
+        encoder_hidden_states=encoder_states,
+        self_attention_mask=self_attn_mask,
+        cross_attention_mask=cross_attn_mask,
+        past_key_values=past_key_values,
+        diffusion=True,
+    )
+    diff_logits = teacher._lm_head(diff_states)
+    diff_logprobs = torch.log_softmax(diff_logits, dim=2, dtype=torch.float32)
+
+    kld = (target.exp() * (target - diff_logprobs)).sum(dim=2)  # [b,s]
+    weight = torch.exp(torch.arange(diff_block, dtype=torch.float32, device=kld.device) / -12)
+    kld.view(bsz, diff_blocks, diff_block).mul_(weight)
+    return (kld * pad_mask).sum() / pad_mask.sum()
+
+
+def train_step(
+    encoder_states: torch.Tensor,  # [b,t,d]
+    encoder_lengths: torch.LongTensor,  # [b]
+    input_ids: torch.LongTensor,  # [b,s]
+    teacher: TransformerDecoderWrapper,
+    student: TransformerDecoderCore,
+    optimizer: torch.optim.AdamW,
+    scheduler: torch.optim.lr_scheduler.LambdaLR,
+) -> dict[str, torch.Tensor]:
+    device = encoder_states.device
+    bsz = encoder_states.shape[0]
+    ar_len = input_ids.shape[1]
+    diff_len = 2048
+    kv_len = encoder_states.shape[1]
+    diff_block = 32
+    ar_blocks = ar_len // diff_block
+    diff_blocks = diff_len // diff_block
+    prompt_len = 9
+
+    pad_id = 2
+    eot_id = 3
+    specials = 254
+    input_lengths = (input_ids > eot_id).cumprod(dim=1).sum(dim=1)
+    weights = (input_ids > specials).bfloat16()
+    weights[:, prompt_len] = 1
+    indices = torch.multinomial(weights, diff_blocks).sort(dim=1).values
+    _create_block_mask = torch.compile(create_block_mask, fullgraph=True, dynamic=False)
+
+    def causal_mask(b, h, q_idx, kv_idx):
+        return kv_idx <= q_idx
+    causal_attn_mask = _create_block_mask(causal_mask, B=None, H=None, Q_LEN=ar_len, KV_LEN=ar_len, device=device)
+
+    def self_mask(b, h, q_idx, kv_idx):
+        block = q_idx // diff_block == kv_idx // diff_block - ar_blocks
+        idx = q_idx // diff_block
+        prefix = kv_idx < indices[b, idx]
+        return (block | prefix) & (idx < input_lengths[b])
+    self_attn_mask = _create_block_mask(self_mask, B=bsz, H=None, Q_LEN=diff_len, KV_LEN=ar_len+diff_len, device=device)
+
+    def cross_mask(b, h, q_idx, kv_idx):
+        return kv_idx < encoder_lengths[b]
+    cross_attn_mask = _create_block_mask(cross_mask, B=bsz, H=None, Q_LEN=diff_len, KV_LEN=kv_len, device=device)
+
+    optimizer.zero_grad()
+    loss = forward(
+        teacher=teacher,
+        student=student,
+        encoder_states=encoder_states,
+        input_ids=input_ids,
+        indices=indices,
+        causal_attn_mask=causal_attn_mask,
+        self_attn_mask=self_attn_mask,
+        cross_attn_mask=cross_attn_mask,
+        bsz=bsz,
+        ar_len=ar_len,
+        diff_len=diff_len,
+        kv_len=kv_len,
+        diff_block=diff_block,
+        diff_blocks=diff_blocks,
+        pad_id=pad_id,
+        eot_id=eot_id,
+    )
+    loss.backward()
+    grad_norm = nn.utils.clip_grad_norm_(student.parameters(), 1.0)
+    optimizer.step()
+    scheduler.step()
+
+    return {
+        'train/loss': loss.detach().clone(),
+        'train/grad_norm': grad_norm.detach().clone(),
+        'train/learning_rate': scheduler.get_last_lr()[0],
+    }
+
+
+def main():
+    model_id = 'efwkjn/cohere-asr-ja'
+    hub_id = 'efwkjn/checkpoints'
+    output_dir = Path('checkpoints')
+    summary_writer = SummaryWriter(log_dir=output_dir/'runs')
+    handler = logging.StreamHandler()
+    formatter = logging.Formatter('%(levelname)s: %(message)s')
+    handler.setFormatter(formatter)
+    LOGGER.addHandler(handler)
+    LOGGER.setLevel(logging.INFO)
+    API.create_bucket(hub_id, private=True, exist_ok=True)
+
+    max_steps = 2**17
+    warmup_steps = 2**13
+    save_steps = 2**10
+    logging_steps = 2**4
+
+    cur_step, teacher, student, optimizer, scheduler = load(model_id, hub_id, warmup_steps, max_steps, output_dir)
+    dataset = None
+    train_metrics: list[dict[str, torch.Tensor | float]] = []
+    time_start = time.perf_counter()
+
+    for batch in tqdm(dataset, initial=cur_step, total=max_steps):
+        cur_step += 1
+        metrics = train_step(
+            **batch,
+            teacher=teacher,
+            student=student,
+            optimizer=optimizer,
+            scheduler=scheduler,
+        )
+
+        if cur_step % logging_steps == 0:
+            metrics['step'] = cur_step
+            train_metrics.append(metrics)
+            if cur_step % (logging_steps * 16) == 0:
+                prev_metrics = train_metrics[-2]
+                s = ' | '.join(f'{k[6:]}: {v.item():.5f}' for k, v in prev_metrics.items()
+                               if k.startswith('train/') and isinstance(v, torch.Tensor))
+                LOGGER.info(f'{prev_metrics["step"]}: {s}')
+
+        if cur_step % save_steps == 0:
+            train_time = time.perf_counter() - time_start
+            summary_writer.add_scalar('train/time', train_time, cur_step)
+            for m in train_metrics:
+                step = m.pop('step')
+                for k, v in m.items():
+                    summary_writer.add_scalar(k, v, step)
+            summary_writer.flush()
+            train_metrics = []
+            time_start = time.perf_counter()
+
+            checkpoint = output_dir/f'{cur_step}.pt'
+            torch.save({
+                'step': cur_step,
+                'model': student.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'scheduler': scheduler.state_dict(),
+            }, checkpoint)
+            API.run_as_future(API.sync_bucket, str(output_dir), f'hf://buckets/{hub_id}', ignore_times=True)
+
+
+if __name__ == '__main__':
+    main()