Fix pulled from Instruct model

config.json

@@ -1,8 +1,7 @@
 {
-  "architectures": [
-    "StableDiffcoderForCausalLM"
-  ],
+  "architectures": ["StableDiffcoderForCausalLM"],
   "auto_map": {
+    "AutoModel": "modeling_stable_diffcoder.StableDiffcoderForCausalLM",
     "AutoModelForCausalLM": "modeling_stable_diffcoder.StableDiffcoderForCausalLM"
   },
   "attention_bias": false,
@@ -21,11 +20,11 @@
   "num_hidden_layers": 32,
   "num_key_value_heads": 8,
   "resid_pdrop": 0.1,
-  "rms_norm_eps": 1e-06,
+  "rms_norm_eps": 1e-6,
   "rope_theta": 500000.0,
   "tie_word_embeddings": false,
   "torch_dtype": "bfloat16",
   "transformers_version": "5.3.0",
   "use_cache": true,
   "vocab_size": 155136
-}
+}

modeling_stable_diffcoder.py

@@ -137,8 +137,10 @@ class StableDiffcoderForCausalLM(LlamaForCausalLM):
         prompt_length = input_ids.shape[1]
         gen_block_list = [block_length for _ in range(gen_blocks)]
 
-        res_block = block_length - (prompt_length % block_length)
-        if res_block > 0:
+        # Fix 3: Only handle residual blocks if the prompt length is NOT cleanly divisible
+        remainder = prompt_length % block_length
+        if remainder != 0:
+            res_block = block_length - remainder
             gen_block_list = [res_block] + gen_block_list
             gen_block_list[-1] = block_length - res_block
             gen_blocks += 1
@@ -156,16 +158,23 @@ 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],
                 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 +191,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 +214,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],