Fix bug for https://github.com/ByteDance-Seed/Stable-DiffCoder/issues/1

modeling_stable_diffcoder.py

@@ -156,16 +156,20 @@ class StableDiffcoderForCausalLM(LlamaForCausalLM):
         nfe = 0
         final_flag = False
         prefill_length = prompt_length // block_length * block_length
+        
         if prefill_length > 0:
             cur_attn_mask = block_diffusion_attention_mask[
                 ..., :prefill_length, :prefill_length
             ]
+            # Fix 1: Explicitly pass cache_position for newer transformers prefill
+            cache_pos = torch.arange(prefill_length, device=x.device)
             self(
                 x[:, :prefill_length],
                 past_key_values=past_key_values,
                 attention_mask=cur_attn_mask,
                 use_cache=True,
-            ).past_key_values
+                cache_position=cache_pos,
+            )
 
         for block_id, block_size in enumerate(gen_block_list):
             block_start = (
@@ -182,7 +186,7 @@ class StableDiffcoderForCausalLM(LlamaForCausalLM):
             replace_position[:, block_start:block_end] = True
 
             for token_count in num_transfer_tokens:
-                if token_count:
+                if token_count > 0:
                     nfe += 1
                     mask_map = x[:, block_start:block_end] == mask_id
                     attention_mask = block_diffusion_attention_mask[
@@ -205,22 +209,28 @@ class StableDiffcoderForCausalLM(LlamaForCausalLM):
                         remasking,
                         mask_map,
                         x[:, block_start:block_end],
-                        token_count if threshold is None else None,
+                        token_count.item() if threshold is None else None, 
                         threshold,
-                        shift=False,
+                        shift=shift,
                     )
                     x[:, block_start:block_end][transfer_map] = x0[transfer_map]
 
                 if (x[:, block_start:block_end] == mask_id).sum() == 0:
+                    
+                   # Fix 2: Calculate where the generated tokens ACTUALLY start in this block
+                    gen_start = max(block_start, prompt_length)
+                    
                     if (
                         eos_id is not None
-                        and (x[:, block_start:block_end] == eos_id).sum() > 0
+                        and gen_start < block_end
+                        and (x[:, gen_start:block_end] == eos_id).sum() > 0
                     ):
                         final_flag = True
                         x = x[:, :block_end]
-                        eos_pos = (x == eos_id).nonzero(as_tuple=True)[1][0].item()
+                        eos_pos = (x[:, gen_start:block_end] == eos_id).nonzero(as_tuple=True)[1][0].item() + gen_start
                         x[0, eos_pos:] = eos_id
                         break
+                        
                     nfe += 1
                     self(
                         x[:, block_start:block_end],
@@ -231,7 +241,7 @@ class StableDiffcoderForCausalLM(LlamaForCausalLM):
                         use_cache=True,
                         cache_position=replace_position.nonzero(as_tuple=True)[1],
                     )
-                    break
+                    break 
 
             if final_flag:
                 break