KV Cache in BLT Generation.

Thank you for your insight work! I find kv cache in your generate.py file, but it can't run on BLT generation in:

self.setup_prefilling(lengths=lengths)
prefill_out = self.model.forward(
tokens,
tok_idx=self.prefill_tok_id,
mask=self.prefill_mask,
attn_impl="flex_attention",
)

reporting:

TypeError: ByteLatentTransformer.forward() got an unexpected keyword argument 'mask'

TypeError: ByteLatentTransformer.forward() got an unexpected keyword argument 'tok_idx'

How can I use kv cache in blt generation?