Optimize GDN state updates for decode

QEfficient/transformers/models/modeling_auto.py

@@ -128,8 +128,10 @@ def _get_gdn_full_state_update_kwargs(
     qaic_config: dict | None,
     *,
     prefill_only: bool,
+    prefill_seq_len: int | None = None,
     batch_size: int = 1,
     full_batch_size: int = 1,
+    kv_cache_batch_size: int = 1,
 ) -> dict:
     if not (qaic_config and qaic_config.get("gdn_full_state_update") is True):
         return {}
@@ -141,8 +143,10 @@ def _get_gdn_full_state_update_kwargs(
     return get_kwargs(
         qaic_config=qaic_config,
         prefill_only=prefill_only,
+        prefill_seq_len=prefill_seq_len,
         batch_size=batch_size,
         full_batch_size=full_batch_size,
+        kv_cache_batch_size=kv_cache_batch_size,
     )
 
 
@@ -1709,6 +1713,7 @@ def export(
             self.model,
             qaic_config,
             prefill_only=prefill_only,
+            prefill_seq_len=prefill_seq_len,
         )
         # TODO: move this to a DA Serving utility class
         if self.model.config.model_type in SPECIALIZED_DISAGG_SERVING_MODEL_ARCH:
@@ -2056,8 +2061,10 @@ def compile(
                 self.model,
                 qaic_config,
                 prefill_only=prefill_only,
+                prefill_seq_len=prefill_seq_len,
                 batch_size=batch_size,
                 full_batch_size=full_batch_size,
+                kv_cache_batch_size=kv_cache_batch_size,
             )
 
         if layerwise:

QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py

@@ -978,8 +978,9 @@ def forward(
         if cache_params is not None:
             conv_state_all = cache_params.conv_states[self.layer_idx]
             recurrent_state_all = cache_params.recurrent_states[self.layer_idx]
-            # Singleton disaggregated prefill replaces these states in full; decode still routes partial updates.
-            state_batch_index = None if self.gdn_full_state_update else batch_index
+            # The decoder wrapper clears batch_index only after placing a full
+            # compute batch into physical cache-slot order.
+            state_batch_index = batch_index
 
             # Continuous batching path: gather only active rows, then scatter updates back.
             if state_batch_index is not None:
@@ -1810,6 +1811,12 @@ def _get_attention_blocking_config(self) -> Optional[AttentionBlockingConfig]:
                 return blocking_config
         return None
 
+    def _uses_gdn_full_state_update(self) -> bool:
+        return any(
+            bool(getattr(getattr(layer, "linear_attn", None), "gdn_full_state_update", False))
+            for layer in getattr(self.language_model, "layers", ())
+        )
+
     def get_onnx_past_key_value_names(self, layer_idx: int, layer_state=None) -> List[str]:
         if self.config.text_config.layer_types[layer_idx] == "full_attention":
             return [f"past_key.{layer_idx}", f"past_value.{layer_idx}"]
@@ -1836,12 +1843,14 @@ def forward(
             and blocking_config is not None
             and bool(blocking_config.batch_fold)
         )
+        full_batch_state_update_cb = batch_index is not None and self._uses_gdn_full_state_update()
+        physical_batch_cb = batch_fold_cb or full_batch_state_update_cb
         gdn_num_head_blocks = get_gdn_num_head_blocks(blocking_config, batch_fold_cb)
         layerwise = is_layerwise_active()
         first_layer_window = not layerwise or QEffQwen3_5MoeTextModel._start == 0
 
-        if batch_fold_cb:
-            # Folded attention expects physical batch rows at the model input.
+        if physical_batch_cb:
+            # Full-batch retained-state updates expect physical cache-slot order.
             if first_layer_window:
                 if input_ids is not None:
                     input_ids = _batch_index_scatter(input_ids, batch_index)
@@ -1880,7 +1889,7 @@ def forward(
             )
             logit_index = position_ids[0].to(torch.int32).argmax(1, keepdim=True)
             hidden_states = outputs.last_hidden_state[torch.arange(position_ids[0].shape[0]).view(-1, 1), logit_index]
-            if batch_fold_cb:
+            if physical_batch_cb:
                 hidden_states = _batch_index_gather(hidden_states, batch_index)
             logits = self.model.lm_head(hidden_states)
             image_idx = (indices1.max() + 1).unsqueeze(0).unsqueeze(0)
@@ -1930,7 +1939,7 @@ def forward(
             )
             logit_index = position_ids[0].to(torch.int32).argmax(1, keepdim=True)
             hidden_states = outputs.last_hidden_state[torch.arange(position_ids[0].shape[0]).view(-1, 1), logit_index]
-            if batch_fold_cb:
+            if physical_batch_cb:
                 hidden_states = _batch_index_gather(hidden_states, batch_index)
             logits = self.model.lm_head(hidden_states)
             return logits, outputs.past_key_values
@@ -2058,15 +2067,30 @@ def get_gdn_full_state_update_kwargs(
         qaic_config: dict | None = None,
         *,
         prefill_only: bool,
+        prefill_seq_len: int | None = None,
         batch_size: int = 1,
         full_batch_size: int = 1,
+        kv_cache_batch_size: int = 1,
     ) -> dict:
         if not (qaic_config and qaic_config.get("gdn_full_state_update") is True):
             return {}
-        if prefill_only is not True or batch_size != 1 or full_batch_size != 1:
+        full_batch_size = batch_size if full_batch_size is None else full_batch_size
+        kv_cache_batch_size = full_batch_size if kv_cache_batch_size is None else kv_cache_batch_size
+        singleton_graph = batch_size == full_batch_size == kv_cache_batch_size == 1
+        full_batch_decode = (
+            prefill_only is False
+            and prefill_seq_len == 1
+            and full_batch_size > 0
+            and full_batch_size == kv_cache_batch_size
+        )
+        if (
+            batch_size < 1
+            or (prefill_only is not True and prefill_only is not False)
+            or not (singleton_graph or full_batch_decode)
+        ):
             raise ValueError(
-                "qaic_config['gdn_full_state_update'] requires qwen3_5_moe prefill with "
-                "prefill_only=True, batch_size=1, and full_batch_size=1"
+                "qaic_config['gdn_full_state_update'] requires qwen3_5_moe singleton prefill/decode or a "
+                "decode-only graph with prefill_seq_len=1 and full_batch_size=kv_cache_batch_size"
             )
         return {"gdn_full_state_update": True}
 
@@ -2232,8 +2256,9 @@ def get_onnx_dynamic_axes(
     ):
         num_layers = self.config.text_config.num_hidden_layers
         batch_axis_name = "full_batch_size" if continuous_batching else "batch_size"
-        input_batch_axis = "full_batch_size" if continuous_batching and batch_fold else "batch_size"
-        linear_state_batch_axis = "batch_size" if gdn_full_state_update else batch_axis_name
+        physical_batch_update = continuous_batching and (batch_fold or gdn_full_state_update)
+        input_batch_axis = "full_batch_size" if physical_batch_update else "batch_size"
+        linear_state_batch_axis = input_batch_axis if gdn_full_state_update else batch_axis_name
 
         vision_dynamic_axes = {
             "pixel_values": {0: "grid_height", 1: "grid_width"},
@@ -2287,7 +2312,7 @@ def get_dummy_inputs(
         fbs = constants.ONNX_EXPORT_EXAMPLE_FBS
         batch_fold = kwargs.pop("batch_fold", False)
         gdn_full_state_update = kwargs.pop("gdn_full_state_update", False)
-        if continuous_batching and batch_fold:
+        if continuous_batching and (batch_fold or gdn_full_state_update):
             bs = fbs
         inputs_shapes["input_ids"] = (bs, dummy_seq_len)
 

examples/image_text_to_text/models/qwen3_5_moe/qwen3_5_disagg_mode_release.py

@@ -95,7 +95,7 @@ def wrapped(*args, **kwargs):
 
 # For faster execution user can run with lesser layers, For Testing Purpose Only
 config.vision_config.depth = 5
-config.text_config.num_hidden_layers = 4
+config.text_config.num_hidden_layers = 8
 config.torch_dtype = "float16"
 layer_types = list(getattr(config.text_config, "layer_types", []))
 if len(layer_types) < config.text_config.num_hidden_layers:
@@ -162,7 +162,7 @@ def _update_retained_states(target_inputs, source_outputs):
 qaic_config = {
     "blocking_mode": "prefill_online",
     "num_kv_blocks": 16,
-    "num_q_blocks": 4,
+    "num_q_blocks": 2,
     "n_rep_chunk": 1,
     "skip_kv": True,
     "gdn_chunk_size": gdn_chunk_size,
@@ -176,6 +176,7 @@ def _update_retained_states(target_inputs, source_outputs):
 #     "blocking_mode": "kv_headpar",
 #     "num_kv_blocks": 8,
 #     "skip_kv": True,
+#     "gdn_full_state_update": True,
 # }
 
 # CL 14K BSZ512