From a5ba081fbb8ffd6983450f4e680bcf2da6630fb2 Mon Sep 17 00:00:00 2001 From: gilfordting <114829710+gilfordting@users.noreply.github.com> Date: Fri, 14 Aug 2026 20:10:09 -0400 Subject: [PATCH] [Spec] Wire DFLASH aux-hidden capture into the Qwen3.5 text-only wrapper (#34771) Co-authored-by: Claude Fable 5 --- python/sglang/srt/models/qwen3_5_text.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/python/sglang/srt/models/qwen3_5_text.py b/python/sglang/srt/models/qwen3_5_text.py index 8b5e041da..f045c2ce0 100644 --- a/python/sglang/srt/models/qwen3_5_text.py +++ b/python/sglang/srt/models/qwen3_5_text.py @@ -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,