remove hacky cache size, add comment for magic number

Author: Gasoonjia

torchchat/generate.py

@@ -604,7 +604,8 @@ def generate(
                     self.is_torchtune_model
                     or self.model.config.model_type == ModelType.Flamingo
                 ):
-                    model.setup_caches(batch_size=1, dtype=self.dtype, encoder_max_seq_len=6404, decoder_max_seq_len=max_seq_length-1)
+                    # 6404 is one-gpu affordable max_seq_length for single image input
+                    model.setup_caches(batch_size=1, dtype=self.dtype, encoder_max_seq_len=6404, decoder_max_seq_len=T_new)
                 else:
                     model.setup_caches(max_batch_size=1, max_seq_length=max_seq_length)
                 if is_speculative and draft_model is not model:
@@ -758,14 +759,15 @@ def chat(
                 data = transform({"messages": messages}, inference=True)
                 batch = padded_collate_tiled_images_and_mask([data], pad_direction="left", pad_max_images=1)
                 seq_len = len(data["tokens"])
+                total_response_length = seq_len + generator_args.max_new_tokens
                 batch["causal_mask"] = torch.tril(
                                             torch.ones(
-                                                size=(generator_args.max_new_tokens, generator_args.max_new_tokens),
+                                                size=(total_response_length, total_response_length),
                                                 dtype=torch.bool,
                                             )
                                         )
                 batch["encoder_mask"] = batch["encoder_mask"][:, :seq_len]
-                encoded = batch["tokens"]
+                encoded = batch["tokens"].view(-1)
 
         else:
             encoded = self.encode_tokens(