Add unsqueeze to attention bias computation in DogeAttention

algo-home · algo-home · commit 7c4b102dcf27 · 2025-10-13T16:16:39.000+08:00
diff --git a/examples/modeling/modeling_doge.py b/examples/modeling/modeling_doge.py
@@ -217,7 +217,7 @@ def forward(
         dt_states = self.dt_proj(
             value_states.transpose(1, 2).reshape(value_states.shape[0], value_states.shape[-2], -1)
         )
-        attn_bias = torch.exp(self.A * F.softplus(dt_states)).transpose(-1, -2).to(hidden_states.dtype)
+        attn_bias = torch.exp(self.A * F.softplus(dt_states)).transpose(-1, -2).unsqueeze(-2).to(hidden_states.dtype)
 
         attention_interface: Callable = flash_dynamic_mask_attention_forward
 

Original file line number	Diff line number	Diff line change
`@@ -217,7 +217,7 @@ def forward(`
`217`	`217`	`dt_states = self.dt_proj(`
`218`	`218`	`value_states.transpose(1, 2).reshape(value_states.shape[0], value_states.shape[-2], -1)`
`219`	`219`	`)`
`220`		`- attn_bias = torch.exp(self.A * F.softplus(dt_states)).transpose(-1, -2).to(hidden_states.dtype)`
	`220`	`+ attn_bias = torch.exp(self.A * F.softplus(dt_states)).transpose(-1, -2).unsqueeze(-2).to(hidden_states.dtype)`
`221`	`221`
`222`	`222`	`attention_interface: Callable = flash_dynamic_mask_attention_forward`
`223`	`223`