[Spec] Wire DFLASH aux-hidden capture into the Qwen3.5 text-only wrapper (#34771)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
gilfordting
2026-08-14 17:10:09 -07:00
committed by GitHub
co-authored by Claude Fable 5
parent 6f005e4da1
commit a5ba081fbb
+10
View File
@@ -118,6 +118,16 @@ class Qwen3_5ForCausalLM(nn.Module):
torch.cuda.empty_cache()
torch.cuda.synchronize()
def set_dflash_layers_to_capture(self, layers_to_capture: list[int]):
if not self.pp_group.is_last_rank:
return
if layers_to_capture is None:
raise ValueError(
"DFLASH requires explicit layer ids for aux hidden capture."
)
self.capture_aux_hidden_states = True
self.model.set_dflash_layers_to_capture(layers_to_capture)
@torch.no_grad()
def forward(
self,