def collate_fn(batch):
            """
            Collate function that creates a flat representation for variable length flash attention.
            """
            # Separate inputs and targets
            inputs, targets = zip(*batch)
            
            # Get sequence lengths
            seq_lens = torch.tensor([len(x) for x in inputs], dtype=torch.int32)
            
            # Concatenate inputs and targets into single tensors
            flat_inputs = torch.cat(inputs)
            flat_targets = torch.cat(targets)
            
            # Create cumulative sequence lengths tensor
            cu_seqlens = torch.zeros(len(seq_lens) + 1, dtype=torch.int32)
            cu_seqlens[1:] = torch.cumsum(seq_lens, dim=0)
            
            # Calculate max sequence length for this batch
            max_seqlen = seq_lens.max().item()
            
            return flat_inputs, flat_targets, seq_lens, cu_seqlens, max_seqlen

            Collate function that creates a flat representation for variable length flash attention.
            
flash_attn/flash_attn_triton_amd/train.py