Add AutoModel register and code polish

config.json

@@ -3,6 +3,7 @@
     "StableDiffcoderForCausalLM"
   ],
   "auto_map": {
+    "AutoModel": "modeling_stable_diffcoder.StableDiffcoderForCausalLM",
     "AutoModelForCausalLM": "modeling_stable_diffcoder.StableDiffcoderForCausalLM"
   },
   "attention_bias": false,

modeling_stable_diffcoder.py

@@ -158,12 +158,15 @@ 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
+            # actually not necessary since transformers will automatically generate it for prefilling
+            # if unspecified, but the official `generate` method does pass it,
+            # so we follow that for consistency and to avoid potential issues in future transformers updates
             cache_pos = torch.arange(prefill_length, device=x.device)
             self(
                 x[:, :prefill_length],
@@ -211,17 +214,17 @@ class StableDiffcoderForCausalLM(LlamaForCausalLM):
                         remasking,
                         mask_map,
                         x[:, block_start:block_end],
-                        token_count.item() if threshold is None else None, 
+                        token_count.item() if threshold is None else None,
                         threshold,
                         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 gen_start < block_end
@@ -232,7 +235,7 @@ class StableDiffcoderForCausalLM(LlamaForCausalLM):
                         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],
@@ -243,7 +246,7 @@ class StableDiffcoderForCausalLM(LlamaForCausalLM):
                         use_cache=True,
                         cache_position=replace_position.nonzero(as_tuple=True)[1],
                     )
-                    break 
+                    break
 
             if final_flag:
                 break