From 983a77d273d2a63b4d7e6dc31eed0a03dbd63140 Mon Sep 17 00:00:00 2001 From: Vahid Janfaza Date: Wed, 30 Sep 2026 17:38:17 -0700 Subject: [PATCH 1/9] Add DFlash support for Qwen3.5 hybrid VLMs Enable DFlash for dense Qwen3.5 vision-language targets with fixed 16-token verification and accepted-prefix cache replay. DFlash2 and Qwen3.5-MoE support are intentionally excluded. Co-authored-by: FNU Annanya Signed-off-by: Vahid Janfaza --- QEfficient/generation/dflash_generation.py | 112 +++++++-- .../dflash_draft/modeling_dflash_draft.py | 38 ++- .../transformers/models/modeling_auto.py | 1 + .../models/qwen3_5/modeling_qwen3_5.py | 155 +++++++++++- examples/performance/dflash/README.md | 13 + .../dflash/basic_inference_vision.py | 23 +- examples/performance/dflash/utils.py | 43 +++- .../spd/test_dflash_spd_inference.py | 211 +++++++++++++++++ .../unit_test/models/test_model_quickcheck.py | 140 +++++++++++ .../transforms/test_dflash_transform.py | 224 ++++++++++++++++++ 10 files changed, 913 insertions(+), 47 deletions(-) diff --git a/QEfficient/generation/dflash_generation.py b/QEfficient/generation/dflash_generation.py index 02984ddf9f..c554403f4e 100644 --- a/QEfficient/generation/dflash_generation.py +++ b/QEfficient/generation/dflash_generation.py @@ -22,6 +22,10 @@ from qwen_vl_utils import process_vision_info from QEfficient.generation.cloud_infer import QAICInferenceSession +from QEfficient.transformers.models.qwen3_5.modeling_qwen3_5 import ( + QWEN3_5_DFLASH_CHUNK_SIZE, + QEffQwen3_5ForConditionalGeneration, +) from QEfficient.transformers.models.qwen3_vl.modeling_qwen3_vl import ( QEffQwen3VLForConditionalGeneration, ) @@ -49,6 +53,8 @@ def __init__(self, block_size: int = 10): self.block_size = block_size self.total_prefill_time = 0.0 self.tlm_decode_time = 0.0 + self.tlm_verify_time = 0.0 + self.tlm_commit_time = 0.0 self.dlm_decode_time = 0.0 self.total_accepted_tokens = 0 self.total_rejected_tokens = 0 @@ -119,6 +125,7 @@ def _run_spd_core( generated_ids: np.ndarray, metrics: SpecDecodingMetrics, prefill_start_time: float | None = None, + tlm_commit: Callable[[np.ndarray, np.ndarray, int], None] | None = None, ) -> SpecDecodingMetrics: """Run model-independent DFlash prefill, verification, and cache advancement. @@ -212,7 +219,9 @@ def run_dlm_prefill_block(hidden_states, position_ids, sub_start, sub_len): tlm_decode_start = time.time() tlm_logits, target_hidden = tlm_decode(dlm_candidates, dlm_inputs["position_ids"]) - metrics.tlm_decode_time += time.time() - tlm_decode_start + verify_time = time.time() - tlm_decode_start + metrics.tlm_verify_time += verify_time + metrics.tlm_decode_time += verify_time accepted_length = 0 rejected_flag = False @@ -249,6 +258,16 @@ def run_dlm_prefill_block(hidden_states, position_ids, sub_start, sub_len): if any(tok_id in eos_token_ids for tok_id in this_iter_gen_ids): continue_generation = False + if tlm_commit is not None: + # A commit call reuses the session's output buffers. Preserve the + # verification features that the draft model consumes next. + target_hidden = target_hidden.copy() + commit_start = time.time() + tlm_commit(dlm_candidates, dlm_inputs["position_ids"], accepted_length + 1) + commit_time = time.time() - commit_start + metrics.tlm_commit_time += commit_time + metrics.tlm_decode_time += commit_time + if not continue_generation or gen_idx >= generation_limit: break @@ -560,11 +579,16 @@ def _get_rope_model_qwen3_vl(config): key = id(config) if key not in _rope_model_cache_qwen3_vl: with torch.device("meta"): - _rope_model_cache_qwen3_vl[key] = QEffQwen3VLForConditionalGeneration._from_config(config) + model_cls = ( + QEffQwen3_5ForConditionalGeneration + if config.model_type == "qwen3_5" + else QEffQwen3VLForConditionalGeneration + ) + _rope_model_cache_qwen3_vl[key] = model_cls._from_config(config) return _rope_model_cache_qwen3_vl[key] -def compute_position_ids_qwen3_vl(input_ids, attention_mask, image_grid_thw, config): +def compute_position_ids_qwen3_vl(input_ids, attention_mask, image_grid_thw, config, mm_token_type_ids=None): """Delegates to QEffQwen3VLForConditionalGeneration.prepare_inputs_for_generation for M-RoPE construction. With image_grid_thw=None (text-only) this reduces to a plain arange position broadcast to all 4 rows with rope_deltas=0. @@ -572,6 +596,14 @@ def compute_position_ids_qwen3_vl(input_ids, attention_mask, image_grid_thw, con """ rope_model = _get_rope_model_qwen3_vl(config) batch_size, seq_len = input_ids.shape + if config.model_type == "qwen3_5": + if mm_token_type_ids is None: + mm_token_type_ids = (input_ids == config.image_token_id).to(torch.int64) + mrope, deltas = rope_model.model.get_rope_index( + input_ids, mm_token_type_ids, image_grid_thw=image_grid_thw, attention_mask=attention_mask + ) + flat = torch.arange(seq_len).view(1, 1, -1).expand(1, batch_size, -1) + return torch.cat((flat, mrope), dim=0), deltas inputs = rope_model.prepare_inputs_for_generation( inputs={ "input_ids": input_ids, @@ -651,7 +683,17 @@ def run_spd_inference_qwen3_vl( compiled_width: int = 536, batch_size: int = 1, ) -> SpecDecodingMetrics: - """Run Qwen3-VL DFlash inference for a text-only or single-image prompt.""" + """Run Qwen3-VL or Qwen3.5 DFlash for a text-only or single-image prompt.""" + is_qwen35 = tlm_config.model_type == "qwen3_5" + if is_qwen35 and (batch_size != 1 or "cache_commit" not in tlm_session.input_names): + raise ValueError("Qwen3.5 DFlash requires batch size 1 and a TLM QPC with cache_commit input.") + if is_qwen35 and ( + prompt_chunk_size not in (QWEN3_5_DFLASH_CHUNK_SIZE, 64) or not 1 < block_size <= prompt_chunk_size + ): + raise ValueError( + "Qwen3.5 DFlash requires 16-token target calls (or a legacy 64-token QPC) and a block that fits; " + "rebuild with --prefill_seq_len 16." + ) metrics = SpecDecodingMetrics(block_size=block_size) image_processing_start = time.time() @@ -667,6 +709,7 @@ def run_spd_inference_qwen3_vl( processor_inputs["attention_mask"], image_grid_thw, tlm_config, + mm_token_type_ids=processor_inputs.get("mm_token_type_ids"), ) position_ids = torch.nn.functional.pad(position_ids, (0, padded_len - input_ids_length), value=-1) input_ids = torch.nn.functional.pad( @@ -678,35 +721,47 @@ def run_spd_inference_qwen3_vl( position_ids_np = position_ids.numpy() flat_position_ids = position_ids_np[0] rope_delta = int(rope_deltas.reshape(-1)[0]) - generated_ids = _initialize_generated_ids(batch_size, ctx_len, padded_len, tokenizer.pad_token_id) + # Verification always consumes a complete block, including near the length limit. + generation_ctx_len = ctx_len - block_size if is_qwen35 else ctx_len + generated_ids = _initialize_generated_ids(batch_size, generation_ctx_len, padded_len, tokenizer.pad_token_id) metrics.image_processing_time = time.time() - image_processing_start prefill_start = time.time() vision_binding = tlm_session.bindings[tlm_session.binding_index_map["vision_embeds"]] vision_dtype = tlm_session.aic_to_np_dtype_mapping.get(vision_binding.type, np.dtype(np.float32)) - deepstack_binding = tlm_session.bindings[tlm_session.binding_index_map["deepstack_features"]] - deepstack_dtype = tlm_session.aic_to_np_dtype_mapping.get(deepstack_binding.type, np.dtype(np.float32)) + has_deepstack = "deepstack_features" in tlm_session.input_names + if has_deepstack: + deepstack_binding = tlm_session.bindings[tlm_session.binding_index_map["deepstack_features"]] + deepstack_dtype = tlm_session.aic_to_np_dtype_mapping.get(deepstack_binding.type, np.dtype(np.float32)) if image is not None: if vision_session is None: raise ValueError("An image was given but no vision_session was provided.") - vision_embeds, deepstack_features = run_vision_encoder_qwen3_vl( - vision_session, - processor_inputs["pixel_values"].numpy(), - image_grid_thw.numpy(), + vision_results = vision_session.run( + { + "pixel_values": processor_inputs["pixel_values"].numpy().astype(np.float16), + "image_grid_thw": image_grid_thw.numpy(), + } ) + vision_embeds = vision_results["vision_embeds"] + if has_deepstack: + deepstack_features = vision_results["deepstack_features"] metrics.vision_prefill_time = time.time() - prefill_start else: vision_embeds = np.zeros(tuple(int(dim) for dim in vision_binding.dims), dtype=vision_dtype) - deepstack_features = np.zeros(tuple(int(dim) for dim in deepstack_binding.dims), dtype=deepstack_dtype) + if has_deepstack: + deepstack_features = np.zeros(tuple(int(dim) for dim in deepstack_binding.dims), dtype=deepstack_dtype) metrics.vision_prefill_time = 0.0 vision_outputs = { "vision_embeds": vision_embeds.astype(vision_dtype), - "deepstack_features": deepstack_features.astype(deepstack_dtype), } + if has_deepstack: + vision_outputs["deepstack_features"] = deepstack_features.astype(deepstack_dtype) tlm_session.set_buffers(vision_outputs) lang_extra = {"image_idx": np.array([[0]], dtype=np.int64)} + if is_qwen35: + lang_extra["cache_commit"] = np.ones(1, dtype=np.int64) tlm_session.set_buffers({"logits": np.zeros((batch_size, prompt_chunk_size), dtype=np.int32)}) tlm_session.set_buffers({"hidden_states": np.zeros((batch_size, prompt_chunk_size, hidden_size), dtype=np.float32)}) @@ -722,19 +777,35 @@ def tlm_prefill(chunk_start, chunk_end): return outputs["logits"], outputs["hidden_states"] def prepare_decode(): - tlm_session.set_buffers({"logits": np.zeros((batch_size, block_size), dtype=np.int32)}) - tlm_session.set_buffers({"hidden_states": np.zeros((batch_size, block_size, hidden_size), dtype=np.float32)}) + target_seq_len = prompt_chunk_size if is_qwen35 else block_size + tlm_session.set_buffers({"logits": np.zeros((batch_size, target_seq_len), dtype=np.int32)}) + tlm_session.set_buffers( + {"hidden_states": np.zeros((batch_size, target_seq_len, hidden_size), dtype=np.float32)} + ) tlm_session.skip_buffers(vision_outputs.keys()) + if is_qwen35: + lang_extra["cache_commit"] = np.zeros(1, dtype=np.int64) + + def decode_inputs(draft_ids, decode_position_ids, valid_tokens): + positions = build_decode_position_ids_qwen3_vl(decode_position_ids, rope_delta) + if is_qwen35: + padding = prompt_chunk_size - block_size + draft_ids = np.pad(draft_ids, ((0, 0), (0, padding)), constant_values=tokenizer.pad_token_id) + positions = np.pad(positions, ((0, 0), (0, 0), (0, padding)), constant_values=-1) + positions[:, :, valid_tokens:] = -1 + return {"input_ids": draft_ids, "position_ids": positions, **lang_extra} def tlm_decode(draft_ids, decode_position_ids): - outputs = tlm_session.run( + outputs = tlm_session.run(decode_inputs(draft_ids, decode_position_ids, block_size)) + return outputs["logits"][:, :block_size], outputs["hidden_states"][:, :block_size] + + def tlm_commit(draft_ids, decode_position_ids, count): + tlm_session.run( { - "input_ids": draft_ids, - "position_ids": build_decode_position_ids_qwen3_vl(decode_position_ids, rope_delta), - **lang_extra, + **decode_inputs(draft_ids, decode_position_ids, count), + "cache_commit": np.ones(1, dtype=np.int64), } ) - return outputs["logits"], outputs["hidden_states"] eos_token_ids = {tokenizer.eos_token_id} if tokenizer.eos_token_id is not None else set() return _run_spd_core( @@ -755,4 +826,5 @@ def tlm_decode(draft_ids, decode_position_ids): generated_ids=generated_ids, metrics=metrics, prefill_start_time=prefill_start, + tlm_commit=tlm_commit if is_qwen35 else None, ) diff --git a/QEfficient/transformers/models/dflash_draft/modeling_dflash_draft.py b/QEfficient/transformers/models/dflash_draft/modeling_dflash_draft.py index 5284aae8c1..53275ec92a 100644 --- a/QEfficient/transformers/models/dflash_draft/modeling_dflash_draft.py +++ b/QEfficient/transformers/models/dflash_draft/modeling_dflash_draft.py @@ -36,6 +36,7 @@ def _create_mask( # valid_kv_length: int, sliding_window: int | None = None, start_index: int | None = 0, + is_causal: bool = False, ): """ Args: @@ -44,6 +45,7 @@ def _create_mask( valid_kv_length: number of valid KV cache positions sliding_window: optional local attention window start_index: offset into KV cache (default 0) + is_causal: restrict each query to its own position and earlier keys Returns: attention_mask: [1, 1, num_queries, target_length] @@ -52,13 +54,24 @@ def _create_mask( num_queries = position_ids.shape[1] # = block_size (B) bsz = position_ids.shape[0] - kv_positions = torch.arange(start_index, start_index + target_length) + kv_positions = torch.arange(start_index, start_index + target_length, device=position_ids.device) row_max = position_ids.max(dim=-1, keepdim=True).values # [bsz, 1] valid_kv_mask = kv_positions.view(1, target_length) > (start_index + row_max) # [bsz, target_length] # ---- Step 2: Expand to [bsz, num_queries, target_length] ---- attention_mask = valid_kv_mask.unsqueeze(1).expand(bsz, num_queries, target_length) + # Full DFlash layers see the whole noise block. Local layers use a per-query + # window, with causality selected from the draft config as in upstream DFlash. + query_positions = (start_index + position_ids).unsqueeze(-1) + key_positions = kv_positions.view(1, 1, target_length) + if is_causal: + attention_mask = attention_mask | (key_positions > query_positions) + if sliding_window is not None: + attention_mask = attention_mask | (key_positions <= query_positions - sliding_window) + if not is_causal: + attention_mask = attention_mask | (key_positions >= query_positions + sliding_window) + # ---- Step 4: Add head dimension ---- # Final shape: [bsz, 1, B, 3B] — broadcasts over heads in eager attention. attention_mask = attention_mask.unsqueeze(1) @@ -226,6 +239,11 @@ class QEffDFlashAttention(Qwen3Attention): def __qeff_init__(self): self.dflash_dlm = True + layer_types = getattr(self.config, "layer_types", None) + layer_type = layer_types[self.layer_idx] if layer_types else "full_attention" + is_causal = getattr(self.config, "is_causal", None) + self.is_causal = layer_type == "sliding_attention" if is_causal is None else bool(is_causal) + self.sliding_window = self.config.sliding_window if layer_type == "sliding_attention" else None def forward( self, @@ -425,9 +443,17 @@ def forward( position_ids = cache_position.unsqueeze(0) target_length = attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else past_seen_tokens - causal_mask = _create_mask( - position_ids=position_ids, target_length=target_length, sliding_window=self.config.sliding_window - ) + attention_masks = {} + for decoder_layer in self.layers: + attn = decoder_layer.self_attn + mask_key = (attn.is_causal, attn.sliding_window) + if mask_key not in attention_masks: + attention_masks[mask_key] = _create_mask( + position_ids=position_ids, + target_length=target_length, + sliding_window=attn.sliding_window, + is_causal=attn.is_causal, + ) hidden_states = noise_embeds # Pass the whole rotary table down: `qeff_apply_rope_two_streams` gathers it once per stream @@ -448,7 +474,9 @@ def forward( hidden_states, target_hidden=target_hidden, position_ids_target=position_ids_target, - attention_mask=causal_mask, + attention_mask=attention_masks[ + (decoder_layer.self_attn.is_causal, decoder_layer.self_attn.sliding_window) + ], position_ids=position_ids, past_key_value=past_key_values, comp_ctx_lengths=comp_ctx_lengths, diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index cd2d8ab689..8386c0866e 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -3804,6 +3804,7 @@ def __init__( if self.dflash_dlm: self.model, _ = DFlashTransform.apply(self.model, qaic_config) self.model, _ = DFlashDLMTransform.apply(self.model, qaic_config) + self.hash_params["dflash_attention"] = "per_layer_window_v1" if self.dflash_tlm: self.model, _ = DFlashTLMTransform.apply(self.model, qaic_config) diff --git a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py index aea2788a6e..2b46bbcf67 100644 --- a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -57,6 +57,7 @@ is_layerwise_active, resolve_layer_window, ) +from QEfficient.transformers.spd.dflash import compute_dflash_target_hidden_states from QEfficient.utils import constants from QEfficient.utils._utils import IOInfo, get_padding_shape_from_config from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE @@ -64,6 +65,7 @@ logger = QEFFLogger.get_logger("MODEL") QWEN3_5_ROPE_CACHE_EXPORT_CAP = 76800 +QWEN3_5_DFLASH_CHUNK_SIZE = 16 def _expand_mrope_position_ids(position_ids, cache_position, batch_size): @@ -515,9 +517,12 @@ def qeff_torch_causal_conv1d_update( shifted_conv_state = shifted_conv_state * valid_decode + conv_state_flat * (1 - valid_decode) # Prefill and padded chunks retain the general last-valid-position gather. - # Derive zeros from an existing state view; torch.zeros emits a dynamic ConstantOfShape in subfunctions. - zeros = conv_state_flat[:, 0, :].to(pos_ids.dtype) * 0 - order = torch.argsort(torch.cat([zeros, pos_ids], dim=1), dim=1) + # Give cached entries distinct chronological keys. Ties can reorder history + # during short accepted-prefix replay, and PyTorch/ONNX break them differently. + cached_order = torch.arange(state_len, dtype=pos_ids.dtype, device=pos_ids.device) + cached_order = cached_order.unsqueeze(0).expand(pos_ids.shape[0], -1) + token_order = torch.where(pos_ids >= 0, pos_ids + state_len, -1) + order = torch.argsort(torch.cat([cached_order, token_order], dim=1), dim=1) last_positions = order[:, -state_len:] ctx_idx = last_positions.to(torch.long).unsqueeze(1).expand(-1, hidden_size, -1) gathered_conv_state = torch.gather(hidden_states_new, dim=2, index=ctx_idx) @@ -808,6 +813,19 @@ def _solve_chunk_attn(self, attn: torch.Tensor, mask: torch.Tensor, eye: torch.T return self._solve_chunk_attn_scaled_newton_schulz(attn, mask, eye, chunk_size) return self._solve_chunk_attn_recursive_sns(attn, mask, eye, chunk_size) + @staticmethod + def _dflash_triangular_inverse(attn, eye, chunk_size): + if chunk_size < 2 or chunk_size & (chunk_size - 1): + raise ValueError("DFlash matrix solve requires a power-of-two chunk size >= 2.") + # For strictly lower triangular A, A**C = 0. Doubling evaluates + # I + A + ... + A**(C-1) exactly for power-of-two chunk lengths. + result = attn + eye + power = attn + for _ in range(chunk_size.bit_length() - 2): + power = power @ power + result = result + result @ power + return result + def torch_chunk_gated_delta_rule_qeff( self, query, @@ -824,6 +842,7 @@ def torch_chunk_gated_delta_rule_qeff( mask_strict=None, ones_lower=None, eye=None, + use_matrix_solve=False, ): initial_dtype = query.dtype # if use_qk_l2norm_in_kernel: @@ -895,7 +914,10 @@ def torch_chunk_gated_delta_rule_qeff( attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) I_base = eye if eye is not None else torch.eye(chunk_size, dtype=attn.dtype) - attn = self._solve_chunk_attn(attn=attn, mask=mask, eye=I_base, chunk_size=chunk_size) + if use_matrix_solve: + attn = self._dflash_triangular_inverse(attn, I_base.to(dtype=attn.dtype, device=attn.device), chunk_size) + else: + attn = self._solve_chunk_attn(attn=attn, mask=mask, eye=I_base, chunk_size=chunk_size) value = attn @ v_beta k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) @@ -957,6 +979,7 @@ def forward( batch_index: torch.LongTensor | None = None, batch_fold: bool = False, gdn_num_head_blocks: int = 1, + dflash_chunk_size: int | None = None, ): batch_size, seq_len, _ = hidden_states.shape @@ -1078,7 +1101,29 @@ def forward( # ── Recurrent State ─────────────────────────────────── if cache_params is not None: - if batch_fold: + if dflash_chunk_size is not None: + # DFlash uses a fixed multi-token shape, including prefix replay. + # Select at export time so no one-token recurrent branch is traced. + if self.gdn_chunk_size < dflash_chunk_size: + raise ValueError(f"Qwen3.5 DFlash requires gdn_chunk_size >= {dflash_chunk_size}.") + core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + position_ids=position_ids, + chunk_size=dflash_chunk_size, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + mask_causal=self._mask_causal[:dflash_chunk_size, :dflash_chunk_size], + mask_strict=self._mask_strict[:dflash_chunk_size, :dflash_chunk_size], + ones_lower=self._ones_lower[:dflash_chunk_size, :dflash_chunk_size], + eye=self._eye[:dflash_chunk_size, :dflash_chunk_size], + use_matrix_solve=getattr(self, "dflash_matrix_solve", True), + ) + elif batch_fold: # Folded decode is exported with a static one-token sequence. core_attn_out, last_recurrent_state = self._recurrent_step_batched( query, @@ -1187,6 +1232,7 @@ def forward( gdn_num_head_blocks: int = 1, use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + dflash_chunk_size: int | None = None, **kwargs, ) -> torch.FloatTensor: del use_cache @@ -1203,6 +1249,7 @@ def forward( batch_index=batch_index, batch_fold=batch_fold, gdn_num_head_blocks=gdn_num_head_blocks, + dflash_chunk_size=dflash_chunk_size, ) else: hidden_states, _ = self.self_attn( @@ -1254,6 +1301,7 @@ def forward( use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, output_hidden_states: bool | None = None, + cache_commit: torch.Tensor | None = None, **kwargs, ) -> BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): @@ -1272,12 +1320,29 @@ def forward( elif use_cache and past_key_values is None: past_key_values = QEffQwen3_5DynamicCache(self.config) + target_layer_ids = getattr(self, "target_layer_ids", None) + if target_layer_ids: + if is_layerwise_active() or kwargs.get("layer_indices_to_run") is not None: + raise ValueError("Qwen3.5 DFlash requires the full target model; layerwise execution is unsupported.") + if cache_commit is None: + raise ValueError( + "Qwen3.5 DFlash requires cache_commit: 0 for verification, 1 for accepted-prefix replay." + ) + if batch_index is not None or batch_fold: + raise ValueError("Qwen3.5 DFlash currently supports non-continuous batch size 1 only.") + if past_key_values is None or any( + past_key_values.conv_states[i] is None or past_key_values.recurrent_states[i] is None + for i, kind in enumerate(self.config.layer_types) + if kind == "linear_attention" + ): + raise ValueError("Qwen3.5 DFlash requires preallocated convolution and recurrent states.") + if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) start, end = resolve_layer_window(QEffQwen3_5TextModel, len(self.layers)) + past_seen_tokens = past_key_values.get_seq_length(layer_idx=start) if past_key_values is not None else 0 if cache_position is None: - past_seen_tokens = past_key_values.get_seq_length(layer_idx=start) if past_key_values is not None else 0 active_layer_type = self.config.layer_types[start] if start < len(self.config.layer_types) else None if is_layerwise_active() and position_ids is not None and active_layer_type == "linear_attention": text_position_ids = position_ids[0] if position_ids.ndim == 3 else position_ids @@ -1288,6 +1353,19 @@ def forward( ) position_ids = _expand_mrope_position_ids(position_ids, cache_position, inputs_embeds.shape[0]) text_position_ids = position_ids[0] if position_ids.ndim == 3 else position_ids + if target_layer_ids: + # A new prompt may reuse an already-loaded QPC. Unlike attention KV, + # recurrent state has no positional mask to hide the previous request. + reset = (text_position_ids[:, :1] == 0).reshape(-1, 1, 1, 1) + for i, kind in enumerate(self.config.layer_types): + if kind == "linear_attention": + conv = past_key_values.conv_states[i] + recurrent = past_key_values.recurrent_states[i] + conv_reset = reset if conv.ndim == 4 else reset.squeeze(-1) + past_key_values.conv_states[i] = torch.where(conv_reset, torch.zeros_like(conv), conv) + past_key_values.recurrent_states[i] = torch.where(reset, torch.zeros_like(recurrent), recurrent) + committed_conv = list(past_key_values.conv_states) + committed_recurrent = list(past_key_values.recurrent_states) rotary_position_ids = ( position_ids[1:] if position_ids.ndim == 3 and position_ids.shape[0] == 4 else position_ids ) @@ -1314,6 +1392,7 @@ def forward( ) position_embeddings = (cos, sin) all_hidden_states = () if output_hidden_states else None + target_hidden_list = [] layer_indices_to_run = kwargs.get("layer_indices_to_run", None) for layer_idx, decoder_layer in enumerate(self.layers): @@ -1323,6 +1402,8 @@ def forward( continue if output_hidden_states: all_hidden_states += (hidden_states,) + if target_layer_ids and layer_idx in target_layer_ids: + target_hidden_list.append(hidden_states) layer_mask = linear_attn_mask if decoder_layer.layer_type == "linear_attention" else causal_mask hidden_states = decoder_layer( @@ -1337,16 +1418,31 @@ def forward( gdn_num_head_blocks=gdn_num_head_blocks, use_cache=use_cache, cache_position=cache_position, + dflash_chunk_size=getattr(self, "dflash_chunk_size", QWEN3_5_DFLASH_CHUNK_SIZE) + if target_layer_ids + else None, **kwargs, ) - # break - + if target_layer_ids and self.config.num_hidden_layers in target_layer_ids: + target_hidden_list.append(hidden_states) if is_last_layer_window(QEffQwen3_5TextModel, len(self.layers)): hidden_states = self.norm(hidden_states) if output_hidden_states: all_hidden_states += (hidden_states,) + if target_layer_ids: + # Verification may overwrite attention KV slots, but recurrent states + # cannot be rolled back by positions. Only an accepted-prefix pass commits them. + for i, kind in enumerate(self.config.layer_types): + if kind == "linear_attention": + past_key_values.conv_states[i] = torch.where( + cache_commit.bool(), past_key_values.conv_states[i], committed_conv[i] + ) + past_key_values.recurrent_states[i] = torch.where( + cache_commit.bool(), past_key_values.recurrent_states[i], committed_recurrent[i] + ) + if return_legacy_cache: past_key_values = past_key_values.to_legacy_cache() @@ -1355,6 +1451,9 @@ def forward( return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None, + hidden_states=compute_dflash_target_hidden_states(target_hidden_list, self.fc, self.hidden_norm) + if target_layer_ids + else all_hidden_states, ) @@ -1472,7 +1571,9 @@ def forward( **kwargs, ) - if position_ids is None: + if getattr(self.model, "target_layer_ids", None): + hidden_states = outputs.last_hidden_state + elif position_ids is None: hidden_states = outputs.last_hidden_state[:, -1:, :] else: text_position_ids = position_ids[0] if position_ids.ndim == 3 else position_ids @@ -1815,6 +1916,7 @@ def forward( past_key_values=None, batch_index: Optional[torch.LongTensor] = None, comp_ctx_lengths: Optional[List[int]] = None, + cache_commit: torch.Tensor | None = None, ): # Continuous batching also supplies batch_index for prefill. Folded # attention is decode-only, so do not infer it from batch_index alone. @@ -1853,13 +1955,23 @@ def forward( batch_fold=batch_fold_cb, gdn_num_head_blocks=gdn_num_head_blocks, use_cache=True, + cache_commit=cache_commit, ) + image_idx = (indices1.max() + 1).unsqueeze(0).unsqueeze(0) + if getattr(self.language_model, "target_layer_ids", None): + logits = self.model.lm_head(outputs.last_hidden_state).float() + return ( + logits.argmax(dim=-1).to(torch.int32), + vision_embeds, + image_idx, + outputs.hidden_states, + outputs.past_key_values, + ) 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: 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) return logits, vision_embeds, image_idx, outputs.past_key_values[: len(past_key_values)] @@ -2010,6 +2122,10 @@ def get_specializations( num_frames = [num_frames] * len(height) if isinstance(num_frames, int) else num_frames prefill_seq_len = prefill_seq_len if prefill_seq_len else 128 + if getattr(language_model, "target_layer_ids", None): + chunk_size = getattr(language_model, "dflash_chunk_size", QWEN3_5_DFLASH_CHUNK_SIZE) + if prefill_seq_len != chunk_size: + raise ValueError(f"Qwen3.5 DFlash requires prefill_seq_len={chunk_size} for its gated-delta chunk.") ctx_len = ctx_len if ctx_len else constants.INTERN_CTX_LEN kv_cache_batch_size = kv_cache_batch_size or full_batch_size or batch_size channel = 3 @@ -2133,6 +2249,9 @@ def get_onnx_dynamic_axes( "position_ids": {1: input_batch_axis, 2: "seq_len"}, "vision_embeds": {0: "vision_batch_size", 1: "vision_size"}, } + language_model = getattr(getattr(self, "model", None), "language_model", None) + if getattr(language_model, "target_layer_ids", None): + lang_dynamic_axes["hidden_states"] = {0: input_batch_axis, 1: "seq_len"} for i in range(num_layers): if self.config.text_config.layer_types[i] == "full_attention": @@ -2169,6 +2288,18 @@ def get_dummy_inputs( inputs_shapes = {} dummy_seq_len = int(kwargs.get("prefill_seq_len", 32)) + is_dflash = bool(getattr(self.model.language_model, "target_layer_ids", None)) + if is_dflash: + if not kv_offload or continuous_batching: + raise ValueError("Qwen3.5 DFlash requires kv_offload=True and continuous_batching=False.") + if not 1 < dummy_seq_len <= 64: + raise ValueError( + "Qwen3.5 DFlash requires 2 <= prefill_seq_len <= 64 for the exported gated-delta graph." + ) + # The shared exporter supplies a 32-token dummy. Trace only one + # chunk: Python chunk loops do not generalize across chunk counts. + chunk_size = getattr(self.model.language_model, "dflash_chunk_size", QWEN3_5_DFLASH_CHUNK_SIZE) + dummy_seq_len = min(dummy_seq_len, chunk_size) bs = kwargs.get("batch_size", constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE) if bs > 1: bs = 2 @@ -2211,6 +2342,8 @@ def get_dummy_inputs( .repeat(4, 1, 1) ) lang_inputs["image_idx"] = torch.zeros((inputs_shapes["image_idx"]), dtype=torch.int64) + if is_dflash: + lang_inputs["cache_commit"] = torch.ones(1, dtype=torch.int64) kv_cache_shape = get_padding_shape_from_config( config=self.model.config.text_config, @@ -2269,6 +2402,8 @@ def get_dummy_inputs( def get_output_names(self, kv_offload: bool = False): vision_output_names = ["vision_embeds"] lang_output_names = ["logits"] + if getattr(self.model.language_model, "target_layer_ids", None): + lang_output_names.append("hidden_states") for layer_idx, layer_type in enumerate(self.model.config.text_config.layer_types): if layer_idx >= self.model.config.text_config.num_hidden_layers: break diff --git a/examples/performance/dflash/README.md b/examples/performance/dflash/README.md index 65f4c9af27..2e775efc19 100644 --- a/examples/performance/dflash/README.md +++ b/examples/performance/dflash/README.md @@ -76,6 +76,19 @@ python basic_inference_vision.py --model_name Qwen3-VL-32B-Instruct \ `--height`/`--width` (qwen3-vl only) set the vision encoder's compiled input resolution; they default to `354`/`536` when omitted. +### Qwen3.5-4B + +Qwen3.5 uses the VLM entry point with its DFlash target and draft checkpoints: + +```bash +python basic_inference_vision.py --model_name Qwen3.5-4B \ + --tlm_devices 40,41,42,43 --dlm_devices 44,45,46,47 --vision_devices 48,49,50,51 \ + --ctx_len 4096 --prefill_seq_len 16 \ + --image --image_prompt "Describe this image in detail." +``` + +Use `--prompt "..."` instead of `--image --image_prompt ...` for text-only input. + --- ## `--model_name` diff --git a/examples/performance/dflash/basic_inference_vision.py b/examples/performance/dflash/basic_inference_vision.py index e16c42fe54..a58b943e2e 100644 --- a/examples/performance/dflash/basic_inference_vision.py +++ b/examples/performance/dflash/basic_inference_vision.py @@ -23,6 +23,7 @@ compile_gemma_vlm_qpcs, compile_qwen3vl_vlm_dlm_qpc, compile_qwen3vl_vlm_qpcs, + get_dflash_block_size, get_spd_prompt_chunk_size, load_spd_sessions, resolve_model_name, @@ -76,7 +77,7 @@ def parse_args(): parser.add_argument("--vision_devices", type=parse_device_list, default=[0, 1, 2, 3]) parser.add_argument("--tlm_cores", type=int, default=8) parser.add_argument("--dlm_cores", type=int, default=8) - parser.add_argument("--ctx_len", type=int, default=2048) + parser.add_argument("--ctx_len", type=int, default=4096) parser.add_argument("--prefill_seq_len", type=int, default=128) parser.add_argument("--generation_len", type=int, default=1800) parser.add_argument("--iteration", type=int, default=300) @@ -170,7 +171,7 @@ def _run_gemma( vocab_size=config.vocab_size, prompt_chunk_size=prompt_chunk_size, ctx_len=args.ctx_len, - block_size=config.block_size, + block_size=get_dflash_block_size(config), max_iterations=args.iteration, hidden_size=config.hidden_size, generation_len=args.generation_len, @@ -223,7 +224,7 @@ def _run_qwen3_vl( vocab_size=config.vocab_size, prompt_chunk_size=prompt_chunk_size, ctx_len=args.ctx_len, - block_size=config.block_size, + block_size=get_dflash_block_size(config), max_iterations=args.iteration, hidden_size=config.hidden_size, generation_len=args.generation_len, @@ -248,7 +249,7 @@ def main(): raise ValueError(f"No default TLM HF path for '{args.model_name}'. Pass --tlm_hf_path.") tlm_config = AutoConfig.from_pretrained(tlm_repo, trust_remote_code=True, token=args.hf_token) - is_qwen3vl = tlm_config.model_type == "qwen3_vl" + is_qwen3vl = tlm_config.model_type in {"qwen3_vl", "qwen3_5"} if args.tlm_qpc and args.vision_qpc: tlm_qpc, vision_qpc = args.tlm_qpc, args.vision_qpc @@ -296,7 +297,8 @@ def main(): dlm_session, tlm_session = load_spd_sessions(tlm_qpc, dlm_qpc, args.tlm_devices, args.dlm_devices) prompt_chunk_size = get_spd_prompt_chunk_size(tlm_session) - validate_spd_decode_specialization(tlm_session, config.block_size) + target_decode_length = prompt_chunk_size if tlm_config.model_type == "qwen3_5" else get_dflash_block_size(config) + validate_spd_decode_specialization(tlm_session, target_decode_length) if is_qwen3vl: metrics, output_extra = _run_qwen3_vl( @@ -336,7 +338,16 @@ def main(): print("=" * width) print(f" {'Acceptance Rate (tok/iter)':<30} {metrics.acceptance_rate():>6.2f}") print(f" {'DLM Throughput (tok/s)':<30} {metrics.dlm_tok_rate():>6.1f}") - print(f" {'TLM Throughput (tok/s)':<30} {metrics.tlm_tok_rate():>6.1f}") + if tlm_config.model_type == "qwen3_5": + round_rate = metrics.num_total_iters / metrics.tlm_decode_time if metrics.tlm_decode_time > 0 else 0.0 + print(f" {'TLM Rounds (round/s)':<30} {round_rate:>6.1f}") + else: + print(f" {'TLM Throughput (tok/s)':<30} {metrics.tlm_tok_rate():>6.1f}") + if metrics.tlm_commit_time > 0 and metrics.tlm_verify_time > 0 and metrics.num_total_iters > 0: + print(f" {'TLM Verify (ms/call)':<30} {1000 * metrics.tlm_verify_time / metrics.num_total_iters:>6.1f}") + print(f" {'TLM Replay (ms/call)':<30} {1000 * metrics.tlm_commit_time / metrics.num_total_iters:>6.1f}") + verified_rate = metrics.block_size * metrics.num_total_iters / metrics.tlm_verify_time + print(f" {'TLM Verified positions/s':<30} {verified_rate:>6.1f}") print(f" {'SPD Decode Speed (tok/s)':<30} {metrics.spd_tok_rate():>6.1f}") print(f" {'Generated tokens':<30} {metrics.total_generated_tokens:>6}") print(f" {'Iterations':<30} {metrics.num_total_iters:>6}") diff --git a/examples/performance/dflash/utils.py b/examples/performance/dflash/utils.py index c023366f98..a9f81c0e57 100644 --- a/examples/performance/dflash/utils.py +++ b/examples/performance/dflash/utils.py @@ -21,6 +21,7 @@ QAICInferenceSession, is_retained_state_name, ) +from QEfficient.transformers.models.qwen3_5.modeling_qwen3_5 import QWEN3_5_DFLASH_CHUNK_SIZE from QEfficient.utils.logging_utils import QEFFLogger logger = QEFFLogger.get_logger("INFRA") @@ -32,6 +33,8 @@ MODEL_MAP = { "gemma-4-31B-it": ("google/gemma-4-31B-it", "z-lab/gemma-4-31B-it-DFlash"), "Qwen3-4B": ("Qwen/Qwen3-4B", "z-lab/Qwen3-4B-DFlash-b16"), + "Qwen3.5-4B": ("Qwen/Qwen3.5-4B", "z-lab/Qwen3.5-4B-DFlash"), + "Qwen3.6-27B": ("Qwen/Qwen3.6-27B", "z-lab/Qwen3.6-27B-DFlash"), "Qwen3-8B": ("Qwen/Qwen3-8B", "z-lab/Qwen3-8B-DFlash-b16"), "Llama-3.1-8B-Instruct": ("meta-llama/Llama-3.1-8B-Instruct", "z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat"), "Qwen3-VL-32B-Instruct": ( @@ -81,7 +84,10 @@ def validate_spd_decode_specialization(tlm_session: QAICInferenceSession, block_ """Require a TLM specialization that verifies one complete DFlash block.""" input_index = tlm_session.binding_index_map["input_ids"] sequence_lengths = sorted({shape[input_index][1][1] for shape in tlm_session.allowed_shapes}) - logger.info(f"TLM input_ids seq_lens (allowed shapes): {sequence_lengths}") + if not sequence_lengths: + # Single-specialization QPCs report their fixed shape only in bindings. + sequence_lengths = [tlm_session.bindings[input_index].dims[1]] + logger.info(f"TLM input_ids seq_lens: {sequence_lengths}") if block_size not in sequence_lengths: raise ValueError( f"TLM QPC has no decode specialization with seq_len={block_size} " @@ -160,11 +166,19 @@ def load_dflash_checkpoint(dflash_model_path: str) -> tuple[dict, dict]: return state_dict, cfg +def get_dflash_block_size(config): + config = config if isinstance(config, dict) else config.to_dict() + block_size = config.get("block_size") or config.get("dflash_config", {}).get("block_size") + if not isinstance(block_size, int) or block_size < 2: + raise ValueError("DFlash config must provide block_size >= 2 (top-level or under dflash_config).") + return block_size + + def read_dlm_meta(dlm_repo: str, hf_token: str | None = None): """Load a DFlash checkpoint and return (state_dict, target_layer_ids, block_size).""" state_dict, cfg = load_dflash_checkpoint(dlm_repo) target_layer_ids = cfg.get("dflash_config", {}).get("target_layer_ids", []) - block_size = cfg.get("block_size", None) + block_size = get_dflash_block_size(cfg) return state_dict, target_layer_ids, block_size @@ -337,7 +351,7 @@ def compile_qwen3vl_vlm_qpcs( width: int | None = None, hf_token: str | None = None, ) -> tuple[str, str]: - """Build the qwen3-vl TLM — vision encoder + language decoder — and compile it for + """Build the Qwen3-VL or Qwen3.5 TLM vision encoder + language decoder for SPD. Returns ``(lang_qpc, vision_qpc)``. Same DFlashTLMTransform-attaches-random / inject-real-weights pattern as @@ -358,6 +372,16 @@ def compile_qwen3vl_vlm_qpcs( logger.info(f"[compile_qwen3vl_tlm] base={tlm_repo} dlm={dlm_repo} block_size={block_size}") config = AutoConfig.from_pretrained(tlm_repo, trust_remote_code=True, token=hf_token) + is_qwen35 = config.model_type == "qwen3_5" + if is_qwen35: + if prefill_seq_len != QWEN3_5_DFLASH_CHUNK_SIZE: + logger.warning( + f"Qwen3.5 DFlash uses prefill_seq_len={QWEN3_5_DFLASH_CHUNK_SIZE} for its fixed gated-delta chunk." + ) + prefill_seq_len = QWEN3_5_DFLASH_CHUNK_SIZE + if block_size != QWEN3_5_DFLASH_CHUNK_SIZE: + raise ValueError(f"Qwen3.5 DFlash currently requires block_size={QWEN3_5_DFLASH_CHUNK_SIZE}.") + logger.warning("Qwen3.5 DFlash uses accepted-prefix replay; target timing includes the extra commit pass.") tlm_qeff = QEFFAutoModelForImageTextToText.from_pretrained( tlm_repo, config=config, @@ -380,6 +404,9 @@ def compile_qwen3vl_vlm_qpcs( text_model.fc.weight.data.div_(s) text_model.hidden_norm.weight.data.copy_(state_dict["hidden_norm.weight"].to(torch.float32)) logger.info(f"[compile_qwen3vl_tlm] fc/hidden_norm injected (fc scale s={s:.6f})") + if is_qwen35: + # The fixed chunk and matrix solve change the graph, not its I/O names. + tlm_qeff.lang_model.hash_params["dflash_cache_protocol"] = "qwen35_chunk16_matrix_replay_v3" compile_kwargs = { "prefill_seq_len": prefill_seq_len, @@ -387,19 +414,23 @@ def compile_qwen3vl_vlm_qpcs( "num_cores": num_cores, "num_devices": num_devices, "mxfp6_matmul": True, - "mxint8_kv_cache": True, + # Keep hybrid verification/replay state in FP16 until MXINT8 is validated. + "mxint8_kv_cache": not is_qwen35, "mos": 1, "use_onnx_subfunctions": False, "split_model_io": True, "dflash_block_size": block_size, } + if is_qwen35: + # Verification and replay use the same 16-token shape as prefill. + compile_kwargs["prefill_only"] = True if height is not None: compile_kwargs["height"] = height if width is not None: compile_kwargs["width"] = width qpc = tlm_qeff.compile(**compile_kwargs) - lang_qpc = str(qpc["lang_qpc_path"]) + lang_qpc = str(qpc["lang_prefill_qpc_path"] if is_qwen35 else qpc["lang_qpc_path"]) vision_qpc = str(qpc["vision_qpc_path"]) logger.info(f"[compile_qwen3vl_tlm] lang_qpc={lang_qpc} vision_qpc={vision_qpc}") return lang_qpc, vision_qpc @@ -572,7 +603,7 @@ def compile_qwen3vl_vlm_dlm_qpc( num_devices: int, hf_token: str | None = None, ) -> str: - """Build + compile the DFlash DLM (draft) for the qwen3-vl TLM. + """Build + compile the DFlash DLM for a Qwen3-VL or Qwen3.5 TLM. Same VLM-aware lm_head/embed injection as ``compile_gemma_vlm_dlm_qpc``, kept as a separate function since qwen3-vl's config is not registered under diff --git a/tests/transformers/spd/test_dflash_spd_inference.py b/tests/transformers/spd/test_dflash_spd_inference.py index 9a0191cb87..395325cc5f 100644 --- a/tests/transformers/spd/test_dflash_spd_inference.py +++ b/tests/transformers/spd/test_dflash_spd_inference.py @@ -13,6 +13,9 @@ """ import os +import runpy +import sys +from types import SimpleNamespace import numpy as np import pytest @@ -22,7 +25,9 @@ from QEfficient.generation.cloud_infer import QAICInferenceSession from QEfficient.generation.dflash_generation import ( SpecDecodingMetrics, + _run_spd_core, run_spd_inference_gemma4, + run_spd_inference_qwen3_vl, run_spd_inference_single, ) @@ -42,6 +47,167 @@ MASK_TOKEN_ID = 2 +@pytest.mark.parametrize("ctx_len", [4096, 8192]) +def test_qwen35_compile_accepts_context_larger_than_draft_window(monkeypatch, ctx_len): + from examples.performance.dflash import utils + + class ReachedWeightLoading(Exception): + pass + + monkeypatch.setattr( + utils.AutoConfig, + "from_pretrained", + lambda repo, **kwargs: SimpleNamespace(model_type="qwen3_5", sliding_window=2048), + ) + monkeypatch.setattr(utils, "read_dlm_meta", lambda *args: ({}, [1], 16)) + + def stop_loading(*args, **kwargs): + raise ReachedWeightLoading + + monkeypatch.setattr(utils.QEFFAutoModelForImageTextToText, "from_pretrained", stop_loading) + with pytest.raises(ReachedWeightLoading): + utils.compile_qwen3vl_vlm_qpcs( + "Qwen/Qwen3.6-27B", + "z-lab/Qwen3.6-27B-DFlash", + ctx_len=ctx_len, + prefill_seq_len=16, + num_cores=8, + num_devices=4, + ) + + +@pytest.mark.parametrize("ctx_len", [4096, 8192]) +def test_qwen35_vision_cli_accepts_long_context(monkeypatch, ctx_len): + from examples.performance.dflash import utils + + class ReachedCompilation(Exception): + pass + + monkeypatch.setitem(sys.modules, "utils", utils) + monkeypatch.setattr(sys, "path", sys.path.copy()) + monkeypatch.setattr( + sys, + "argv", + ["basic_inference_vision.py", "--model_name", "Qwen/Qwen3.6-27B", "--image", "--ctx_len", str(ctx_len)], + ) + monkeypatch.setattr( + utils.AutoConfig, "from_pretrained", lambda *args, **kwargs: SimpleNamespace(model_type="qwen3_5") + ) + + def stop_compile(*args, **kwargs): + assert kwargs["ctx_len"] == ctx_len + raise ReachedCompilation + + monkeypatch.setattr(utils, "compile_qwen3vl_vlm_qpcs", stop_compile) + with pytest.raises(ReachedCompilation): + runpy.run_path(os.path.join(os.path.dirname(utils.__file__), "basic_inference_vision.py"), run_name="__main__") + + +@pytest.mark.parametrize("sequence_lengths", [[], [16, 64]]) +def test_dflash_decode_specialization_accepts_fixed_and_dynamic_shapes(sequence_lengths): + from examples.performance.dflash.utils import validate_spd_decode_specialization + + session = SimpleNamespace( + binding_index_map={"input_ids": 0}, + bindings=[SimpleNamespace(dims=[1, 64])], + allowed_shapes=[[(8, [1, length])] for length in sequence_lengths], + ) + validate_spd_decode_specialization(session, 64) + with pytest.raises(ValueError, match="seq_len=32"): + validate_spd_decode_specialization(session, 32) + + +@pytest.mark.parametrize("prompt_chunk_size", [8, 32, 128]) +def test_qwen3_5_runtime_rejects_unsupported_prefill_qpc(prompt_chunk_size): + with pytest.raises(ValueError, match="rebuild with --prefill_seq_len 16"): + run_spd_inference_qwen3_vl( + prompt_text="test", + tokenizer=None, + processor=None, + tlm_config=SimpleNamespace(model_type="qwen3_5"), + dlm_session=None, + tlm_session=SimpleNamespace(input_names=["cache_commit"]), + mask_token_id=MASK_TOKEN_ID, + vocab_size=VOCAB_SIZE, + prompt_chunk_size=prompt_chunk_size, + ) + + +@pytest.mark.parametrize("target_seq_len", [16, 64]) +def test_qwen3_5_runtime_pads_verification_and_commit(monkeypatch, target_seq_len): + import torch + + from QEfficient.generation import dflash_generation as generation + + calls = [] + buffers = [] + + def run(feeds): + calls.append({key: value.copy() for key, value in feeds.items()}) + return { + "logits": np.zeros((1, target_seq_len), dtype=np.int32), + "hidden_states": np.zeros((1, target_seq_len, 32), dtype=np.float32), + "image_idx_output": np.zeros((1, 1), dtype=np.int64), + } + + session = SimpleNamespace( + input_names=["cache_commit", "vision_embeds"], + binding_index_map={"vision_embeds": 0}, + bindings=[SimpleNamespace(type=0, dims=(1, 1, 32))], + aic_to_np_dtype_mapping={0: np.dtype(np.float32)}, + run=run, + set_buffers=buffers.append, + skip_buffers=lambda names: None, + ) + monkeypatch.setattr( + generation, + "build_inputs_qwen3_vl", + lambda *a, **k: {"input_ids": torch.ones((1, 5), dtype=torch.int64), "attention_mask": torch.ones((1, 5))}, + ) + monkeypatch.setattr( + generation, + "compute_position_ids_qwen3_vl", + lambda *a, **k: (torch.arange(5).view(1, 1, 5).repeat(4, 1, 1), torch.zeros((1, 1), dtype=torch.int64)), + ) + + def core(**kwargs): + kwargs["tlm_prefill"](0, target_seq_len) + kwargs["prepare_decode"]() + ids = np.arange(16, dtype=np.int64).reshape(1, 16) + positions = np.arange(5, 21, dtype=np.int64).reshape(1, 16) + logits, hidden = kwargs["tlm_decode"](ids, positions) + assert logits.shape == (1, 16) + assert hidden.shape == (1, 16, 32) + verify = calls[-1] + assert verify["input_ids"].shape == (1, target_seq_len) + np.testing.assert_array_equal(verify["position_ids"][0, :, :16], positions) + assert np.all(verify["position_ids"][..., 16:] == -1) + assert verify["cache_commit"].item() == 0 + kwargs["tlm_commit"](ids, positions, 2) + commit = calls[-1] + assert commit["cache_commit"].item() == 1 + np.testing.assert_array_equal(commit["position_ids"][0, :, :2], positions[:, :2]) + assert np.all(commit["position_ids"][..., 2:] == -1) + return kwargs["metrics"] + + monkeypatch.setattr(generation, "_run_spd_core", core) + generation.run_spd_inference_qwen3_vl( + prompt_text="test", + tokenizer=SimpleNamespace(pad_token_id=0, eos_token_id=None), + processor=None, + tlm_config=SimpleNamespace(model_type="qwen3_5"), + dlm_session=None, + tlm_session=session, + mask_token_id=2, + vocab_size=64, + prompt_chunk_size=target_seq_len, + ctx_len=256, + block_size=16, + hidden_size=32, + ) + assert all(value.shape[1] == target_seq_len for item in buffers for key, value in item.items() if key == "logits") + + class _FakeTokenizer: """Minimal stand-in implementing only what run_spd_inference_single calls: __call__(prompts, return_tensors="np", padding=...) and pad/eos token ids. @@ -157,6 +323,51 @@ def test_dflash_acceptance_rate_uses_accepted_tokens(): assert metrics.acceptance_rate() == 2.0 +@pytest.mark.parametrize("accepted", [0, 1, 3]) +def test_dflash_commit_counts_seed_and_preserves_verification_features(accepted): + commits = [] + hidden = np.ones((1, BLOCK_SIZE, HIDDEN_SIZE), dtype=np.float32) + + class Draft(_FakeDLMSession): + def run(self, inputs): + if commits: + np.testing.assert_array_equal(inputs["target_hidden"], 1) + return super().run(inputs) + + def verify(ids, positions): + hidden.fill(1) + logits = np.array([[6, 7, 8, 9]], dtype=np.int32) + if accepted < BLOCK_SIZE - 1: + logits[:, accepted] = 10 + return logits, hidden + + def commit(ids, positions, count): + commits.append(count) + hidden.fill(99) # Simulate session output-buffer reuse by the replay call. + + metrics = _run_spd_core( + input_ids=np.array([[3, 4, 0, 0]]), + flat_position_ids=np.array([[0, 1, -1, -1]]), + tlm_prefill=lambda start, end: (np.array([[0, 5, 0, 0]]), hidden), + prepare_decode=lambda: None, + tlm_decode=verify, + tlm_commit=commit, + dlm_session=Draft(), + mask_token_id=MASK_TOKEN_ID, + vocab_size=VOCAB_SIZE, + prompt_chunk_size=BLOCK_SIZE, + block_size=BLOCK_SIZE, + hidden_size=HIDDEN_SIZE, + generation_len=16, + max_iterations=2, + eos_token_ids=set(), + generated_ids=np.zeros((1, 16), dtype=np.int64), + metrics=SpecDecodingMetrics(block_size=BLOCK_SIZE), + ) + assert commits == [accepted + 1, accepted + 1] + assert metrics.tlm_decode_time == pytest.approx(metrics.tlm_verify_time + metrics.tlm_commit_time) + + def test_dflash_generation_len_caps_partial_block(): metrics = run_spd_inference_single( prompt_text="hi", diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index 8e80235c1d..086afd75ed 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -861,6 +861,133 @@ def _ort_qwen_logits(onnx_path: Path, qeff_model, inputs): return outputs["logits"].astype(np.float32) +def test_qwen3_5_without_dflash_hf_qeff_ort_parity(tmp_path): + _assert_qwen_hf_qeff_ort_parity("qwen3_5", tmp_path) + + +def test_dflash_mixed_window_ort_parity(tmp_path): + """One draft graph must retain per-layer masks as positions cross the local window.""" + from transformers import Qwen3Config, Qwen3ForCausalLM + + torch.manual_seed(27) + config = Qwen3Config( + hidden_size=32, + intermediate_size=64, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=16, + num_hidden_layers=2, + vocab_size=64, + max_position_embeddings=8192, + use_sliding_window=True, + sliding_window=2048, + layer_types=["sliding_attention", "full_attention"], + ) + qeff = QEFFAutoModelForCausalLM(Qwen3ForCausalLM(config).eval(), qaic_config={"dflash_dlm": True}) + onnx_path = _exported_onnx_path(qeff.export(tmp_path / "dflash-window", offload_pt_weights=False)) + session = _ort_session(onnx_path) + session_inputs = [item.name for item in session.get_inputs()] + session_outputs = [item.name for item in session.get_outputs()] + cache = tuple((torch.randn(1, 1, 8192, 16), torch.randn(1, 1, 8192, 16)) for _ in range(2)) + previous = {} + for start, valid in [(16, 16), (2032, 16), (2047, 15), (2048, 1), (2064, 16), (4096, 16)]: + target_positions = torch.arange(start - valid, start - valid + 16).unsqueeze(0) + target_positions[:, valid:] = -1 + inputs = { + "input_ids": torch.randint(0, 64, (1, 16)), + "target_hidden": torch.randn(1, 16, 32), + "position_ids": torch.arange(start, start + 16).unsqueeze(0), + "position_ids_target": target_positions, + "past_key_values": cache, + } + feed = { + name: value.copy() for name, value in _flatten_qwen_ort_inputs(inputs, session_inputs, qeff.model).items() + } + for name in session_inputs: + if name + "_RetainedState" in previous: + feed[name] = previous[name + "_RetainedState"] + with torch.no_grad(): + expected = qeff.model(**inputs) + actual = dict(zip(session_outputs, session.run(session_outputs, feed))) + np.testing.assert_allclose(actual["logits"], expected.logits.numpy(), atol=3e-4, rtol=3e-4) + cache = expected.past_key_values + for layer_idx, layer_cache in enumerate(cache): + for name, tensor in zip((f"past_key.{layer_idx}", f"past_value.{layer_idx}"), layer_cache): + np.testing.assert_allclose(actual[name + "_RetainedState"], tensor.numpy(), atol=3e-4, rtol=3e-4) + previous = actual + + +@pytest.mark.parametrize( + "prefill_length,verification_length,dflash_chunk_size", + [(8, 16, 64), (64, 16, 64), (8, 64, 64), (64, 64, 64), (8, 16, 16), (16, 16, 16)], +) +@pytest.mark.parametrize("matrix_solve", [False, True]) +def test_qwen3_5_dflash_ort_recurrent_commit( + tmp_path, prefill_length, verification_length, dflash_chunk_size, matrix_solve +): + from tests.unit_test.transforms.test_dflash_transform import make_tiny_qwen3_5_vlm + + torch.manual_seed(11) + model = make_tiny_qwen3_5_vlm() + model.config.text_config.max_position_embeddings = 256 + qeff = QEFFAutoModelForImageTextToText(model, kv_offload=True, qaic_config={"target_layer_ids": [1, 3]}) + qeff.lang_model.model.language_model.dflash_chunk_size = dflash_chunk_size + for layer in qeff.lang_model.model.language_model.layers: + if layer.layer_type == "linear_attention": + layer.linear_attn.dflash_matrix_solve = matrix_solve + inputs, output_names, dynamic_axes = _qwen_export_io(qeff, prefill_seq_len=8) + for i, kind in enumerate(qeff.model.config.text_config.layer_types): + if kind == "full_attention": + inputs["past_key_values"][i] = [ + tensor.new_zeros((tensor.shape[0], tensor.shape[1], 256, tensor.shape[3])) + for tensor in inputs["past_key_values"][i] + ] + qeff.lang_model.export( + inputs=deepcopy(inputs), + output_names=output_names, + dynamic_axes=dynamic_axes, + export_dir=tmp_path / "qwen35-dflash", + offload_pt_weights=False, + ) + session = _ort_session(Path(qeff.lang_model.onnx_path)) + session_inputs = [item.name for item in session.get_inputs()] + assert "cache_commit" in session_inputs + decoder = qeff.lang_model.model + cache = deepcopy(inputs["past_key_values"]) + batch_size = inputs["input_ids"].shape[0] + actual = {} + prompt_length = prefill_length - 3 + # Keep one exported graph across different sequence lengths and acceptance lengths. + for start, length, valid, commit in [ + (0, prefill_length, prompt_length, 1), + (prompt_length, verification_length, 16, 0), + (prompt_length, verification_length, 2, 1), + (prompt_length + 2, verification_length, 16, 0), + (prompt_length + 2, verification_length, 1, 1), + (prompt_length + 3, verification_length, 16, 0), + ]: + inputs["input_ids"] = torch.full((batch_size, length), 10 + start, dtype=torch.int64) + positions = torch.arange(start, start + length).view(1, 1, -1).repeat(4, batch_size, 1) + positions[:, :, valid:] = -1 + inputs["position_ids"] = positions + inputs["cache_commit"] = torch.tensor([commit]) + inputs["past_key_values"] = cache + feed = {name: value.copy() for name, value in _flatten_qwen_ort_inputs(inputs, session_inputs, decoder).items()} + if start: + for name in session_inputs: + if name + "_RetainedState" in actual: + feed[name] = actual[name + "_RetainedState"] + with torch.no_grad(): + predicted, _vision, _image_idx, hidden, cache = decoder(**inputs) + actual = dict(zip(output_names, session.run(output_names, feed))) + np.testing.assert_array_equal(actual["logits"][:, :valid], predicted.numpy()[:, :valid]) + np.testing.assert_allclose(actual["hidden_states"][:, :valid], hidden.numpy()[:, :valid], atol=3e-4, rtol=3e-4) + for layer_idx, layer_state in enumerate(cache): + if qeff.model.config.text_config.layer_types[layer_idx] == "linear_attention": + for name, expected in zip(decoder.get_onnx_past_key_value_names(layer_idx), layer_state): + np.testing.assert_allclose(actual[name + "_RetainedState"], expected.numpy(), atol=3e-4, rtol=3e-4) + + def _export_qwen_decoder_onnx(qeff_model, inputs, tmp_path, *, model_type: str, prefill_only: bool, layerwise: bool): export_dir = ( tmp_path / f"{model_type}-{'layerwise' if layerwise else 'default'}-{'prefill' if prefill_only else 'decode'}" @@ -3351,6 +3478,19 @@ def test_qwen3_5_moe_conv_decode_slice_keeps_prefill_gather_path(): torch.testing.assert_close(state, expected_state) +@pytest.mark.parametrize("valid_tokens", [0, 1, 2, 3, 4]) +def test_qwen3_5_conv_replay_preserves_history_order(valid_tokens): + from QEfficient.transformers.models.qwen3_5.modeling_qwen3_5 import qeff_torch_causal_conv1d_update + + state = torch.arange(24, dtype=torch.float32).reshape(1, 2, 3, 4) + hidden = 100 + torch.arange(6 * 64, dtype=torch.float32).reshape(1, 6, 64) + positions = torch.full((4, 1, 64), -1, dtype=torch.int64) + positions[..., :valid_tokens] = torch.arange(40, 40 + valid_tokens) + _, actual = qeff_torch_causal_conv1d_update(hidden, state, torch.ones(6, 4), positions) + expected = torch.cat([state.reshape(1, 6, 4), hidden[..., :valid_tokens]], dim=-1)[..., -4:] + torch.testing.assert_close(actual, expected.reshape_as(state), rtol=0, atol=0) + + def test_qwen3_5_conv_decode_slice_matches_gather_reference(): from QEfficient.transformers.models.qwen3_5.modeling_qwen3_5 import qeff_torch_causal_conv1d_update diff --git a/tests/unit_test/transforms/test_dflash_transform.py b/tests/unit_test/transforms/test_dflash_transform.py index 7d7299321d..30f9bfc8b9 100644 --- a/tests/unit_test/transforms/test_dflash_transform.py +++ b/tests/unit_test/transforms/test_dflash_transform.py @@ -23,6 +23,8 @@ All tests run on CPU only, using tiny in-memory models -- no checkpoints, no Hugging Face Hub access. """ +import copy + import pytest import torch from torch import nn @@ -32,6 +34,165 @@ HIDDEN_SIZE = 64 +def make_tiny_qwen3_5_vlm(): + from transformers import Qwen3_5Config, Qwen3_5ForConditionalGeneration + + config = Qwen3_5Config( + text_config={ + "vocab_size": VOCAB_SIZE, + "hidden_size": 32, + "intermediate_size": 64, + "num_hidden_layers": 4, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 16, + "layer_types": ["linear_attention", "full_attention", "linear_attention", "full_attention"], + "linear_key_head_dim": 8, + "linear_value_head_dim": 8, + "linear_num_key_heads": 2, + "linear_num_value_heads": 4, + "max_position_embeddings": CTX_LEN, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 10000, + "partial_rotary_factor": 1.0, + "mrope_section": [2, 3, 3], + }, + }, + vision_config={ + "depth": 1, + "hidden_size": 32, + "intermediate_size": 64, + "num_heads": 2, + "out_hidden_size": 32, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + }, + image_token_id=3, + video_token_id=4, + vision_start_token_id=5, + vision_end_token_id=6, + ) + config._attn_implementation = "eager" + return Qwen3_5ForConditionalGeneration(config).eval() + + +@pytest.mark.parametrize("accepted", [0, 1, 3]) +@pytest.mark.parametrize("prompt_len", [1, 5]) +@pytest.mark.parametrize("chunk_size", [16, 64]) +@pytest.mark.parametrize("matrix_solve", [False, True]) +def test_qwen3_5_dflash_replay_matches_autoregressive(accepted, prompt_len, chunk_size, matrix_solve, monkeypatch): + """Verify then replay a prefix; next-round results must exclude rejected tokens.""" + from QEfficient import QEFFAutoModelForImageTextToText + + torch.manual_seed(7) + hf_model = make_tiny_qwen3_5_vlm() + reference = copy.deepcopy(hf_model) + qeff = QEFFAutoModelForImageTextToText(hf_model, kv_offload=True, qaic_config={"target_layer_ids": [1, 3]}) + model = qeff.lang_model.model.language_model + model.dflash_chunk_size = chunk_size + for layer in model.layers: + if layer.layer_type == "linear_attention": + layer.linear_attn.dflash_matrix_solve = matrix_solve + + def unexpected_decode(*args, **kwargs): + raise AssertionError("DFlash must not compute the one-token recurrent branch") + + monkeypatch.setattr( + "QEfficient.transformers.models.qwen3_5.modeling_qwen3_5.recurrent_gdn_decode_forward", unexpected_decode + ) + initial = qeff.model.get_dummy_inputs(kv_offload=True)["lang"]["past_key_values"] + tokens = torch.randint(10, VOCAB_SIZE, (1, 16)) + # This standalone runtime uses a single request. + cache = tuple(tuple(t[:1].clone() for t in layer) for layer in initial) + + def run(ids, start, state, commit, valid=None): + positions = torch.arange(start, start + ids.shape[1]).view(1, -1) + if valid is not None: + positions[:, valid:] = -1 + return model( + input_ids=ids, + position_ids=positions, + past_key_values=state, + cache_commit=torch.tensor([commit]), + use_cache=True, + ) + + with torch.no_grad(): + prefill = run(tokens[:, :8], 0, cache, 1, valid=prompt_len) + verify = run(tokens[:, prompt_len : prompt_len + 4], prompt_len, prefill.past_key_values, 0) + for i in (0, 2): + for actual, expected in zip(verify.past_key_values[i], prefill.past_key_values[i]): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + commit_count = accepted + 1 # seed token plus accepted DLM tokens + committed = run( + tokens[:, prompt_len : prompt_len + 4], prompt_len, verify.past_key_values, 1, valid=commit_count + ) + next_pos = prompt_len + commit_count + followup_ids = torch.tensor([[48, 49, 50, 51]]) + followup = run(followup_ids, next_pos, committed.past_key_values, 0) + expected_ids = torch.cat([tokens[:, :next_pos], followup_ids], dim=1) + expected = reference.model.language_model(input_ids=expected_ids, use_cache=False, output_hidden_states=True) + + torch.testing.assert_close(followup.last_hidden_state, expected.last_hidden_state[:, -4:], rtol=2e-4, atol=2e-4) + projection = model.hidden_norm(model.fc(torch.cat([expected.hidden_states[i][:, -4:] for i in [1, 3]], dim=-1))) + torch.testing.assert_close(followup.hidden_states, projection, rtol=2e-4, atol=2e-4) + assert followup.hidden_states.shape == (1, 4, 32) + with torch.no_grad(): + restarted = run(tokens[:, :8], 0, followup.past_key_values, 1, valid=prompt_len) + for i in (0, 2): + for actual, expected in zip(restarted.past_key_values[i], prefill.past_key_values[i]): + torch.testing.assert_close(actual, expected, rtol=2e-4, atol=2e-4) + + +def test_qwen3_5_dflash_requires_explicit_commit(): + from QEfficient import QEFFAutoModelForImageTextToText + + qeff = QEFFAutoModelForImageTextToText( + make_tiny_qwen3_5_vlm(), kv_offload=True, qaic_config={"target_layer_ids": [1, 3]} + ) + inputs = qeff.model.get_dummy_inputs(kv_offload=True)["lang"] + inputs.pop("cache_commit") + with pytest.raises(ValueError, match="requires cache_commit"): + qeff.lang_model.model(**inputs) + + +@pytest.mark.parametrize("prefill_seq_len", [1, 65, 128]) +def test_qwen3_5_dflash_rejects_unsupported_prefill(prefill_seq_len): + from QEfficient import QEFFAutoModelForImageTextToText + + qeff = QEFFAutoModelForImageTextToText( + make_tiny_qwen3_5_vlm(), kv_offload=True, qaic_config={"target_layer_ids": [1, 3]} + ) + with pytest.raises(ValueError, match="prefill_seq_len=16"): + qeff.model.get_specializations(batch_size=1, prefill_seq_len=prefill_seq_len, ctx_len=256, kv_offload=True) + with pytest.raises(ValueError, match="prefill_seq_len <= 64"): + qeff.model.get_dummy_inputs(kv_offload=True, prefill_seq_len=prefill_seq_len) + + +@pytest.mark.parametrize("dummy_seq_len", [16, 32, 64]) +def test_qwen3_5_dflash_export_traces_one_chunk(dummy_seq_len): + from QEfficient import QEFFAutoModelForImageTextToText + + qeff = QEFFAutoModelForImageTextToText( + make_tiny_qwen3_5_vlm(), kv_offload=True, qaic_config={"target_layer_ids": [1, 3]} + ) + inputs = qeff.model.get_dummy_inputs(kv_offload=True, prefill_seq_len=dummy_seq_len)["lang"] + assert inputs["input_ids"].shape[1] == 16 + assert inputs["position_ids"].shape[-1] == 16 + with pytest.raises(ValueError, match="prefill_seq_len=16"): + qeff.model.get_specializations(batch_size=1, prefill_seq_len=64, ctx_len=256, kv_offload=True) + + +@pytest.mark.parametrize("chunk_size", [0, 1, 3, 24]) +def test_qwen3_5_dflash_matrix_solve_requires_power_of_two(chunk_size): + from QEfficient.transformers.models.qwen3_5.modeling_qwen3_5 import QEffQwen3_5GatedDeltaNet + + with pytest.raises(ValueError, match="power-of-two"): + QEffQwen3_5GatedDeltaNet._dflash_triangular_inverse(None, None, chunk_size) + + # --------------------------------------------------------------------------- # Tiny model factories # --------------------------------------------------------------------------- @@ -203,6 +364,69 @@ def test_dflash_draft_forward_smoke(self): assert outputs.logits.shape == (batch, block_size, VOCAB_SIZE) assert torch.isfinite(outputs.logits).all() + @pytest.mark.parametrize("is_causal", [None, False, True]) + def test_mixed_attention_masks_follow_draft_config(self, is_causal): + from QEfficient import QEFFAutoModelForCausalLM + from QEfficient.transformers.cache_utils import QEffDynamicCache + + hf_model = make_tiny_qwen3() + hf_model.config.layer_types = ["sliding_attention", "full_attention"] + hf_model.config.sliding_window = 8 + hf_model.config.is_causal = is_causal + qeff = QEFFAutoModelForCausalLM(hf_model, qaic_config={"dflash_dlm": True}) + assert qeff.hash_params["dflash_attention"] == "per_layer_window_v1" + masks = [] + + def record_mask(module, args, kwargs): + masks.append(kwargs["attention_mask"].clone()) + + handles = [ + layer.self_attn.register_forward_pre_hook(record_mask, with_kwargs=True) + for layer in qeff.model.model.layers + ] + try: + with torch.no_grad(): + qeff.model( + input_ids=torch.ones((1, 4), dtype=torch.int64), + target_hidden=torch.randn(1, 4, HIDDEN_SIZE), + position_ids_target=torch.arange(5, 9).unsqueeze(0), + position_ids=torch.arange(9, 13).unsqueeze(0), + past_key_values=QEffDynamicCache.from_legacy_cache( + _make_legacy_past_key_values(hf_model.config, 1, CTX_LEN) + ), + use_cache=True, + ) + finally: + for handle in handles: + handle.remove() + for i, mask in enumerate(masks): + causal = i == 0 if is_causal is None else is_causal + for row, query in enumerate(range(9, 13)): + for key in range(CTX_LEN): + visible = key <= 12 and (not causal or key <= query) + if i == 0: + visible = visible and query - key < 8 and (causal or key - query < 8) + assert mask[0, 0, row, key].item() == (not visible) + + +@pytest.mark.parametrize("start", [0, 2032, 2047, 2048, 2064]) +@pytest.mark.parametrize("sliding_window,is_causal", [(None, False), (2048, True), (2048, False)]) +def test_dflash_mask_crosses_sliding_window(start, sliding_window, is_causal): + from QEfficient.transformers.models.dflash_draft.modeling_dflash_draft import _create_mask + + positions = torch.arange(start, start + 16).unsqueeze(0) + mask = _create_mask(positions, 4096, sliding_window=sliding_window, is_causal=is_causal) + keys = torch.arange(4096) + for row, query in enumerate(positions[0].tolist()): + visible = keys <= start + 15 + if is_causal: + visible &= keys <= query + if sliding_window: + visible &= query - keys < sliding_window + if not is_causal: + visible &= keys - query < sliding_window + torch.testing.assert_close(mask[0, 0, row], ~visible) + # --------------------------------------------------------------------------- # Tests: target (TLM) side -- DFlashTLMTransform From d394ef80e50d6b15e47457923dd5e009b7d850ab Mon Sep 17 00:00:00 2001 From: Vahid Janfaza Date: Mon, 12 Oct 2026 12:01:09 -0700 Subject: [PATCH 2/9] test: cover invalid DFlash block sizes Cover missing block_size, empty dflash_config, non-integer values, and values below two in both supported config locations. Co-authored-by: FNU Annanya Signed-off-by: Vahid Janfaza --- .../spd/test_dflash_spd_inference.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/transformers/spd/test_dflash_spd_inference.py b/tests/transformers/spd/test_dflash_spd_inference.py index 395325cc5f..c757c45155 100644 --- a/tests/transformers/spd/test_dflash_spd_inference.py +++ b/tests/transformers/spd/test_dflash_spd_inference.py @@ -47,6 +47,26 @@ MASK_TOKEN_ID = 2 +@pytest.mark.parametrize("config", [{}, {"dflash_config": {}}], ids=["missing-block-size", "empty-dflash-config"]) +def test_get_dflash_block_size_rejects_missing_value(config): + from examples.performance.dflash import utils + + with pytest.raises(ValueError, match="block_size >= 2"): + utils.get_dflash_block_size(config) + + +@pytest.mark.parametrize("block_size", [None, "16", 16.0, -1, 0, 1]) +@pytest.mark.parametrize("nested", [False, True], ids=["top-level", "nested"]) +def test_get_dflash_block_size_rejects_invalid_value(block_size, nested): + from examples.performance.dflash import utils + + config = {"block_size": block_size} + if nested: + config = {"dflash_config": config} + with pytest.raises(ValueError, match="block_size >= 2"): + utils.get_dflash_block_size(config) + + @pytest.mark.parametrize("ctx_len", [4096, 8192]) def test_qwen35_compile_accepts_context_larger_than_draft_window(monkeypatch, ctx_len): from examples.performance.dflash import utils From 4b1fba2b26a289195163f1135cc8c6891222fc08 Mon Sep 17 00:00:00 2001 From: Vahid Janfaza Date: Mon, 12 Oct 2026 15:35:36 -0700 Subject: [PATCH 3/9] Address DFlash ValueError test cases and add Qwen3.5-MoE support Extend the existing Qwen3.5 DFlash approach to Qwen3.5-MoE, reusing the same hybrid-cache handling and accepted-prefix replay. Keep the existing MoE routing and expert-parallel implementation unchanged. Co-authored-by: FNU Annanya Signed-off-by: Vahid Janfaza --- QEfficient/generation/dflash_generation.py | 35 ++- .../transformers/models/modeling_auto.py | 23 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 202 +++++++++++++++--- QEfficient/transformers/spd/dflash.py | 33 +++ .../dflash/basic_inference_vision.py | 4 +- examples/performance/dflash/utils.py | 19 +- .../spd/test_dflash_spd_inference.py | 96 ++++++++- .../unit_test/models/test_model_quickcheck.py | 98 ++++++++- .../transforms/test_dflash_transform.py | 108 +++++++++- 9 files changed, 552 insertions(+), 66 deletions(-) diff --git a/QEfficient/generation/dflash_generation.py b/QEfficient/generation/dflash_generation.py index c554403f4e..77e04a0977 100644 --- a/QEfficient/generation/dflash_generation.py +++ b/QEfficient/generation/dflash_generation.py @@ -584,6 +584,11 @@ def _get_rope_model_qwen3_vl(config): if config.model_type == "qwen3_5" else QEffQwen3VLForConditionalGeneration ) + if config.model_type == "qwen3_5_moe": + # Native HF construction initializes the MoE expert backend on meta. + from transformers import Qwen3_5MoeForConditionalGeneration + + model_cls = Qwen3_5MoeForConditionalGeneration _rope_model_cache_qwen3_vl[key] = model_cls._from_config(config) return _rope_model_cache_qwen3_vl[key] @@ -596,7 +601,7 @@ def compute_position_ids_qwen3_vl(input_ids, attention_mask, image_grid_thw, con """ rope_model = _get_rope_model_qwen3_vl(config) batch_size, seq_len = input_ids.shape - if config.model_type == "qwen3_5": + if config.model_type in {"qwen3_5", "qwen3_5_moe"}: if mm_token_type_ids is None: mm_token_type_ids = (input_ids == config.image_token_id).to(torch.int64) mrope, deltas = rope_model.model.get_rope_index( @@ -662,6 +667,20 @@ def run_vision_encoder_qwen3_vl(vision_session, pixel_values, image_grid_thw): return vision_outputs["vision_embeds"], vision_outputs["deepstack_features"] +def get_qwen35_target_seq_len(tlm_session: QAICInferenceSession, block_size: int) -> int: + """Require block-sized verification, independently of prompt prefill.""" + index = tlm_session.binding_index_map["input_ids"] + lengths = {shape[index][1][1] for shape in tlm_session.allowed_shapes} + if not lengths: + lengths = {tlm_session.bindings[index].dims[1]} + if block_size in lengths: + return block_size + raise ValueError( + f"Qwen3.5-MoE DFlash requires a {block_size}-token verification specialization " + f"(found {sorted(lengths)}). Recompile the TLM with dflash_block_size={block_size}." + ) + + def run_spd_inference_qwen3_vl( prompt_text: str, tokenizer, @@ -683,17 +702,22 @@ def run_spd_inference_qwen3_vl( compiled_width: int = 536, batch_size: int = 1, ) -> SpecDecodingMetrics: - """Run Qwen3-VL or Qwen3.5 DFlash for a text-only or single-image prompt.""" - is_qwen35 = tlm_config.model_type == "qwen3_5" + """Run Qwen3-VL, Qwen3.5 or Qwen3.5-MoE DFlash for a text or image prompt.""" + is_qwen35 = tlm_config.model_type in {"qwen3_5", "qwen3_5_moe"} if is_qwen35 and (batch_size != 1 or "cache_commit" not in tlm_session.input_names): raise ValueError("Qwen3.5 DFlash requires batch size 1 and a TLM QPC with cache_commit input.") - if is_qwen35 and ( + if tlm_config.model_type == "qwen3_5" and ( prompt_chunk_size not in (QWEN3_5_DFLASH_CHUNK_SIZE, 64) or not 1 < block_size <= prompt_chunk_size ): raise ValueError( "Qwen3.5 DFlash requires 16-token target calls (or a legacy 64-token QPC) and a block that fits; " "rebuild with --prefill_seq_len 16." ) + target_seq_len = prompt_chunk_size if is_qwen35 else block_size + if tlm_config.model_type == "qwen3_5_moe": + if block_size < 2 or prompt_chunk_size < block_size or prompt_chunk_size % block_size: + raise ValueError("Qwen3.5-MoE DFlash requires prefill length to be a positive multiple of block_size >= 2.") + target_seq_len = get_qwen35_target_seq_len(tlm_session, block_size) metrics = SpecDecodingMetrics(block_size=block_size) image_processing_start = time.time() @@ -777,7 +801,6 @@ def tlm_prefill(chunk_start, chunk_end): return outputs["logits"], outputs["hidden_states"] def prepare_decode(): - target_seq_len = prompt_chunk_size if is_qwen35 else block_size tlm_session.set_buffers({"logits": np.zeros((batch_size, target_seq_len), dtype=np.int32)}) tlm_session.set_buffers( {"hidden_states": np.zeros((batch_size, target_seq_len, hidden_size), dtype=np.float32)} @@ -789,7 +812,7 @@ def prepare_decode(): def decode_inputs(draft_ids, decode_position_ids, valid_tokens): positions = build_decode_position_ids_qwen3_vl(decode_position_ids, rope_delta) if is_qwen35: - padding = prompt_chunk_size - block_size + padding = target_seq_len - block_size draft_ids = np.pad(draft_ids, ((0, 0), (0, padding)), constant_values=tokenizer.pad_token_id) positions = np.pad(positions, ((0, 0), (0, 0), (0, padding)), constant_values=-1) positions[:, :, valid_tokens:] = -1 diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 8386c0866e..00185e5bdf 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -1594,6 +1594,10 @@ def __init__( self.model, _ = DFlashTLMTransform.apply(self.model, qaic_config) # ---Sampling--- + if self.dflash_tlm and self.model.config.model_type == "qwen3_5_moe": + self.lang_model.hash_params["dflash_cache_protocol"] = "qwen35_moe_dynamic_replay_v1" + self.lang_model.hash_params["dflash_chunk_size"] = 16 + # Note: SamplerTransform should be applied after all other transforms # are done. The role of the sampler is to just add nodes at the output of the # previous transform function. @@ -2093,6 +2097,17 @@ def compile( full_batch_size=full_batch_size, ) + if self.dflash_tlm and self.model.config.model_type == "qwen3_5_moe": + language_model = self.lang_model.model.language_model + previous_chunk_size = getattr(language_model, "dflash_chunk_size", 16) + chunk_size = previous_chunk_size if dflash_block_size is None else dflash_block_size + if not isinstance(chunk_size, int) or chunk_size < 2 or chunk_size & (chunk_size - 1): + raise ValueError("Qwen3.5-MoE DFlash requires a power-of-two dflash_block_size >= 2.") + if chunk_size != previous_chunk_size and self.lang_model.onnx_path is not None and lang_onnx_path is None: + raise ValueError("Changing dflash_block_size requires a fresh model instance to re-export the TLM.") + language_model.dflash_chunk_size = chunk_size + self.lang_model.hash_params["dflash_chunk_size"] = chunk_size + if layerwise: if skip_lang and not skip_vision: vision_wrapper = self._build_layerwise_vision_wrapper() @@ -2333,7 +2348,13 @@ def compile( # but one CCL value is silently dropped. lang_specs = specializations["lang"] if prefill_only: - if self.comp_ctx_lengths_prefill is not None: + if self.dflash_tlm and self.model.config.model_type == "qwen3_5_moe": + # Prefill and block-sized verification share the chunk graph. + specializations = [] + for spec in lang_specs: + if spec not in specializations: + specializations.append(spec) + elif self.comp_ctx_lengths_prefill is not None: specializations = lang_specs[: len(self.comp_ctx_lengths_prefill)] else: specializations = lang_specs[:1] diff --git a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index e602a7136d..65f39a674b 100644 --- a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -75,6 +75,7 @@ delete_module_attrs, silu_glu_mlp, ) +from QEfficient.transformers.spd.dflash import compute_dflash_target_hidden_states, dflash_chunk_gated_delta_scan from QEfficient.utils import constants from QEfficient.utils._utils import IOInfo, get_padding_shape_from_config from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE @@ -82,6 +83,7 @@ logger = QEFFLogger.get_logger("MODEL") QWEN3_5_MOE_ROPE_CACHE_EXPORT_CAP = 76800 +QWEN3_5_MOE_DFLASH_CHUNK_SIZE = 16 def _expand_mrope_position_ids(position_ids, cache_position, batch_size): @@ -504,6 +506,7 @@ def qeff_torch_causal_conv1d_update( weight: torch.Tensor, position_ids: torch.Tensor, bias: Optional[torch.Tensor] = None, + dflash: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: _, hidden_size, seq_len = hidden_states.shape grouped_conv_state = conv_state.ndim == 4 @@ -532,9 +535,16 @@ def qeff_torch_causal_conv1d_update( shifted_conv_state = shifted_conv_state * valid_decode + conv_state_flat * (1 - valid_decode) # Prefill and padded chunks retain the general last-valid-position gather. - # Derive zeros from an existing state view; torch.zeros emits a dynamic ConstantOfShape in subfunctions. - zeros = conv_state_flat[:, 0, :].to(pos_ids.dtype) * 0 - order = torch.argsort(torch.cat([zeros, pos_ids], dim=1), dim=1) + if dflash: + # Distinct chronological keys preserve cached history during padded prefix replay. + cached_order = torch.arange(state_len, dtype=pos_ids.dtype, device=pos_ids.device) + cached_order = cached_order.unsqueeze(0).expand(pos_ids.shape[0], -1) + token_order = torch.where(pos_ids >= 0, pos_ids + state_len, -1) + order = torch.argsort(torch.cat([cached_order, token_order], dim=1), dim=1) + else: + # Derive zeros from an existing state view; torch.zeros emits a dynamic ConstantOfShape in subfunctions. + zeros = conv_state_flat[:, 0, :].to(pos_ids.dtype) * 0 + order = torch.argsort(torch.cat([zeros, pos_ids], dim=1), dim=1) last_positions = order[:, -state_len:] ctx_idx = last_positions.to(torch.long).unsqueeze(1).expand(-1, hidden_size, -1) gathered_conv_state = torch.gather(hidden_states_new, dim=2, index=ctx_idx) @@ -823,6 +833,19 @@ def _solve_chunk_attn(self, attn: torch.Tensor, mask: torch.Tensor, eye: torch.T return self._solve_chunk_attn_scaled_newton_schulz(attn, mask, eye, chunk_size) return self._solve_chunk_attn_recursive_sns(attn, mask, eye, chunk_size) + @staticmethod + def _dflash_triangular_inverse(attn, eye, chunk_size): + if chunk_size < 2 or chunk_size & (chunk_size - 1): + raise ValueError("DFlash matrix solve requires a power-of-two chunk size >= 2.") + # For strictly lower triangular A, A**C = 0. Doubling evaluates + # I + A + ... + A**(C-1) exactly for power-of-two chunk lengths. + result = attn + eye + power = attn + for _ in range(chunk_size.bit_length() - 2): + power = power @ power + result = result + result @ power + return result + def torch_chunk_gated_delta_rule_qeff( self, query, @@ -839,6 +862,8 @@ def torch_chunk_gated_delta_rule_qeff( mask_strict=None, ones_lower=None, eye=None, + use_matrix_solve=False, + use_dynamic_chunk_loop=False, ): initial_dtype = query.dtype # if use_qk_l2norm_in_kernel: @@ -914,7 +939,10 @@ def torch_chunk_gated_delta_rule_qeff( attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) I_base = eye if eye is not None else torch.eye(chunk_size, dtype=attn.dtype) - attn = self._solve_chunk_attn(attn=attn, mask=mask, eye=I_base, chunk_size=chunk_size) + if use_matrix_solve: + attn = self._dflash_triangular_inverse(attn, I_base.to(dtype=attn.dtype, device=attn.device), chunk_size) + else: + attn = self._solve_chunk_attn(attn=attn, mask=mask, eye=I_base, chunk_size=chunk_size) value = attn @ v_beta k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) @@ -924,21 +952,25 @@ def torch_chunk_gated_delta_rule_qeff( if initial_state is None else initial_state.to(value) ) - core_attn_out = torch.zeros_like(value) - mask = mask_strict - - # for each chunk - for i in range(total_sequence_length // chunk_size): - q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] - attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) - v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state - v_new = v_i - v_prime - attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - core_attn_out[:, :, i] = attn_inter + attn @ v_new - last_recurrent_state = ( - last_recurrent_state * g[:, :, i, -1, None, None].exp() - + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new + if use_dynamic_chunk_loop: + core_attn_out, last_recurrent_state = dflash_chunk_gated_delta_scan( + query, key, value, k_cumdecay, g, decay_mask, mask_strict, last_recurrent_state ) + else: + core_attn_out = torch.zeros_like(value) + mask = mask_strict + + for i in range(total_sequence_length // chunk_size): + q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) + v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state + v_new = v_i - v_prime + attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state + core_attn_out[:, :, i] = attn_inter + attn @ v_new + last_recurrent_state = ( + last_recurrent_state * g[:, :, i, -1, None, None].exp() + + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new + ) if not output_final_state: last_recurrent_state = None @@ -976,6 +1008,7 @@ def forward( batch_index: Optional[torch.LongTensor] = None, batch_fold: bool = False, gdn_num_head_blocks: int = 1, + dflash_chunk_size: int | None = None, ): batch_size, seq_len, _ = hidden_states.shape @@ -1045,6 +1078,7 @@ def forward( self.conv1d.weight.squeeze(1), position_ids, self.conv1d.bias, + dflash=dflash_chunk_size is not None, ) if state_batch_index is not None: if conv_state_all.ndim == 4: @@ -1098,7 +1132,30 @@ def forward( # ── Recurrent State ─────────────────────────────────── if cache_params is not None: - if batch_fold: + if dflash_chunk_size is not None: + # Select chunk execution at export time, including prefix replay. + # The scripted scan also supports larger prompt specializations. + if self.gdn_chunk_size < dflash_chunk_size: + raise ValueError(f"Qwen3.5-MoE DFlash requires gdn_chunk_size >= {dflash_chunk_size}.") + core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + position_ids=position_ids, + chunk_size=dflash_chunk_size, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + mask_causal=self._mask_causal[:dflash_chunk_size, :dflash_chunk_size], + mask_strict=self._mask_strict[:dflash_chunk_size, :dflash_chunk_size], + ones_lower=self._ones_lower[:dflash_chunk_size, :dflash_chunk_size], + eye=self._eye[:dflash_chunk_size, :dflash_chunk_size], + use_matrix_solve=getattr(self, "dflash_matrix_solve", True), + use_dynamic_chunk_loop=True, + ) + elif batch_fold: # Folded decode is exported with a static one-token sequence. core_attn_out, last_recurrent_state = self._recurrent_step_batched( query, @@ -1209,6 +1266,7 @@ def forward( gdn_num_head_blocks: int = 1, use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + dflash_chunk_size: int | None = None, **kwargs, ) -> torch.FloatTensor: del use_cache @@ -1225,6 +1283,7 @@ def forward( batch_index=batch_index, batch_fold=batch_fold, gdn_num_head_blocks=gdn_num_head_blocks, + dflash_chunk_size=dflash_chunk_size, ) else: hidden_states, _ = self.self_attn( @@ -1279,6 +1338,7 @@ def forward( use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, output_hidden_states: Optional[bool] = None, + cache_commit: torch.Tensor | None = None, **kwargs, ) -> BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): @@ -1297,12 +1357,32 @@ def forward( elif use_cache and past_key_values is None: past_key_values = QEffQwen3_5MoeDynamicCache(self.config) + target_layer_ids = getattr(self, "target_layer_ids", None) + if target_layer_ids: + if is_layerwise_active() or kwargs.get("layer_indices_to_run") is not None: + raise ValueError( + "Qwen3.5-MoE DFlash requires the full target model; layerwise execution is unsupported." + ) + if cache_commit is None: + raise ValueError( + "Qwen3.5-MoE DFlash requires cache_commit: 0 for verification, 1 for accepted-prefix replay." + ) + if batch_index is not None or batch_fold: + raise ValueError("Qwen3.5-MoE DFlash currently supports non-continuous batch size 1 only.") + if past_key_values is None or any( + past_key_values.conv_states[i] is None or past_key_values.recurrent_states[i] is None + for i, kind in enumerate(self.config.layer_types) + if kind == "linear_attention" + ): + raise ValueError("Qwen3.5-MoE DFlash requires preallocated convolution and recurrent states.") + if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) start, end = resolve_layer_window(QEffQwen3_5MoeTextModel, len(self.layers)) - if cache_position is None: + if cache_position is None or target_layer_ids: past_seen_tokens = past_key_values.get_seq_length(layer_idx=start) if past_key_values is not None else 0 + if cache_position is None: active_layer_type = self.config.layer_types[start] if start < len(self.config.layer_types) else None if is_layerwise_active() and position_ids is not None and active_layer_type == "linear_attention": text_position_ids = position_ids[0] if position_ids.ndim == 3 else position_ids @@ -1313,6 +1393,18 @@ def forward( ) position_ids = _expand_mrope_position_ids(position_ids, cache_position, inputs_embeds.shape[0]) text_position_ids = position_ids[0] if position_ids.ndim == 3 else position_ids + if target_layer_ids: + # Reset the saved state too: verification must not restore an old request. + reset = (text_position_ids[:, :1] == 0).reshape(-1, 1, 1, 1) + for i, kind in enumerate(self.config.layer_types): + if kind == "linear_attention": + conv = past_key_values.conv_states[i] + recurrent = past_key_values.recurrent_states[i] + conv_reset = reset if conv.ndim == 4 else reset.squeeze(-1) + past_key_values.conv_states[i] = torch.where(conv_reset, torch.zeros_like(conv), conv) + past_key_values.recurrent_states[i] = torch.where(reset, torch.zeros_like(recurrent), recurrent) + committed_conv = list(past_key_values.conv_states) + committed_recurrent = list(past_key_values.recurrent_states) rotary_position_ids = ( position_ids[1:] if position_ids.ndim == 3 and position_ids.shape[0] == 4 else position_ids ) @@ -1345,6 +1437,7 @@ def forward( ) position_embeddings = (cos, sin) all_hidden_states = () if output_hidden_states else None + target_hidden_list = [] layer_indices_to_run = kwargs.get("layer_indices_to_run", None) for layer_idx, decoder_layer in enumerate(self.layers): @@ -1354,6 +1447,8 @@ def forward( continue if output_hidden_states: all_hidden_states += (hidden_states,) + if target_layer_ids and layer_idx in target_layer_ids: + target_hidden_list.append(hidden_states) layer_mask = linear_attn_mask if decoder_layer.layer_type == "linear_attention" else causal_mask hidden_states = decoder_layer( @@ -1368,16 +1463,30 @@ def forward( gdn_num_head_blocks=gdn_num_head_blocks, use_cache=use_cache, cache_position=cache_position, + dflash_chunk_size=getattr(self, "dflash_chunk_size", QWEN3_5_MOE_DFLASH_CHUNK_SIZE) + if target_layer_ids + else None, **kwargs, ) - # break - + if target_layer_ids and self.config.num_hidden_layers in target_layer_ids: + target_hidden_list.append(hidden_states) if is_last_layer_window(QEffQwen3_5MoeTextModel, len(self.layers)): hidden_states = self.norm(hidden_states) if output_hidden_states: all_hidden_states += (hidden_states,) + if target_layer_ids: + # Attention KV uses positional overwrite; recurrent states require replay. + for i, kind in enumerate(self.config.layer_types): + if kind == "linear_attention": + past_key_values.conv_states[i] = torch.where( + cache_commit.bool(), past_key_values.conv_states[i], committed_conv[i] + ) + past_key_values.recurrent_states[i] = torch.where( + cache_commit.bool(), past_key_values.recurrent_states[i], committed_recurrent[i] + ) + if return_legacy_cache: past_key_values = past_key_values.to_legacy_cache() @@ -1386,6 +1495,9 @@ def forward( return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None, + hidden_states=compute_dflash_target_hidden_states(target_hidden_list, self.fc, self.hidden_norm) + if target_layer_ids + else None, ) @@ -1503,7 +1615,9 @@ def forward( **kwargs, ) - if position_ids is None: + if getattr(self.model, "target_layer_ids", None): + hidden_states = outputs.last_hidden_state + elif position_ids is None: hidden_states = outputs.last_hidden_state[:, -1:, :] else: text_position_ids = position_ids[0] if position_ids.ndim == 3 else position_ids @@ -1846,7 +1960,10 @@ def forward( past_key_values=None, batch_index: Optional[torch.LongTensor] = None, comp_ctx_lengths: Optional[List[int]] = None, + cache_commit: torch.Tensor | None = None, ): + if getattr(self.language_model, "target_layer_ids", None) and is_layerwise_active(): + raise ValueError("Qwen3.5-MoE DFlash requires the full target model; layerwise execution is unsupported.") # Continuous batching also supplies batch_index for prefill. Folded # attention is decode-only, so do not infer it from batch_index alone. seq_len = input_ids.shape[1] if input_ids is not None else position_ids.shape[-1] @@ -1898,7 +2015,18 @@ def forward( batch_fold=batch_fold_cb, gdn_num_head_blocks=gdn_num_head_blocks, use_cache=True, + cache_commit=cache_commit, ) + if getattr(self.language_model, "target_layer_ids", None): + logits = self.model.lm_head(outputs.last_hidden_state).float() + image_idx = (indices1.max() + 1).unsqueeze(0).unsqueeze(0) + return ( + logits.argmax(dim=-1).to(torch.int32), + vision_embeds, + image_idx, + outputs.hidden_states, + outputs.past_key_values, + ) 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: @@ -2141,6 +2269,12 @@ def get_specializations( num_frames = [num_frames] * len(height) if isinstance(num_frames, int) else num_frames prefill_seq_len = prefill_seq_len if prefill_seq_len else 128 + if getattr(language_model, "target_layer_ids", None): + chunk_size = getattr(language_model, "dflash_chunk_size", QWEN3_5_MOE_DFLASH_CHUNK_SIZE) + if prefill_seq_len < chunk_size or prefill_seq_len % chunk_size: + raise ValueError( + f"Qwen3.5-MoE DFlash requires prefill_seq_len to be a positive multiple of {chunk_size}." + ) ctx_len = ctx_len if ctx_len else constants.INTERN_CTX_LEN kv_cache_batch_size = kv_cache_batch_size or full_batch_size or batch_size channel = 3 @@ -2225,15 +2359,16 @@ def _build_lang_spec(seq_len_val, comp_ctx_len=None): spec["full_batch_exec_size"] = full_batch_size return spec + decode_seq_len = chunk_size if getattr(language_model, "target_layer_ids", None) else 1 lang = [] if comp_ctx_lengths_prefill is not None or comp_ctx_lengths_decode is not None: for comp_ctx in comp_ctx_lengths_prefill or []: lang.append(_build_lang_spec(prefill_seq_len, comp_ctx_len=comp_ctx)) for comp_ctx in comp_ctx_lengths_decode or []: - lang.append(_build_lang_spec(1, comp_ctx_len=comp_ctx)) + lang.append(_build_lang_spec(decode_seq_len, comp_ctx_len=comp_ctx)) else: lang.append(_build_lang_spec(prefill_seq_len)) - lang.append(_build_lang_spec(1)) + lang.append(_build_lang_spec(decode_seq_len)) if kv_offload: return {"vision": vision, "lang": lang}, compiler_options @@ -2267,6 +2402,10 @@ def get_onnx_dynamic_axes( "vision_embeds": {0: "vision_batch_size", 1: "vision_size"}, } + language_model = getattr(getattr(self, "model", None), "language_model", None) + if getattr(language_model, "target_layer_ids", None): + lang_dynamic_axes["hidden_states"] = {0: input_batch_axis, 1: "seq_len"} + for i in range(num_layers): if self.config.text_config.layer_types[i] == "full_attention": lang_dynamic_axes[f"past_key.{i}"] = {0: batch_axis_name, 2: "ctx_len"} @@ -2302,6 +2441,15 @@ def get_dummy_inputs( inputs_shapes = {} dummy_seq_len = int(kwargs.get("prefill_seq_len", 32)) + is_dflash = bool(getattr(self.model.language_model, "target_layer_ids", None)) + if is_dflash: + if not kv_offload or continuous_batching: + raise ValueError("Qwen3.5-MoE DFlash requires kv_offload=True and continuous_batching=False.") + if dummy_seq_len < 2: + raise ValueError("Qwen3.5-MoE DFlash requires prefill_seq_len >= 2 for the exported gated-delta graph.") + # A small dummy still exports the scripted loop's dynamic trip count. + chunk_size = getattr(self.model.language_model, "dflash_chunk_size", QWEN3_5_MOE_DFLASH_CHUNK_SIZE) + dummy_seq_len = min(dummy_seq_len, chunk_size) bs = kwargs.get("batch_size", constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE) if bs > 1: bs = 2 @@ -2345,6 +2493,8 @@ def get_dummy_inputs( .repeat(4, 1, 1) ) lang_inputs["image_idx"] = torch.zeros((inputs_shapes["image_idx"]), dtype=torch.int64) + if is_dflash: + lang_inputs["cache_commit"] = torch.ones(1, dtype=torch.int64) kv_cache_shape = get_padding_shape_from_config( config=self.model.config.text_config, @@ -2403,6 +2553,8 @@ def get_dummy_inputs( def get_output_names(self, kv_offload: bool = False): vision_output_names = ["vision_embeds"] lang_output_names = ["logits"] + if getattr(self.model.language_model, "target_layer_ids", None): + lang_output_names.append("hidden_states") for layer_idx, layer_type in enumerate(self.model.config.text_config.layer_types): if layer_idx < self.model.config.text_config.num_hidden_layers: if layer_type == "full_attention": diff --git a/QEfficient/transformers/spd/dflash.py b/QEfficient/transformers/spd/dflash.py index db05f6c8e2..617735d02d 100644 --- a/QEfficient/transformers/spd/dflash.py +++ b/QEfficient/transformers/spd/dflash.py @@ -8,6 +8,39 @@ import torch +@torch.jit.script +def dflash_chunk_gated_delta_scan( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + k_cumdecay: torch.Tensor, + g: torch.Tensor, + decay_mask: torch.Tensor, + mask_strict: torch.Tensor, + state: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Preserve the recurrent chunk loop in ONNX for sequence specialization.""" + output = torch.zeros_like(value) + batch, heads, chunk = query.shape[0], query.shape[1], query.shape[3] + for i in range(query.shape[2]): + q_i = query[:, :, i : i + 1].reshape(batch, heads, chunk, query.shape[4]) + k_i = key[:, :, i : i + 1].reshape(batch, heads, chunk, key.shape[4]) + v_i = value[:, :, i : i + 1].reshape(batch, heads, chunk, value.shape[4]) + k_decay_i = k_cumdecay[:, :, i : i + 1].reshape(batch, heads, chunk, key.shape[4]) + decay_i = decay_mask[:, :, i : i + 1].reshape(batch, heads, chunk, chunk) + g_i = g[:, :, i : i + 1].reshape(batch, heads, chunk) + last_g = g_i[:, :, -1:] + attn = (q_i @ k_i.transpose(-1, -2) * decay_i).masked_fill(mask_strict, 0) + v_prime = k_decay_i @ state + v_new = v_i - v_prime + attn_inter = (q_i * g_i.exp().unsqueeze(-1)) @ state + output[:, :, i : i + 1] = (attn_inter + attn @ v_new).unsqueeze(2) + state = ( + state * last_g.exp().unsqueeze(-1) + (k_i * (last_g - g_i).exp().unsqueeze(-1)).transpose(-1, -2) @ v_new + ) + return output, state + + def compute_dflash_target_hidden_states( target_hidden_list: list[torch.Tensor], fc: torch.nn.Module, diff --git a/examples/performance/dflash/basic_inference_vision.py b/examples/performance/dflash/basic_inference_vision.py index a58b943e2e..49ce8dfea6 100644 --- a/examples/performance/dflash/basic_inference_vision.py +++ b/examples/performance/dflash/basic_inference_vision.py @@ -249,7 +249,7 @@ def main(): raise ValueError(f"No default TLM HF path for '{args.model_name}'. Pass --tlm_hf_path.") tlm_config = AutoConfig.from_pretrained(tlm_repo, trust_remote_code=True, token=args.hf_token) - is_qwen3vl = tlm_config.model_type in {"qwen3_vl", "qwen3_5"} + is_qwen3vl = tlm_config.model_type in {"qwen3_vl", "qwen3_5", "qwen3_5_moe"} if args.tlm_qpc and args.vision_qpc: tlm_qpc, vision_qpc = args.tlm_qpc, args.vision_qpc @@ -338,7 +338,7 @@ def main(): print("=" * width) print(f" {'Acceptance Rate (tok/iter)':<30} {metrics.acceptance_rate():>6.2f}") print(f" {'DLM Throughput (tok/s)':<30} {metrics.dlm_tok_rate():>6.1f}") - if tlm_config.model_type == "qwen3_5": + if tlm_config.model_type in {"qwen3_5", "qwen3_5_moe"}: round_rate = metrics.num_total_iters / metrics.tlm_decode_time if metrics.tlm_decode_time > 0 else 0.0 print(f" {'TLM Rounds (round/s)':<30} {round_rate:>6.1f}") else: diff --git a/examples/performance/dflash/utils.py b/examples/performance/dflash/utils.py index a9f81c0e57..a12177a2ff 100644 --- a/examples/performance/dflash/utils.py +++ b/examples/performance/dflash/utils.py @@ -35,6 +35,7 @@ "Qwen3-4B": ("Qwen/Qwen3-4B", "z-lab/Qwen3-4B-DFlash-b16"), "Qwen3.5-4B": ("Qwen/Qwen3.5-4B", "z-lab/Qwen3.5-4B-DFlash"), "Qwen3.6-27B": ("Qwen/Qwen3.6-27B", "z-lab/Qwen3.6-27B-DFlash"), + "Qwen3.6-35B-A3B": ("Qwen/Qwen3.6-35B-A3B", "z-lab/Qwen3.6-35B-A3B-DFlash"), "Qwen3-8B": ("Qwen/Qwen3-8B", "z-lab/Qwen3-8B-DFlash-b16"), "Llama-3.1-8B-Instruct": ("meta-llama/Llama-3.1-8B-Instruct", "z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat"), "Qwen3-VL-32B-Instruct": ( @@ -373,6 +374,13 @@ def compile_qwen3vl_vlm_qpcs( logger.info(f"[compile_qwen3vl_tlm] base={tlm_repo} dlm={dlm_repo} block_size={block_size}") config = AutoConfig.from_pretrained(tlm_repo, trust_remote_code=True, token=hf_token) is_qwen35 = config.model_type == "qwen3_5" + is_qwen35_moe = config.model_type == "qwen3_5_moe" + qaic_config = {"target_layer_ids": tlm_target_ids} + if is_qwen35_moe: + if block_size < 2 or block_size & (block_size - 1): + raise ValueError("Qwen3.5-MoE DFlash requires a power-of-two draft block_size >= 2.") + if prefill_seq_len < block_size or prefill_seq_len % block_size: + raise ValueError(f"Qwen3.5-MoE DFlash requires prefill_seq_len to be a positive multiple of {block_size}.") if is_qwen35: if prefill_seq_len != QWEN3_5_DFLASH_CHUNK_SIZE: logger.warning( @@ -390,7 +398,7 @@ def compile_qwen3vl_vlm_qpcs( kv_offload=True, ignore_mismatched_sizes=True, token=hf_token, - qaic_config={"target_layer_ids": tlm_target_ids}, + qaic_config=qaic_config, ) # DFlashTLMTransform attached fc/hidden_norm but left them random; copy in the real @@ -421,16 +429,19 @@ def compile_qwen3vl_vlm_qpcs( "split_model_io": True, "dflash_block_size": block_size, } - if is_qwen35: - # Verification and replay use the same 16-token shape as prefill. + if is_qwen35 or is_qwen35_moe: + # Hybrid verification and accepted-prefix replay use chunk attention. compile_kwargs["prefill_only"] = True + if is_qwen35_moe: + compile_kwargs["qaic_config"] = qaic_config + compile_kwargs["enable_chunking"] = True if height is not None: compile_kwargs["height"] = height if width is not None: compile_kwargs["width"] = width qpc = tlm_qeff.compile(**compile_kwargs) - lang_qpc = str(qpc["lang_prefill_qpc_path"] if is_qwen35 else qpc["lang_qpc_path"]) + lang_qpc = str(qpc["lang_prefill_qpc_path"] if is_qwen35 or is_qwen35_moe else qpc["lang_qpc_path"]) vision_qpc = str(qpc["vision_qpc_path"]) logger.info(f"[compile_qwen3vl_tlm] lang_qpc={lang_qpc} vision_qpc={vision_qpc}") return lang_qpc, vision_qpc diff --git a/tests/transformers/spd/test_dflash_spd_inference.py b/tests/transformers/spd/test_dflash_spd_inference.py index c757c45155..fa3cfc506c 100644 --- a/tests/transformers/spd/test_dflash_spd_inference.py +++ b/tests/transformers/spd/test_dflash_spd_inference.py @@ -97,7 +97,14 @@ def stop_loading(*args, **kwargs): @pytest.mark.parametrize("ctx_len", [4096, 8192]) -def test_qwen35_vision_cli_accepts_long_context(monkeypatch, ctx_len): +@pytest.mark.parametrize( + "model_type,model_name", + [ + ("qwen3_5", "Qwen/Qwen3.6-27B"), + ("qwen3_5_moe", "Qwen/Qwen3.6-35B-A3B"), + ], +) +def test_qwen35_vision_cli_accepts_long_context(monkeypatch, ctx_len, model_type, model_name): from examples.performance.dflash import utils class ReachedCompilation(Exception): @@ -108,10 +115,10 @@ class ReachedCompilation(Exception): monkeypatch.setattr( sys, "argv", - ["basic_inference_vision.py", "--model_name", "Qwen/Qwen3.6-27B", "--image", "--ctx_len", str(ctx_len)], + ["basic_inference_vision.py", "--model_name", model_name, "--image", "--ctx_len", str(ctx_len)], ) monkeypatch.setattr( - utils.AutoConfig, "from_pretrained", lambda *args, **kwargs: SimpleNamespace(model_type="qwen3_5") + utils.AutoConfig, "from_pretrained", lambda *args, **kwargs: SimpleNamespace(model_type=model_type) ) def stop_compile(*args, **kwargs): @@ -153,8 +160,16 @@ def test_qwen3_5_runtime_rejects_unsupported_prefill_qpc(prompt_chunk_size): ) -@pytest.mark.parametrize("target_seq_len", [16, 64]) -def test_qwen3_5_runtime_pads_verification_and_commit(monkeypatch, target_seq_len): +@pytest.mark.parametrize( + "model_type,prompt_chunk_size,target_seq_len", + [ + ("qwen3_5", 16, 16), + ("qwen3_5", 64, 64), + ("qwen3_5_moe", 16, 16), + ("qwen3_5_moe", 128, 16), + ], +) +def test_qwen3_5_runtime_pads_verification_and_commit(monkeypatch, model_type, prompt_chunk_size, target_seq_len): import torch from QEfficient.generation import dflash_generation as generation @@ -172,8 +187,9 @@ def run(feeds): session = SimpleNamespace( input_names=["cache_commit", "vision_embeds"], - binding_index_map={"vision_embeds": 0}, - bindings=[SimpleNamespace(type=0, dims=(1, 1, 32))], + binding_index_map={"vision_embeds": 0, "input_ids": 1}, + bindings=[SimpleNamespace(type=0, dims=(1, 1, 32)), SimpleNamespace(dims=(1, prompt_chunk_size))], + allowed_shapes=[[(0, [1, 1, 32]), (8, [1, length])] for length in {prompt_chunk_size, target_seq_len}], aic_to_np_dtype_mapping={0: np.dtype(np.float32)}, run=run, set_buffers=buffers.append, @@ -191,7 +207,9 @@ def run(feeds): ) def core(**kwargs): - kwargs["tlm_prefill"](0, target_seq_len) + kwargs["tlm_prefill"](0, prompt_chunk_size) + assert calls[-1]["input_ids"].shape[1] == prompt_chunk_size + assert calls[-1]["cache_commit"].item() == 1 kwargs["prepare_decode"]() ids = np.arange(16, dtype=np.int64).reshape(1, 16) positions = np.arange(5, 21, dtype=np.int64).reshape(1, 16) @@ -215,17 +233,73 @@ def core(**kwargs): prompt_text="test", tokenizer=SimpleNamespace(pad_token_id=0, eos_token_id=None), processor=None, - tlm_config=SimpleNamespace(model_type="qwen3_5"), + tlm_config=SimpleNamespace(model_type=model_type), dlm_session=None, tlm_session=session, mask_token_id=2, vocab_size=64, - prompt_chunk_size=target_seq_len, + prompt_chunk_size=prompt_chunk_size, ctx_len=256, block_size=16, hidden_size=32, ) - assert all(value.shape[1] == target_seq_len for item in buffers for key, value in item.items() if key == "logits") + assert [item["logits"].shape[1] for item in buffers if "logits" in item] == [prompt_chunk_size, target_seq_len] + + +@pytest.mark.parametrize("with_image", [False, True]) +def test_qwen35_moe_runtime_mrope_matches_hf(with_image): + import torch + + from QEfficient.generation.dflash_generation import compute_position_ids_qwen3_vl + from tests.unit_test.transforms.test_dflash_transform import make_tiny_qwen3_5_vlm + + model = make_tiny_qwen3_5_vlm(moe=True) + ids = torch.tensor([[10, 5, 3, 3, 3, 3, 6, 11]] if with_image else [[10, 11, 12]]) + mask = torch.ones_like(ids) + types = (ids == model.config.image_token_id).long() + grid = torch.tensor([[1, 4, 4]]) if with_image else None + expected, expected_delta = model.model.get_rope_index(ids, types, image_grid_thw=grid, attention_mask=mask) + actual, delta = compute_position_ids_qwen3_vl(ids, mask, grid, model.config, types) + torch.testing.assert_close(actual[0], torch.arange(ids.shape[1]).view(1, -1)) + torch.testing.assert_close(actual[1:], expected) + torch.testing.assert_close(delta, expected_delta) + + +@pytest.mark.parametrize("block_size", [8, 16]) +def test_qwen35_moe_compile_preserves_dflash_metadata(monkeypatch, block_size): + import torch + + from examples.performance.dflash import utils + + calls = [] + text_model = SimpleNamespace(fc=torch.nn.Linear(8, 4, bias=False), hidden_norm=torch.nn.LayerNorm(4)) + qaic_config = {"target_layer_ids": [2, 4]} + + def compile_model(**kwargs): + calls.append(kwargs) + return {"lang_prefill_qpc_path": "target.qpc", "vision_qpc_path": "vision.qpc"} + + def load_model(*args, **kwargs): + assert kwargs["qaic_config"] == qaic_config + assert kwargs["dtype"] == "float32" + return SimpleNamespace( + lang_model=SimpleNamespace(model=SimpleNamespace(language_model=text_model)), compile=compile_model + ) + + monkeypatch.setattr(utils.AutoConfig, "from_pretrained", lambda *a, **k: SimpleNamespace(model_type="qwen3_5_moe")) + monkeypatch.setattr( + utils, + "read_dlm_meta", + lambda *a: ({"fc.weight": torch.ones(4, 8), "hidden_norm.weight": torch.ones(4)}, [1, 3], block_size), + ) + monkeypatch.setattr(utils.QEFFAutoModelForImageTextToText, "from_pretrained", load_model) + assert utils.compile_qwen3vl_vlm_qpcs( + "target", "draft", ctx_len=8192, prefill_seq_len=128, num_cores=8, num_devices=4 + ) == ("target.qpc", "vision.qpc") + assert calls[0]["qaic_config"] == qaic_config + assert calls[0]["prefill_only"] and calls[0]["enable_chunking"] + assert calls[0]["prefill_seq_len"] == 128 and calls[0]["dflash_block_size"] == block_size + assert calls[0]["mxint8_kv_cache"] class _FakeTokenizer: diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index 086afd75ed..a531be93e9 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -861,8 +861,51 @@ def _ort_qwen_logits(onnx_path: Path, qeff_model, inputs): return outputs["logits"].astype(np.float32) -def test_qwen3_5_without_dflash_hf_qeff_ort_parity(tmp_path): - _assert_qwen_hf_qeff_ort_parity("qwen3_5", tmp_path) +@pytest.mark.parametrize("model_type", ["qwen3_5", "qwen3_5_moe"]) +def test_qwen3_5_without_dflash_hf_qeff_ort_parity(tmp_path, model_type): + _assert_qwen_hf_qeff_ort_parity(model_type, tmp_path) + + +@pytest.mark.parametrize("output_hidden_states", [False, True]) +@pytest.mark.parametrize("explicit_cache_position", [False, True]) +def test_qwen3_5_moe_without_dflash_preserves_forward_path(monkeypatch, output_hidden_states, explicit_cache_position): + from QEfficient.transformers.models.qwen3_5_moe import modeling_qwen3_5_moe as modeling + from tests.unit_test.transforms.test_dflash_transform import make_tiny_qwen3_5_vlm + + qeff = QEFFAutoModelForImageTextToText(make_tiny_qwen3_5_vlm(moe=True), kv_offload=True) + model = qeff.lang_model.model.language_model + inputs = qeff.model.get_dummy_inputs(kv_offload=True)["lang"] + conv_update = modeling.qeff_torch_causal_conv1d_update + conv_calls = [] + + def baseline_conv_update(*args, **kwargs): + assert kwargs.get("dflash", False) is False + conv_calls.append(True) + return conv_update(*args, **kwargs) + + def unexpected_dflash(*args, **kwargs): + raise AssertionError("Non-DFlash forward must not use DFlash operations") + + monkeypatch.setattr(modeling, "qeff_torch_causal_conv1d_update", baseline_conv_update) + monkeypatch.setattr(modeling, "compute_dflash_target_hidden_states", unexpected_dflash) + monkeypatch.setattr(modeling, "dflash_chunk_gated_delta_scan", unexpected_dflash) + monkeypatch.setattr(modeling.QEffQwen3_5MoeGatedDeltaNet, "_dflash_triangular_inverse", unexpected_dflash) + with torch.no_grad(): + output = model( + input_ids=inputs["input_ids"], + position_ids=inputs["position_ids"], + past_key_values=inputs["past_key_values"], + attention_mask=torch.ones_like(inputs["input_ids"]), + cache_position=torch.arange(inputs["input_ids"].shape[1]) if explicit_cache_position else None, + output_hidden_states=output_hidden_states, + use_cache=True, + ) + + assert len(conv_calls) == model.config.layer_types.count("linear_attention") + assert output.hidden_states is None + assert "cache_commit" not in inputs + assert "hidden_states" not in qeff.model.get_output_names(kv_offload=True)["lang"] + assert "hidden_states" not in qeff.model.get_onnx_dynamic_axes(kv_offload=True)["lang"] def test_dflash_mixed_window_ort_parity(tmp_path): @@ -918,17 +961,27 @@ def test_dflash_mixed_window_ort_parity(tmp_path): @pytest.mark.parametrize( - "prefill_length,verification_length,dflash_chunk_size", - [(8, 16, 64), (64, 16, 64), (8, 64, 64), (64, 64, 64), (8, 16, 16), (16, 16, 16)], + "moe,prefill_length,verification_length,dflash_chunk_size", + [ + (False, 8, 16, 64), + (False, 64, 16, 64), + (False, 8, 64, 64), + (False, 64, 64, 64), + (False, 8, 16, 16), + (False, 16, 16, 16), + (True, 16, 16, 16), + (True, 128, 16, 16), + (True, 128, 8, 8), + ], ) @pytest.mark.parametrize("matrix_solve", [False, True]) def test_qwen3_5_dflash_ort_recurrent_commit( - tmp_path, prefill_length, verification_length, dflash_chunk_size, matrix_solve + tmp_path, prefill_length, verification_length, dflash_chunk_size, matrix_solve, moe ): from tests.unit_test.transforms.test_dflash_transform import make_tiny_qwen3_5_vlm torch.manual_seed(11) - model = make_tiny_qwen3_5_vlm() + model = make_tiny_qwen3_5_vlm(moe=moe) model.config.text_config.max_position_embeddings = 256 qeff = QEFFAutoModelForImageTextToText(model, kv_offload=True, qaic_config={"target_layer_ids": [1, 3]}) qeff.lang_model.model.language_model.dflash_chunk_size = dflash_chunk_size @@ -957,14 +1010,15 @@ def test_qwen3_5_dflash_ort_recurrent_commit( batch_size = inputs["input_ids"].shape[0] actual = {} prompt_length = prefill_length - 3 + verified_tokens = min(verification_length, 16) # Keep one exported graph across different sequence lengths and acceptance lengths. for start, length, valid, commit in [ (0, prefill_length, prompt_length, 1), - (prompt_length, verification_length, 16, 0), + (prompt_length, verification_length, verified_tokens, 0), (prompt_length, verification_length, 2, 1), - (prompt_length + 2, verification_length, 16, 0), + (prompt_length + 2, verification_length, verified_tokens, 0), (prompt_length + 2, verification_length, 1, 1), - (prompt_length + 3, verification_length, 16, 0), + (prompt_length + 3, verification_length, verified_tokens, 0), ]: inputs["input_ids"] = torch.full((batch_size, length), 10 + start, dtype=torch.int64) positions = torch.arange(start, start + length).view(1, 1, -1).repeat(4, batch_size, 1) @@ -3478,6 +3532,32 @@ def test_qwen3_5_moe_conv_decode_slice_keeps_prefill_gather_path(): torch.testing.assert_close(state, expected_state) +@pytest.mark.parametrize("dflash", [None, False, True]) +@pytest.mark.parametrize("grouped_state", [False, True]) +@pytest.mark.parametrize("valid_tokens", [0, 1, 2, 3, 4]) +def test_qwen3_5_moe_conv_history_dflash_guard(dflash, grouped_state, valid_tokens): + from QEfficient.transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( + qeff_torch_causal_conv1d_update, + ) + + flat_state = torch.arange(24, dtype=torch.float32).reshape(1, 6, 4) + state = flat_state.reshape(1, 2, 3, 4) if grouped_state else flat_state + hidden = 100 + torch.arange(6 * 16, dtype=torch.float32).reshape(1, 6, 16) + positions = torch.full((4, 1, 16), -1, dtype=torch.int64) + positions[..., :valid_tokens] = torch.arange(40, 40 + valid_tokens) + kwargs = {} if dflash is None else {"dflash": dflash} + _, actual = qeff_torch_causal_conv1d_update(hidden, state, torch.ones(6, 4), positions, **kwargs) + + if dflash: + expected = torch.cat([flat_state, hidden[..., :valid_tokens]], dim=-1)[..., -4:] + else: + # Preserve the original baseline ordering, including its tied cached keys. + zeros = flat_state[:, 0, :].to(positions.dtype) * 0 + order = torch.argsort(torch.cat([zeros, positions[0]], dim=1), dim=1) + expected = torch.gather(torch.cat([flat_state, hidden], dim=-1), 2, order[:, None, -4:].expand(-1, 6, -1)) + torch.testing.assert_close(actual, expected.reshape_as(state), rtol=0, atol=0) + + @pytest.mark.parametrize("valid_tokens", [0, 1, 2, 3, 4]) def test_qwen3_5_conv_replay_preserves_history_order(valid_tokens): from QEfficient.transformers.models.qwen3_5.modeling_qwen3_5 import qeff_torch_causal_conv1d_update diff --git a/tests/unit_test/transforms/test_dflash_transform.py b/tests/unit_test/transforms/test_dflash_transform.py index 30f9bfc8b9..5ca55653cc 100644 --- a/tests/unit_test/transforms/test_dflash_transform.py +++ b/tests/unit_test/transforms/test_dflash_transform.py @@ -34,14 +34,31 @@ HIDDEN_SIZE = 64 -def make_tiny_qwen3_5_vlm(): - from transformers import Qwen3_5Config, Qwen3_5ForConditionalGeneration +def make_tiny_qwen3_5_vlm(moe=False): + from transformers import ( + Qwen3_5Config, + Qwen3_5ForConditionalGeneration, + Qwen3_5MoeConfig, + Qwen3_5MoeForConditionalGeneration, + ) - config = Qwen3_5Config( + config_cls = Qwen3_5MoeConfig if moe else Qwen3_5Config + model_cls = Qwen3_5MoeForConditionalGeneration if moe else Qwen3_5ForConditionalGeneration + config = config_cls( text_config={ "vocab_size": VOCAB_SIZE, "hidden_size": 32, "intermediate_size": 64, + **( + { + "num_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": 32, + "shared_expert_intermediate_size": 32, + } + if moe + else {} + ), "num_hidden_layers": 4, "num_attention_heads": 2, "num_key_value_heads": 1, @@ -75,21 +92,39 @@ def make_tiny_qwen3_5_vlm(): vision_end_token_id=6, ) config._attn_implementation = "eager" - return Qwen3_5ForConditionalGeneration(config).eval() + return model_cls(config).eval() @pytest.mark.parametrize("accepted", [0, 1, 3]) @pytest.mark.parametrize("prompt_len", [1, 5]) @pytest.mark.parametrize("chunk_size", [16, 64]) @pytest.mark.parametrize("matrix_solve", [False, True]) -def test_qwen3_5_dflash_replay_matches_autoregressive(accepted, prompt_len, chunk_size, matrix_solve, monkeypatch): +@pytest.mark.parametrize("moe_flavour", [None, "decode_bmm", "expert_parallel"]) +def test_qwen3_5_dflash_replay_matches_autoregressive( + accepted, prompt_len, chunk_size, matrix_solve, moe_flavour, monkeypatch +): """Verify then replay a prefix; next-round results must exclude rejected tokens.""" from QEfficient import QEFFAutoModelForImageTextToText torch.manual_seed(7) - hf_model = make_tiny_qwen3_5_vlm() + moe = moe_flavour is not None + hf_model = make_tiny_qwen3_5_vlm(moe=moe) reference = copy.deepcopy(hf_model) - qeff = QEFFAutoModelForImageTextToText(hf_model, kv_offload=True, qaic_config={"target_layer_ids": [1, 3]}) + qaic_config = {"target_layer_ids": [1, 3]} + if moe: + qaic_config["moe_config"] = {"flavour": moe_flavour} + qeff = QEFFAutoModelForImageTextToText(hf_model, kv_offload=True, qaic_config=qaic_config) + if moe: + qeff.transform( + ctx_len=CTX_LEN, + seq_len=16, + bs=1, + num_devices=1, + aic_num_cores=4, + qaic_config=qaic_config, + prefill_only=True, + prefill_seq_len=16, + ) model = qeff.lang_model.model.language_model model.dflash_chunk_size = chunk_size for layer in model.layers: @@ -99,8 +134,10 @@ def test_qwen3_5_dflash_replay_matches_autoregressive(accepted, prompt_len, chun def unexpected_decode(*args, **kwargs): raise AssertionError("DFlash must not compute the one-token recurrent branch") + model_type = "qwen3_5_moe" if moe else "qwen3_5" monkeypatch.setattr( - "QEfficient.transformers.models.qwen3_5.modeling_qwen3_5.recurrent_gdn_decode_forward", unexpected_decode + f"QEfficient.transformers.models.{model_type}.modeling_{model_type}.recurrent_gdn_decode_forward", + unexpected_decode, ) initial = qeff.model.get_dummy_inputs(kv_offload=True)["lang"]["past_key_values"] tokens = torch.randint(10, VOCAB_SIZE, (1, 16)) @@ -146,6 +183,61 @@ def run(ids, start, state, commit, valid=None): torch.testing.assert_close(actual, expected, rtol=2e-4, atol=2e-4) +@pytest.mark.parametrize( + "moe_config", [None, {"cores_per_expert": 2, "tree_reduce": True, "expert_parallel_chunk_size": 16}] +) +def test_qwen3_5_moe_dflash_preserves_upstream_expert_parallel(moe_config): + from transformers import Qwen3_5MoeForConditionalGeneration + + from QEfficient import QEFFAutoModelForImageTextToText + from QEfficient.transformers.moe import MoEFlavour + from QEfficient.transformers.moe.block import QEffMoEBlockMixin + + torch.manual_seed(19) + config = make_tiny_qwen3_5_vlm(moe=True).config + config.text_config.num_experts = 256 + config.text_config.num_experts_per_tok = 8 + reference = Qwen3_5MoeForConditionalGeneration(config).eval() + qaic_config = {"target_layer_ids": [1, 3]} + if moe_config is not None: + qaic_config["moe_config"] = moe_config + qeff = QEFFAutoModelForImageTextToText(copy.deepcopy(reference), kv_offload=True, qaic_config=qaic_config) + qeff.transform( + ctx_len=256, + seq_len=128, + bs=1, + num_devices=4, + aic_num_cores=8, + qaic_config=qaic_config, + prefill_only=True, + prefill_seq_len=128, + ) + options = moe_config or {} + cores_per_expert = options.get("cores_per_expert", 1) + packed_chunks = 8 if moe_config is not None else 1 + language_model = qeff.lang_model.model.language_model + assert language_model.target_layer_ids == [1, 3] + for layer in language_model.layers: + block = layer.mlp + assert block.forward.__func__ is QEffMoEBlockMixin.forward + assert block._moe_flavour is MoEFlavour.EXPERT_PARALLEL + assert block.num_devices == 4 and block.total_avl_cores == 32 + assert block.num_parallelized_experts == 32 // cores_per_expert + assert block.num_pipeline_stages == 8 * cores_per_expert + assert block.tree_reduce is options.get("tree_reduce", False) + assert block.expert_parallel_num_packed_chunks == packed_chunks + assert qeff.lang_model.hash_params["moe_prefill_flavour"] == "expert_parallel" + assert qeff.lang_model.hash_params["moe_prefill_num_packed_chunks"] == packed_chunks + + # Check routed and shared experts at both verification and prompt-prefill lengths. + with torch.no_grad(): + for seq_len in (16, 128): + hidden = torch.randn(1, seq_len, config.text_config.hidden_size) + actual = language_model.layers[0].mlp(hidden) + expected = reference.model.language_model.layers[0].mlp(hidden) + torch.testing.assert_close(actual, expected, atol=2e-4, rtol=2e-4) + + def test_qwen3_5_dflash_requires_explicit_commit(): from QEfficient import QEFFAutoModelForImageTextToText From 50284684e2898027092c8fc73dda8aa31de31930 Mon Sep 17 00:00:00 2001 From: Vahid Janfaza Date: Mon, 12 Oct 2026 22:54:53 -0700 Subject: [PATCH 4/9] Fix DFlash prefill with num_logits_to_keep Select only the last valid position before the LM head during DFlash prefill, while retaining all verification positions and DLM target features. Keep non-DFlash behavior unchanged. Co-authored-by: FNU Annanya Co-authored-by: Mohammed Faizanulla Azeez Signed-off-by: Vahid Janfaza --- QEfficient/generation/dflash_generation.py | 39 ++++++++++++++----- .../models/gemma4/modeling_gemma4.py | 23 +++++++++-- .../models/llama/modeling_llama.py | 4 ++ .../transformers/models/modeling_auto.py | 31 ++++++++++----- .../models/qwen3/modeling_qwen3.py | 4 ++ .../models/qwen3_5/modeling_qwen3_5.py | 33 +++++++++++++--- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 32 ++++++++++++--- .../models/qwen3_vl/modeling_qwen3_vl.py | 19 ++++++++- examples/performance/dflash/README.md | 7 ++++ .../spd/test_dflash_spd_inference.py | 7 +--- .../unit_test/models/test_model_quickcheck.py | 5 +++ 11 files changed, 164 insertions(+), 40 deletions(-) diff --git a/QEfficient/generation/dflash_generation.py b/QEfficient/generation/dflash_generation.py index 77e04a0977..4f481996c7 100644 --- a/QEfficient/generation/dflash_generation.py +++ b/QEfficient/generation/dflash_generation.py @@ -106,6 +106,12 @@ def _prepare_text_inputs(tokenizer, prompt_text: str, prompt_chunk_size: int) -> return inputs["input_ids"], position_ids +def _dflash_logit_selection_inputs(tlm_session, num_logits: int) -> dict[str, np.ndarray]: + if "num_logits_to_keep" not in getattr(tlm_session, "input_names", ()): + return {} + return {"num_logits_to_keep": np.zeros((num_logits, 1), dtype=np.int64)} + + def _run_spd_core( *, input_ids: np.ndarray, @@ -179,7 +185,8 @@ def run_dlm_prefill_block(hidden_states, position_ids, sub_start, sub_len): continue last_prefill_pos_in_chunk = int(flat_pos_chunk.argmax()) - new_tlm_token = tlm_logits[:, last_prefill_pos_in_chunk] + logit_index = 0 if tlm_logits.shape[1] == 1 else last_prefill_pos_in_chunk + new_tlm_token = tlm_logits[:, logit_index] last_sub_start = (last_prefill_pos_in_chunk // block_size) * block_size for sub_start in range(0, last_sub_start, block_size): @@ -311,7 +318,10 @@ def run_spd_inference_single( generated_ids = _initialize_generated_ids(batch_size, ctx_len, input_ids.shape[1], tokenizer.pad_token_id) metrics = SpecDecodingMetrics(block_size=block_size) - tlm_session.set_buffers({"logits": np.zeros((batch_size, prompt_chunk_size), dtype=np.int32)}) + prefill_selection = _dflash_logit_selection_inputs(tlm_session, 1) + decode_selection = _dflash_logit_selection_inputs(tlm_session, block_size) + prefill_logits = 1 if prefill_selection else prompt_chunk_size + tlm_session.set_buffers({"logits": np.zeros((batch_size, prefill_logits), dtype=np.int32)}) tlm_session.set_buffers({"hidden_states": np.zeros((batch_size, prompt_chunk_size, hidden_size), dtype=np.float32)}) def tlm_prefill(chunk_start, chunk_end): @@ -319,6 +329,7 @@ def tlm_prefill(chunk_start, chunk_end): { "input_ids": input_ids[:, chunk_start:chunk_end], "position_ids": position_ids[:, chunk_start:chunk_end], + **prefill_selection, } ) return outputs["logits"], outputs["hidden_states"] @@ -328,7 +339,7 @@ def prepare_decode(): tlm_session.set_buffers({"hidden_states": np.zeros((batch_size, block_size, hidden_size), dtype=np.float32)}) def tlm_decode(draft_ids, decode_position_ids): - outputs = tlm_session.run({"input_ids": draft_ids, "position_ids": decode_position_ids}) + outputs = tlm_session.run({"input_ids": draft_ids, "position_ids": decode_position_ids, **decode_selection}) return outputs["logits"], outputs["hidden_states"] eos_token_ids = {tokenizer.eos_token_id} if tokenizer.eos_token_id is not None else set() @@ -394,6 +405,9 @@ def run_spd_inference_gemma4( generated_ids = _initialize_generated_ids(batch_size, ctx_len, padded_ids.shape[1], tokenizer.pad_token_id) metrics = SpecDecodingMetrics(block_size=block_size) tlm_hidden_name = "target_hidden_states" if "target_hidden_states" in tlm_session.output_names else "hidden_states" + prefill_selection = _dflash_logit_selection_inputs(tlm_session, 1) + decode_selection = _dflash_logit_selection_inputs(tlm_session, block_size) + prefill_logits = 1 if prefill_selection else prompt_chunk_size tlm_inputs = set(tlm_session.input_names) tlm_needs_vision = "vision_embeds" in tlm_inputs @@ -420,9 +434,9 @@ def run_spd_inference_gemma4( tlm_image_idx = np.array([[0]], dtype=np.int64) - def tlm_run(ids, model_position_ids, mm=None): + def tlm_run(ids, model_position_ids, selection, mm=None): nonlocal tlm_image_idx - feeds = {"input_ids": ids, "position_ids": model_position_ids} + feeds = {"input_ids": ids, "position_ids": model_position_ids, **selection} if tlm_needs_image_idx: feeds["image_idx"] = tlm_image_idx if tlm_needs_mm_ids: @@ -441,9 +455,9 @@ def tlm_run(ids, model_position_ids, mm=None): return token_ids, outputs[tlm_hidden_name] if tlm_logits_is_full: - tlm_session.set_buffers({"logits": np.zeros((batch_size, prompt_chunk_size, vocab_size), dtype=np.float32)}) + tlm_session.set_buffers({"logits": np.zeros((batch_size, prefill_logits, vocab_size), dtype=np.float32)}) else: - tlm_session.set_buffers({"logits": np.zeros((batch_size, prompt_chunk_size), dtype=np.int32)}) + tlm_session.set_buffers({"logits": np.zeros((batch_size, prefill_logits), dtype=np.int32)}) tlm_session.set_buffers({tlm_hidden_name: np.zeros((batch_size, prompt_chunk_size, hidden_size), dtype=np.float32)}) def tlm_prefill(chunk_start, chunk_end): @@ -451,6 +465,7 @@ def tlm_prefill(chunk_start, chunk_end): return tlm_run( padded_ids[:, chunk_start:chunk_end], position_ids[:, chunk_start:chunk_end], + prefill_selection, mm=mm_chunk, ) @@ -464,7 +479,7 @@ def prepare_decode(): tlm_session.skip_buffers(["vision_embeds"]) def tlm_decode(draft_ids, decode_position_ids): - return tlm_run(draft_ids, decode_position_ids) + return tlm_run(draft_ids, decode_position_ids, decode_selection) eos_token_ids = {tokenizer.eos_token_id} if tokenizer.eos_token_id is not None else set() return _run_spd_core( @@ -786,7 +801,10 @@ def run_spd_inference_qwen3_vl( lang_extra = {"image_idx": np.array([[0]], dtype=np.int64)} if is_qwen35: lang_extra["cache_commit"] = np.ones(1, dtype=np.int64) - tlm_session.set_buffers({"logits": np.zeros((batch_size, prompt_chunk_size), dtype=np.int32)}) + prefill_selection = _dflash_logit_selection_inputs(tlm_session, 1) + decode_selection = _dflash_logit_selection_inputs(tlm_session, target_seq_len) + prefill_logits = 1 if prefill_selection else prompt_chunk_size + tlm_session.set_buffers({"logits": np.zeros((batch_size, prefill_logits), dtype=np.int32)}) tlm_session.set_buffers({"hidden_states": np.zeros((batch_size, prompt_chunk_size, hidden_size), dtype=np.float32)}) def tlm_prefill(chunk_start, chunk_end): @@ -794,6 +812,7 @@ def tlm_prefill(chunk_start, chunk_end): { "input_ids": input_ids_np[:, chunk_start:chunk_end], "position_ids": position_ids_np[:, :, chunk_start:chunk_end], + **prefill_selection, **lang_extra, } ) @@ -816,7 +835,7 @@ def decode_inputs(draft_ids, decode_position_ids, valid_tokens): draft_ids = np.pad(draft_ids, ((0, 0), (0, padding)), constant_values=tokenizer.pad_token_id) positions = np.pad(positions, ((0, 0), (0, 0), (0, padding)), constant_values=-1) positions[:, :, valid_tokens:] = -1 - return {"input_ids": draft_ids, "position_ids": positions, **lang_extra} + return {"input_ids": draft_ids, "position_ids": positions, **decode_selection, **lang_extra} def tlm_decode(draft_ids, decode_position_ids): outputs = tlm_session.run(decode_inputs(draft_ids, decode_position_ids, block_size)) diff --git a/QEfficient/transformers/models/gemma4/modeling_gemma4.py b/QEfficient/transformers/models/gemma4/modeling_gemma4.py index c88ab7a04f..5f2220e632 100644 --- a/QEfficient/transformers/models/gemma4/modeling_gemma4.py +++ b/QEfficient/transformers/models/gemma4/modeling_gemma4.py @@ -49,6 +49,7 @@ silu_glu_mlp, ) from QEfficient.transformers.spd.dflash import compute_dflash_target_hidden_states +from QEfficient.transformers.spd.spd_transform_forward import filter_hidden_states from QEfficient.utils import constants _FP16_CLAMP_MIN = -65504.0 @@ -1114,6 +1115,7 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, + num_logits_to_keep: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs, ) -> CausalLMOutputWithPast: @@ -1129,8 +1131,11 @@ def forward( ) hidden_states = outputs.last_hidden_state - # DFlash TLM emits full per-position logits; the plain path keeps last-token logits. - if not getattr(self.model, "target_layer_ids", None): + # DFlash selects one prefill position or the full verification block. + if getattr(self.model, "target_layer_ids", None): + if num_logits_to_keep is not None: + hidden_states = filter_hidden_states(hidden_states, position_ids, num_logits_to_keep) + else: if position_ids is not None: logit_index = position_ids.to(torch.int32).argmax(1, keepdim=True) hidden_states = hidden_states[torch.arange(position_ids.shape[0]).view(-1, 1), logit_index] @@ -1170,6 +1175,7 @@ def forward( mm_token_type_ids=None, batch_index: torch.LongTensor | None = None, comp_ctx_lengths: list[int] | None = None, + num_logits_to_keep: torch.LongTensor | None = None, **kwargs, ): del kwargs @@ -1234,9 +1240,11 @@ def forward( ) finally: _DISABLE_EXPORT_FP16_CLAMP = restore_disable_clamp - # DFlash TLM emits full per-position logits; the plain path keeps last-token logits. + # DFlash selects one prefill position or the full verification block. if getattr(self.language_model, "target_layer_ids", None): hidden_states = outputs[0] + if num_logits_to_keep is not None: + hidden_states = filter_hidden_states(hidden_states, position_ids, num_logits_to_keep) else: logit_index = position_ids.to(torch.int32).argmax(1, keepdim=True) hidden_states = outputs[0][torch.arange(position_ids.shape[0]).view(-1, 1), logit_index] @@ -1384,6 +1392,9 @@ def get_specializations( prefill_seq_len = prefill_seq_len if prefill_seq_len else 32 ctx_len = ctx_len if ctx_len else constants.INTERN_CTX_LEN max_patches = self._get_vision_max_patches() + dflash_block_size = None + if getattr(self.model.language_model, "target_layer_ids", None): + dflash_block_size = compiler_options.pop("dflash_block_size", None) user_vision_size = compiler_options.pop("vision_size", None) if user_vision_size: if user_vision_size >= ctx_len: @@ -1403,6 +1414,8 @@ def build_lang_prefill_spec(comp_ctx_lengths: int | None = None): "vision_batch_size": vision_batch_size, "vision_size": vision_size, } + if dflash_block_size is not None: + spec["num_logits_to_keep"] = 1 if comp_ctx_lengths is not None: spec["comp_ctx_lengths"] = comp_ctx_lengths if continuous_batching: @@ -1416,12 +1429,14 @@ def build_lang_prefill_spec(comp_ctx_lengths: int | None = None): def build_lang_decode_spec(comp_ctx_lengths: int | None = None): spec = { "batch_size": full_batch_size if continuous_batching else batch_size, - "seq_len": compiler_options.pop("dflash_block_size", None) or "1", + "seq_len": dflash_block_size or "1", "ctx_len": ctx_len, "sliding_window": self.model.language_model.config.sliding_window, "vision_batch_size": vision_batch_size, "vision_size": vision_size, } + if dflash_block_size is not None: + spec["num_logits_to_keep"] = dflash_block_size if comp_ctx_lengths is not None: spec["comp_ctx_lengths"] = comp_ctx_lengths if continuous_batching: diff --git a/QEfficient/transformers/models/llama/modeling_llama.py b/QEfficient/transformers/models/llama/modeling_llama.py index 940ff058c1..2c48f4cee0 100644 --- a/QEfficient/transformers/models/llama/modeling_llama.py +++ b/QEfficient/transformers/models/llama/modeling_llama.py @@ -32,6 +32,7 @@ from QEfficient.transformers.cache_utils import QEffDynamicCache from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask from QEfficient.transformers.spd.dflash import compute_dflash_target_hidden_states +from QEfficient.transformers.spd.spd_transform_forward import filter_hidden_states from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE @@ -369,6 +370,7 @@ def forward( use_cache: bool | None = None, output_hidden_states: bool | None = None, cache_position: torch.LongTensor | None = None, + num_logits_to_keep: torch.LongTensor | None = None, **kwargs, ) -> tuple | CausalLMOutputWithPast: output_hidden_states = ( @@ -401,6 +403,8 @@ def forward( if getattr(self.model, "target_layer_ids", None): target_hidden = outputs.hidden_states hidden_states = outputs.last_hidden_state + if num_logits_to_keep is not None: + hidden_states = filter_hidden_states(hidden_states, position_ids, num_logits_to_keep) logits = self.lm_head(hidden_states).float() predicted_token_ids = logits.argmax(dim=-1).to(torch.int32) return CausalLMOutputWithPast( diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 00185e5bdf..54312b1e0f 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -1779,6 +1779,13 @@ def export( ) dynamic_axes = self.model.get_onnx_dynamic_axes(**dynamic_axes_kwargs) output_names = self.model.get_output_names(kv_offload=True) + if self.dflash_tlm: + nlk = constants.ONNX_EXPORT_EXAMPLE_NLK + inputs["lang"]["num_logits_to_keep"] = torch.arange(nlk).view(nlk, 1) + dynamic_axes["lang"]["num_logits_to_keep"] = {0: "num_logits_to_keep"} + dynamic_axes["lang"]["logits"] = {0: "batch_size", 1: "num_logits_to_keep"} + hidden_name = "target_hidden_states" if "target_hidden_states" in output_names["lang"] else "hidden_states" + dynamic_axes["lang"][hidden_name] = {0: "batch_size", 1: "seq_len"} # Prefix only the language-side KV-cache retained buffers (vision buffers are untouched). output_names = apply_kv_cache_prefix(output_names, validate_kv_cache_prefix(kv_cache_prefix)) if self.lang_model.qaic_config is not None and self.lang_model.qaic_config.get("include_sampler", False): @@ -2086,6 +2093,10 @@ def compile( compiler_options["artifacts"] = True if skip_lang and skip_vision: raise ValueError("Expected at least one of 'skip_lang' or 'skip_vision' to be False") + if self.dflash_tlm and dflash_block_size is None and self.model.config.model_type in {"qwen3_5", "qwen3_5_moe"}: + dflash_block_size = getattr(self.model.model.language_model, "dflash_chunk_size", 16) + if self.dflash_tlm and not skip_lang and (type(dflash_block_size) is not int or dflash_block_size < 2): + raise ValueError("DFlash TLM requires an integer dflash_block_size >= 2.") reject_legacy_moe_prefill_packed_chunk_size(compiler_options) _ignore_public_mdp_ts_num_devices(compiler_options) if qaic_config and qaic_config.get("gdn_full_state_update") is True: @@ -2221,6 +2232,8 @@ def compile( prefill_seq_len=prefill_seq_len, ) + if self.dflash_tlm: + compiler_options["dflash_block_size"] = dflash_block_size specializations, compiler_options = self.model.get_specializations( batch_size=batch_size, prefill_seq_len=prefill_seq_len, @@ -2236,12 +2249,6 @@ def compile( **compiler_options, ) - # DFlash TLM: override the lang decode spec's seq_len to dflash_block_size. - if getattr(self, "dflash_tlm", None) and dflash_block_size is not None: - for spec in specializations["lang"]: - if str(spec.get("seq_len")) == "1": - spec["seq_len"] = dflash_block_size - custom_io_vision = {} target_dtype = getattr(self.model.config, "torch_dtype", torch.float32) kv_cache_dtype = "mxint8" if mxint8_kv_cache else CUSTOM_IO_DTYPE_MAP[target_dtype] @@ -4429,7 +4436,7 @@ def export( example_inputs["batch_index"] = torch.arange(bs).view(bs, 1) dynamic_axes["batch_index"] = {0: "batch_size"} - if self.is_tlm: + if self.is_tlm or self.dflash_tlm: nlk = constants.ONNX_EXPORT_EXAMPLE_NLK # Number of Logits to Keep example_inputs["num_logits_to_keep"] = torch.arange(nlk).view(nlk, 1) dynamic_axes["num_logits_to_keep"] = {0: "num_logits_to_keep"} @@ -4499,6 +4506,8 @@ def _legacyify_cache(obj): if getattr(self, "dflash_tlm", None): output_names.append("hidden_states") + dynamic_axes["logits"] = {0: "batch_size", 1: "num_logits_to_keep"} + dynamic_axes["hidden_states"] = {0: "batch_size", 1: "seq_len"} if QEFFBaseModel._layerwise_active: return self._export_layerwise( @@ -4583,7 +4592,7 @@ def build_prefill_specialization( } if comp_ctx_lengths is not None: spec["comp_ctx_lengths"] = comp_ctx_lengths - spec["num_logits_to_keep"] = 1 if self.is_tlm else None + spec["num_logits_to_keep"] = 1 if self.is_tlm or self.dflash_tlm else None if self.continuous_batching: spec["full_batch_size"] = kv_cache_batch_size else: @@ -4637,7 +4646,7 @@ def build_decode_specialization( of the prefill specialization (e.g., if prefill_seq_len is 1 and not continuous batching). """ decode_seq_len = (num_speculative_tokens + 1) if self.is_tlm else 1 - if decode_seq_len == prefill_seq_len and not self.continuous_batching: + if decode_seq_len == prefill_seq_len and not self.continuous_batching and not self.dflash_tlm: return None if hasattr(self.model, "get_specializations"): spec = self.model.get_specializations( @@ -4658,6 +4667,10 @@ def build_decode_specialization( if getattr(self, "dflash_tlm", None) or getattr(self, "dflash_dlm", None): spec["seq_len"] = dflash_block_size + if self.dflash_tlm: + if type(dflash_block_size) is not int or dflash_block_size < 2: + raise ValueError("DFlash TLM requires an integer dflash_block_size >= 2.") + spec["num_logits_to_keep"] = dflash_block_size if self.continuous_batching: spec["full_batch_size"] = kv_cache_batch_size diff --git a/QEfficient/transformers/models/qwen3/modeling_qwen3.py b/QEfficient/transformers/models/qwen3/modeling_qwen3.py index 8c6a9624e7..35b0841294 100644 --- a/QEfficient/transformers/models/qwen3/modeling_qwen3.py +++ b/QEfficient/transformers/models/qwen3/modeling_qwen3.py @@ -35,6 +35,7 @@ from QEfficient.transformers.cache_utils import QEffDynamicCache from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask from QEfficient.transformers.spd.dflash import compute_dflash_target_hidden_states +from QEfficient.transformers.spd.spd_transform_forward import filter_hidden_states from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE @@ -410,6 +411,7 @@ def forward( use_cache: bool | None = None, output_hidden_states: bool | None = None, cache_position: torch.LongTensor | None = None, + num_logits_to_keep: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs, ) -> tuple | CausalLMOutputWithPast: @@ -441,6 +443,8 @@ def forward( if getattr(self.model, "target_layer_ids", None): target_hidden = outputs.hidden_states hidden_states = outputs.last_hidden_state + if num_logits_to_keep is not None: + hidden_states = filter_hidden_states(hidden_states, position_ids, num_logits_to_keep) logits = self.lm_head(hidden_states).float() predicted_token_ids = logits.argmax(dim=-1).to(torch.int32) return CausalLMOutputWithPast( diff --git a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py index 2b46bbcf67..a599fb7704 100644 --- a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -58,6 +58,7 @@ resolve_layer_window, ) from QEfficient.transformers.spd.dflash import compute_dflash_target_hidden_states +from QEfficient.transformers.spd.spd_transform_forward import filter_hidden_states from QEfficient.utils import constants from QEfficient.utils._utils import IOInfo, get_padding_shape_from_config from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE @@ -1553,6 +1554,7 @@ def forward( use_cache: Optional[bool] = None, output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: torch.LongTensor | None = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs, ) -> CausalLMOutputWithPast: @@ -1573,6 +1575,9 @@ def forward( if getattr(self.model, "target_layer_ids", None): hidden_states = outputs.last_hidden_state + if num_logits_to_keep is not None: + text_position_ids = position_ids[0] if position_ids.ndim == 3 else position_ids + hidden_states = filter_hidden_states(hidden_states, text_position_ids, num_logits_to_keep) elif position_ids is None: hidden_states = outputs.last_hidden_state[:, -1:, :] else: @@ -1917,6 +1922,7 @@ def forward( batch_index: Optional[torch.LongTensor] = None, comp_ctx_lengths: Optional[List[int]] = None, cache_commit: torch.Tensor | None = None, + num_logits_to_keep: torch.LongTensor | None = None, ): # Continuous batching also supplies batch_index for prefill. Folded # attention is decode-only, so do not infer it from batch_index alone. @@ -1959,7 +1965,10 @@ def forward( ) image_idx = (indices1.max() + 1).unsqueeze(0).unsqueeze(0) if getattr(self.language_model, "target_layer_ids", None): - logits = self.model.lm_head(outputs.last_hidden_state).float() + hidden_states = outputs.last_hidden_state + if num_logits_to_keep is not None: + hidden_states = filter_hidden_states(hidden_states, position_ids[0], num_logits_to_keep) + logits = self.model.lm_head(hidden_states).float() return ( logits.argmax(dim=-1).to(torch.int32), vision_embeds, @@ -2096,6 +2105,9 @@ def get_specializations( vision_batch_size = batch_size if vision_batch_size is None else vision_batch_size language_model = getattr(self.model, "language_model", None) + if getattr(language_model, "target_layer_ids", None): + # Hybrid verification and replay retain the language model's chunk width. + compiler_options.pop("dflash_block_size", None) blocking_config = None for layer in getattr(language_model, "layers", ()): attention = getattr(layer, "self_attn", None) @@ -2210,15 +2222,26 @@ def _build_lang_spec(seq_len_val, comp_ctx_len=None): spec["full_batch_exec_size"] = full_batch_size return spec + decode_seq_len = chunk_size if getattr(language_model, "target_layer_ids", None) else 1 lang = [] if comp_ctx_lengths_prefill is not None or comp_ctx_lengths_decode is not None: for comp_ctx in comp_ctx_lengths_prefill or []: - lang.append(_build_lang_spec(prefill_seq_len, comp_ctx_len=comp_ctx)) + prefill_spec = _build_lang_spec(prefill_seq_len, comp_ctx_len=comp_ctx) + if getattr(language_model, "target_layer_ids", None): + prefill_spec["num_logits_to_keep"] = 1 + lang.append(prefill_spec) for comp_ctx in comp_ctx_lengths_decode or []: - lang.append(_build_lang_spec(1, comp_ctx_len=comp_ctx)) + decode_spec = _build_lang_spec(decode_seq_len, comp_ctx_len=comp_ctx) + if getattr(language_model, "target_layer_ids", None): + decode_spec["num_logits_to_keep"] = decode_seq_len + lang.append(decode_spec) else: - lang.append(_build_lang_spec(prefill_seq_len)) - lang.append(_build_lang_spec(1)) + prefill_spec = _build_lang_spec(prefill_seq_len) + decode_spec = _build_lang_spec(decode_seq_len) + if getattr(language_model, "target_layer_ids", None): + prefill_spec["num_logits_to_keep"] = 1 + decode_spec["num_logits_to_keep"] = decode_seq_len + lang = [prefill_spec, decode_spec] if kv_offload: return {"vision": vision, "lang": lang}, compiler_options diff --git a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index 65f39a674b..88fec29940 100644 --- a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -76,6 +76,7 @@ silu_glu_mlp, ) from QEfficient.transformers.spd.dflash import compute_dflash_target_hidden_states, dflash_chunk_gated_delta_scan +from QEfficient.transformers.spd.spd_transform_forward import filter_hidden_states from QEfficient.utils import constants from QEfficient.utils._utils import IOInfo, get_padding_shape_from_config from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE @@ -1597,6 +1598,7 @@ def forward( use_cache: Optional[bool] = None, output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: torch.LongTensor | None = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs, ) -> CausalLMOutputWithPast: @@ -1617,6 +1619,9 @@ def forward( if getattr(self.model, "target_layer_ids", None): hidden_states = outputs.last_hidden_state + if num_logits_to_keep is not None: + text_position_ids = position_ids[0] if position_ids.ndim == 3 else position_ids + hidden_states = filter_hidden_states(hidden_states, text_position_ids, num_logits_to_keep) elif position_ids is None: hidden_states = outputs.last_hidden_state[:, -1:, :] else: @@ -1961,6 +1966,7 @@ def forward( batch_index: Optional[torch.LongTensor] = None, comp_ctx_lengths: Optional[List[int]] = None, cache_commit: torch.Tensor | None = None, + num_logits_to_keep: torch.LongTensor | None = None, ): if getattr(self.language_model, "target_layer_ids", None) and is_layerwise_active(): raise ValueError("Qwen3.5-MoE DFlash requires the full target model; layerwise execution is unsupported.") @@ -2018,7 +2024,10 @@ def forward( cache_commit=cache_commit, ) if getattr(self.language_model, "target_layer_ids", None): - logits = self.model.lm_head(outputs.last_hidden_state).float() + hidden_states = outputs.last_hidden_state + if num_logits_to_keep is not None: + hidden_states = filter_hidden_states(hidden_states, position_ids[0], num_logits_to_keep) + logits = self.model.lm_head(hidden_states).float() image_idx = (indices1.max() + 1).unsqueeze(0).unsqueeze(0) return ( logits.argmax(dim=-1).to(torch.int32), @@ -2243,6 +2252,9 @@ def get_specializations( language_model = getattr(getattr(self, "model", None), "language_model", None) layers = getattr(language_model, "layers", ()) + if getattr(language_model, "target_layer_ids", None): + # Hybrid verification and replay retain the language model's chunk width. + compiler_options.pop("dflash_block_size", None) blocking_config = None for layer in layers: attention = getattr(layer, "self_attn", None) @@ -2363,12 +2375,22 @@ def _build_lang_spec(seq_len_val, comp_ctx_len=None): lang = [] if comp_ctx_lengths_prefill is not None or comp_ctx_lengths_decode is not None: for comp_ctx in comp_ctx_lengths_prefill or []: - lang.append(_build_lang_spec(prefill_seq_len, comp_ctx_len=comp_ctx)) + prefill_spec = _build_lang_spec(prefill_seq_len, comp_ctx_len=comp_ctx) + if getattr(language_model, "target_layer_ids", None): + prefill_spec["num_logits_to_keep"] = 1 + lang.append(prefill_spec) for comp_ctx in comp_ctx_lengths_decode or []: - lang.append(_build_lang_spec(decode_seq_len, comp_ctx_len=comp_ctx)) + decode_spec = _build_lang_spec(decode_seq_len, comp_ctx_len=comp_ctx) + if getattr(language_model, "target_layer_ids", None): + decode_spec["num_logits_to_keep"] = decode_seq_len + lang.append(decode_spec) else: - lang.append(_build_lang_spec(prefill_seq_len)) - lang.append(_build_lang_spec(decode_seq_len)) + prefill_spec = _build_lang_spec(prefill_seq_len) + decode_spec = _build_lang_spec(decode_seq_len) + if getattr(language_model, "target_layer_ids", None): + prefill_spec["num_logits_to_keep"] = 1 + decode_spec["num_logits_to_keep"] = decode_seq_len + lang = [prefill_spec, decode_spec] if kv_offload: return {"vision": vision, "lang": lang}, compiler_options diff --git a/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py index d6db4fcdb8..2c34605f8c 100644 --- a/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -40,6 +40,7 @@ from QEfficient.transformers.cache_utils import QEffDynamicCache from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask from QEfficient.transformers.spd.dflash import compute_dflash_target_hidden_states +from QEfficient.transformers.spd.spd_transform_forward import filter_hidden_states from QEfficient.utils import constants from QEfficient.utils._utils import IOInfo, get_padding_shape_from_config from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE @@ -740,6 +741,7 @@ def forward( slot_id: torch.LongTensor | None = None, batch_index: torch.LongTensor | None = None, comp_ctx_lengths: list[int] | None = None, + num_logits_to_keep: torch.LongTensor | None = None, ): inputs_embeds = self.model.get_input_embeddings()(input_ids) _B, _N, C = inputs_embeds.shape @@ -780,7 +782,10 @@ def forward( if getattr(self.language_model, "target_layer_ids", None): target_hidden = outputs.hidden_states - logits = self.model.lm_head(outputs.last_hidden_state).float() + hidden_states = outputs.last_hidden_state + if num_logits_to_keep is not None: + hidden_states = filter_hidden_states(hidden_states, position_ids[0], num_logits_to_keep) + logits = self.model.lm_head(hidden_states).float() predicted_token_ids = logits.argmax(dim=-1).to(torch.int32) return ( predicted_token_ids, @@ -1025,6 +1030,9 @@ def get_specializations( ): comp_ctx_lengths_prefill = compiler_options.pop("comp_ctx_lengths_prefill", None) comp_ctx_lengths_decode = compiler_options.pop("comp_ctx_lengths_decode", None) + dflash_block_size = None + if getattr(self.model.language_model, "target_layer_ids", None): + dflash_block_size = compiler_options.pop("dflash_block_size", None) # Preserve the legacy shape when callers do not request a separate vision batch. vision_batch_size = batch_size if vision_batch_size is None else vision_batch_size if height is None or width is None: @@ -1121,6 +1129,8 @@ def get_specializations( if full_batch_size: lang_prefill["full_batch_exec_size"] = full_batch_size + if dflash_block_size is not None: + lang_prefill["num_logits_to_keep"] = 1 lang.append(lang_prefill) for i in range(len(comp_ctx_lengths_decode or [])): @@ -1139,6 +1149,9 @@ def get_specializations( else: lang_decode["batch_size"] = kv_cache_batch_size + if dflash_block_size is not None: + lang_decode["num_logits_to_keep"] = dflash_block_size + lang_decode["seq_len"] = dflash_block_size lang.append(lang_decode) else: lang_prefill = { @@ -1187,6 +1200,10 @@ def get_specializations( lang_decode["total_num_kv_blocks"] = kv_cache_batch_size * num_kv_blocks lang_decode["kv_block_size"] = -(-ctx_len // num_kv_blocks) + if dflash_block_size is not None: + lang_prefill["num_logits_to_keep"] = 1 + lang_decode["num_logits_to_keep"] = dflash_block_size + lang_decode["seq_len"] = dflash_block_size lang = [lang_prefill, lang_decode] specializations = {} diff --git a/examples/performance/dflash/README.md b/examples/performance/dflash/README.md index 2e775efc19..7d3b8e0c5a 100644 --- a/examples/performance/dflash/README.md +++ b/examples/performance/dflash/README.md @@ -121,6 +121,13 @@ python basic_inference_vision.py --model_name gemma-4-31B-it \ ## Common flags +DFlash targets use the existing SpD `num_logits_to_keep` input to project only +the last valid position during prompt prefill. Verification still predicts every +block position, and all hidden states needed by the DLM are preserved. This +applies automatically to Qwen3, Llama, Gemma4, Qwen3-VL, Qwen3.5, and Qwen3.5-MoE. +Existing TLM QPCs remain usable but need re-export/compilation for this change; +DLM and vision QPCs are unchanged. + | Flag | Default | Notes | |---|---|---| | `--tlm_devices` | `0,1,2,3` | TLM device IDs | diff --git a/tests/transformers/spd/test_dflash_spd_inference.py b/tests/transformers/spd/test_dflash_spd_inference.py index fa3cfc506c..e7fd47763b 100644 --- a/tests/transformers/spd/test_dflash_spd_inference.py +++ b/tests/transformers/spd/test_dflash_spd_inference.py @@ -34,12 +34,7 @@ VOCAB_SIZE = 64 HIDDEN_SIZE = 64 BLOCK_SIZE = 4 # DFlash decode block size; also the DLM's prefill_seq_len -# TLM prompt-chunk size for prefill. Must differ from BLOCK_SIZE: the TLM's decode -# specialization seq_len is always overridden to BLOCK_SIZE (see -# build_decode_specialization's dflash_tlm branch), so a prefill spec with the same -# seq_len would be indistinguishable from it and the AIC compiler would reject the -# network ("No input that uniquely identifies specialization"). A multiple of -# BLOCK_SIZE keeps the prefill loop's sub-block splitting exact (no remainder). +# Different prefill and verification lengths exercise buffer resizing. PROMPT_CHUNK_SIZE = 2 * BLOCK_SIZE CTX_LEN = 32 PAD_TOKEN_ID = 0 diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index a531be93e9..953c3a0572 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -989,6 +989,9 @@ def test_qwen3_5_dflash_ort_recurrent_commit( if layer.layer_type == "linear_attention": layer.linear_attn.dflash_matrix_solve = matrix_solve inputs, output_names, dynamic_axes = _qwen_export_io(qeff, prefill_seq_len=8) + inputs["num_logits_to_keep"] = torch.zeros(1, 1, dtype=torch.int64) + dynamic_axes["num_logits_to_keep"] = {0: "num_logits_to_keep"} + dynamic_axes["logits"] = {0: "batch_size", 1: "num_logits_to_keep"} for i, kind in enumerate(qeff.model.config.text_config.layer_types): if kind == "full_attention": inputs["past_key_values"][i] = [ @@ -1025,6 +1028,8 @@ def test_qwen3_5_dflash_ort_recurrent_commit( positions[:, :, valid:] = -1 inputs["position_ids"] = positions inputs["cache_commit"] = torch.tensor([commit]) + num_logits = 1 if start == 0 else length + inputs["num_logits_to_keep"] = torch.zeros(num_logits, 1, dtype=torch.int64) inputs["past_key_values"] = cache feed = {name: value.copy() for name, value in _flatten_qwen_ort_inputs(inputs, session_inputs, decoder).items()} if start: From 25221869daded93601c0486a0b48fe257d461379 Mon Sep 17 00:00:00 2001 From: Vahid Janfaza Date: Wed, 14 Oct 2026 11:59:50 -0700 Subject: [PATCH 5/9] Fix variable-length DFlash prefill support for Qwen3.5 dense models Reuse the scripted gated-delta chunk scan for larger prefill lengths while keeping verification and accepted-prefix replay block-sized. Update compile/runtime validation and existing regression tests. Keep non-DFlash behavior unchanged. Co-authored-by: FNU Annanya Signed-off-by: Vahid Janfaza --- QEfficient/generation/dflash_generation.py | 16 ++---- .../transformers/models/modeling_auto.py | 9 ++-- .../models/qwen3_5/modeling_qwen3_5.py | 49 ++++++++++--------- examples/performance/dflash/README.md | 3 +- .../dflash/basic_inference_vision.py | 2 +- examples/performance/dflash/utils.py | 19 ++----- .../spd/test_dflash_spd_inference.py | 30 +++++++----- .../unit_test/models/test_model_quickcheck.py | 10 ++-- .../transforms/test_dflash_transform.py | 35 ++++++++++--- 9 files changed, 96 insertions(+), 77 deletions(-) diff --git a/QEfficient/generation/dflash_generation.py b/QEfficient/generation/dflash_generation.py index 4f481996c7..eee44f28a8 100644 --- a/QEfficient/generation/dflash_generation.py +++ b/QEfficient/generation/dflash_generation.py @@ -23,7 +23,6 @@ from QEfficient.generation.cloud_infer import QAICInferenceSession from QEfficient.transformers.models.qwen3_5.modeling_qwen3_5 import ( - QWEN3_5_DFLASH_CHUNK_SIZE, QEffQwen3_5ForConditionalGeneration, ) from QEfficient.transformers.models.qwen3_vl.modeling_qwen3_vl import ( @@ -691,7 +690,7 @@ def get_qwen35_target_seq_len(tlm_session: QAICInferenceSession, block_size: int if block_size in lengths: return block_size raise ValueError( - f"Qwen3.5-MoE DFlash requires a {block_size}-token verification specialization " + f"Qwen3.5 DFlash requires a {block_size}-token verification specialization " f"(found {sorted(lengths)}). Recompile the TLM with dflash_block_size={block_size}." ) @@ -721,17 +720,10 @@ def run_spd_inference_qwen3_vl( is_qwen35 = tlm_config.model_type in {"qwen3_5", "qwen3_5_moe"} if is_qwen35 and (batch_size != 1 or "cache_commit" not in tlm_session.input_names): raise ValueError("Qwen3.5 DFlash requires batch size 1 and a TLM QPC with cache_commit input.") - if tlm_config.model_type == "qwen3_5" and ( - prompt_chunk_size not in (QWEN3_5_DFLASH_CHUNK_SIZE, 64) or not 1 < block_size <= prompt_chunk_size - ): - raise ValueError( - "Qwen3.5 DFlash requires 16-token target calls (or a legacy 64-token QPC) and a block that fits; " - "rebuild with --prefill_seq_len 16." - ) - target_seq_len = prompt_chunk_size if is_qwen35 else block_size - if tlm_config.model_type == "qwen3_5_moe": + target_seq_len = block_size + if is_qwen35: if block_size < 2 or prompt_chunk_size < block_size or prompt_chunk_size % block_size: - raise ValueError("Qwen3.5-MoE DFlash requires prefill length to be a positive multiple of block_size >= 2.") + raise ValueError("Qwen3.5 DFlash requires prefill length to be a positive multiple of block_size >= 2.") target_seq_len = get_qwen35_target_seq_len(tlm_session, block_size) metrics = SpecDecodingMetrics(block_size=block_size) image_processing_start = time.time() diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 54312b1e0f..5982a7a09b 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -1597,6 +1597,9 @@ def __init__( if self.dflash_tlm and self.model.config.model_type == "qwen3_5_moe": self.lang_model.hash_params["dflash_cache_protocol"] = "qwen35_moe_dynamic_replay_v1" self.lang_model.hash_params["dflash_chunk_size"] = 16 + if self.dflash_tlm and self.model.config.model_type == "qwen3_5": + self.lang_model.hash_params["dflash_cache_protocol"] = "qwen35_dynamic_replay_v1" + self.lang_model.hash_params["dflash_chunk_size"] = 16 # Note: SamplerTransform should be applied after all other transforms # are done. The role of the sampler is to just add nodes at the output of the @@ -2108,12 +2111,12 @@ def compile( full_batch_size=full_batch_size, ) - if self.dflash_tlm and self.model.config.model_type == "qwen3_5_moe": + if self.dflash_tlm and self.model.config.model_type in {"qwen3_5", "qwen3_5_moe"}: language_model = self.lang_model.model.language_model previous_chunk_size = getattr(language_model, "dflash_chunk_size", 16) chunk_size = previous_chunk_size if dflash_block_size is None else dflash_block_size if not isinstance(chunk_size, int) or chunk_size < 2 or chunk_size & (chunk_size - 1): - raise ValueError("Qwen3.5-MoE DFlash requires a power-of-two dflash_block_size >= 2.") + raise ValueError("Qwen3.5 DFlash requires a power-of-two dflash_block_size >= 2.") if chunk_size != previous_chunk_size and self.lang_model.onnx_path is not None and lang_onnx_path is None: raise ValueError("Changing dflash_block_size requires a fresh model instance to re-export the TLM.") language_model.dflash_chunk_size = chunk_size @@ -2355,7 +2358,7 @@ def compile( # but one CCL value is silently dropped. lang_specs = specializations["lang"] if prefill_only: - if self.dflash_tlm and self.model.config.model_type == "qwen3_5_moe": + if self.dflash_tlm and self.model.config.model_type in {"qwen3_5", "qwen3_5_moe"}: # Prefill and block-sized verification share the chunk graph. specializations = [] for spec in lang_specs: diff --git a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py index a599fb7704..276507b303 100644 --- a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -57,7 +57,7 @@ is_layerwise_active, resolve_layer_window, ) -from QEfficient.transformers.spd.dflash import compute_dflash_target_hidden_states +from QEfficient.transformers.spd.dflash import compute_dflash_target_hidden_states, dflash_chunk_gated_delta_scan from QEfficient.transformers.spd.spd_transform_forward import filter_hidden_states from QEfficient.utils import constants from QEfficient.utils._utils import IOInfo, get_padding_shape_from_config @@ -844,6 +844,7 @@ def torch_chunk_gated_delta_rule_qeff( ones_lower=None, eye=None, use_matrix_solve=False, + use_dynamic_chunk_loop=False, ): initial_dtype = query.dtype # if use_qk_l2norm_in_kernel: @@ -928,21 +929,25 @@ def torch_chunk_gated_delta_rule_qeff( if initial_state is None else initial_state.to(value) ) - core_attn_out = torch.zeros_like(value) - mask = mask_strict - - # for each chunk - for i in range(total_sequence_length // chunk_size): - q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] - attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) - v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state - v_new = v_i - v_prime - attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - core_attn_out[:, :, i] = attn_inter + attn @ v_new - last_recurrent_state = ( - last_recurrent_state * g[:, :, i, -1, None, None].exp() - + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new + if use_dynamic_chunk_loop: + core_attn_out, last_recurrent_state = dflash_chunk_gated_delta_scan( + query, key, value, k_cumdecay, g, decay_mask, mask_strict, last_recurrent_state ) + else: + core_attn_out = torch.zeros_like(value) + mask = mask_strict + + for i in range(total_sequence_length // chunk_size): + q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) + v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state + v_new = v_i - v_prime + attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state + core_attn_out[:, :, i] = attn_inter + attn @ v_new + last_recurrent_state = ( + last_recurrent_state * g[:, :, i, -1, None, None].exp() + + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new + ) if not output_final_state: last_recurrent_state = None @@ -1123,6 +1128,7 @@ def forward( ones_lower=self._ones_lower[:dflash_chunk_size, :dflash_chunk_size], eye=self._eye[:dflash_chunk_size, :dflash_chunk_size], use_matrix_solve=getattr(self, "dflash_matrix_solve", True), + use_dynamic_chunk_loop=True, ) elif batch_fold: # Folded decode is exported with a static one-token sequence. @@ -2136,8 +2142,8 @@ def get_specializations( prefill_seq_len = prefill_seq_len if prefill_seq_len else 128 if getattr(language_model, "target_layer_ids", None): chunk_size = getattr(language_model, "dflash_chunk_size", QWEN3_5_DFLASH_CHUNK_SIZE) - if prefill_seq_len != chunk_size: - raise ValueError(f"Qwen3.5 DFlash requires prefill_seq_len={chunk_size} for its gated-delta chunk.") + if prefill_seq_len < chunk_size or prefill_seq_len % chunk_size: + raise ValueError(f"Qwen3.5 DFlash requires prefill_seq_len to be a positive multiple of {chunk_size}.") ctx_len = ctx_len if ctx_len else constants.INTERN_CTX_LEN kv_cache_batch_size = kv_cache_batch_size or full_batch_size or batch_size channel = 3 @@ -2315,12 +2321,9 @@ def get_dummy_inputs( if is_dflash: if not kv_offload or continuous_batching: raise ValueError("Qwen3.5 DFlash requires kv_offload=True and continuous_batching=False.") - if not 1 < dummy_seq_len <= 64: - raise ValueError( - "Qwen3.5 DFlash requires 2 <= prefill_seq_len <= 64 for the exported gated-delta graph." - ) - # The shared exporter supplies a 32-token dummy. Trace only one - # chunk: Python chunk loops do not generalize across chunk counts. + if dummy_seq_len < 2: + raise ValueError("Qwen3.5 DFlash requires prefill_seq_len >= 2 for the exported gated-delta graph.") + # Keep tracing small; the scripted scan preserves the dynamic chunk count. chunk_size = getattr(self.model.language_model, "dflash_chunk_size", QWEN3_5_DFLASH_CHUNK_SIZE) dummy_seq_len = min(dummy_seq_len, chunk_size) bs = kwargs.get("batch_size", constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE) diff --git a/examples/performance/dflash/README.md b/examples/performance/dflash/README.md index 7d3b8e0c5a..5354ba882b 100644 --- a/examples/performance/dflash/README.md +++ b/examples/performance/dflash/README.md @@ -83,11 +83,12 @@ Qwen3.5 uses the VLM entry point with its DFlash target and draft checkpoints: ```bash python basic_inference_vision.py --model_name Qwen3.5-4B \ --tlm_devices 40,41,42,43 --dlm_devices 44,45,46,47 --vision_devices 48,49,50,51 \ - --ctx_len 4096 --prefill_seq_len 16 \ + --ctx_len 4096 --prefill_seq_len 128 \ --image --image_prompt "Describe this image in detail." ``` Use `--prompt "..."` instead of `--image --image_prompt ...` for text-only input. +Prefill length can be any positive multiple of the draft block size; verification and replay remain block-sized. --- diff --git a/examples/performance/dflash/basic_inference_vision.py b/examples/performance/dflash/basic_inference_vision.py index 49ce8dfea6..b5ab87e942 100644 --- a/examples/performance/dflash/basic_inference_vision.py +++ b/examples/performance/dflash/basic_inference_vision.py @@ -297,7 +297,7 @@ def main(): dlm_session, tlm_session = load_spd_sessions(tlm_qpc, dlm_qpc, args.tlm_devices, args.dlm_devices) prompt_chunk_size = get_spd_prompt_chunk_size(tlm_session) - target_decode_length = prompt_chunk_size if tlm_config.model_type == "qwen3_5" else get_dflash_block_size(config) + target_decode_length = get_dflash_block_size(config) validate_spd_decode_specialization(tlm_session, target_decode_length) if is_qwen3vl: diff --git a/examples/performance/dflash/utils.py b/examples/performance/dflash/utils.py index a12177a2ff..cb7e481943 100644 --- a/examples/performance/dflash/utils.py +++ b/examples/performance/dflash/utils.py @@ -21,7 +21,6 @@ QAICInferenceSession, is_retained_state_name, ) -from QEfficient.transformers.models.qwen3_5.modeling_qwen3_5 import QWEN3_5_DFLASH_CHUNK_SIZE from QEfficient.utils.logging_utils import QEFFLogger logger = QEFFLogger.get_logger("INFRA") @@ -376,19 +375,11 @@ def compile_qwen3vl_vlm_qpcs( is_qwen35 = config.model_type == "qwen3_5" is_qwen35_moe = config.model_type == "qwen3_5_moe" qaic_config = {"target_layer_ids": tlm_target_ids} - if is_qwen35_moe: + if is_qwen35 or is_qwen35_moe: if block_size < 2 or block_size & (block_size - 1): - raise ValueError("Qwen3.5-MoE DFlash requires a power-of-two draft block_size >= 2.") + raise ValueError("Qwen3.5 DFlash requires a power-of-two draft block_size >= 2.") if prefill_seq_len < block_size or prefill_seq_len % block_size: - raise ValueError(f"Qwen3.5-MoE DFlash requires prefill_seq_len to be a positive multiple of {block_size}.") - if is_qwen35: - if prefill_seq_len != QWEN3_5_DFLASH_CHUNK_SIZE: - logger.warning( - f"Qwen3.5 DFlash uses prefill_seq_len={QWEN3_5_DFLASH_CHUNK_SIZE} for its fixed gated-delta chunk." - ) - prefill_seq_len = QWEN3_5_DFLASH_CHUNK_SIZE - if block_size != QWEN3_5_DFLASH_CHUNK_SIZE: - raise ValueError(f"Qwen3.5 DFlash currently requires block_size={QWEN3_5_DFLASH_CHUNK_SIZE}.") + raise ValueError(f"Qwen3.5 DFlash requires prefill_seq_len to be a positive multiple of {block_size}.") logger.warning("Qwen3.5 DFlash uses accepted-prefix replay; target timing includes the extra commit pass.") tlm_qeff = QEFFAutoModelForImageTextToText.from_pretrained( tlm_repo, @@ -412,10 +403,6 @@ def compile_qwen3vl_vlm_qpcs( text_model.fc.weight.data.div_(s) text_model.hidden_norm.weight.data.copy_(state_dict["hidden_norm.weight"].to(torch.float32)) logger.info(f"[compile_qwen3vl_tlm] fc/hidden_norm injected (fc scale s={s:.6f})") - if is_qwen35: - # The fixed chunk and matrix solve change the graph, not its I/O names. - tlm_qeff.lang_model.hash_params["dflash_cache_protocol"] = "qwen35_chunk16_matrix_replay_v3" - compile_kwargs = { "prefill_seq_len": prefill_seq_len, "ctx_len": ctx_len, diff --git a/tests/transformers/spd/test_dflash_spd_inference.py b/tests/transformers/spd/test_dflash_spd_inference.py index e7fd47763b..2cfe4f3da9 100644 --- a/tests/transformers/spd/test_dflash_spd_inference.py +++ b/tests/transformers/spd/test_dflash_spd_inference.py @@ -139,9 +139,9 @@ def test_dflash_decode_specialization_accepts_fixed_and_dynamic_shapes(sequence_ validate_spd_decode_specialization(session, 32) -@pytest.mark.parametrize("prompt_chunk_size", [8, 32, 128]) +@pytest.mark.parametrize("prompt_chunk_size", [8, 31, 127]) def test_qwen3_5_runtime_rejects_unsupported_prefill_qpc(prompt_chunk_size): - with pytest.raises(ValueError, match="rebuild with --prefill_seq_len 16"): + with pytest.raises(ValueError, match="positive multiple of block_size"): run_spd_inference_qwen3_vl( prompt_text="test", tokenizer=None, @@ -159,7 +159,9 @@ def test_qwen3_5_runtime_rejects_unsupported_prefill_qpc(prompt_chunk_size): "model_type,prompt_chunk_size,target_seq_len", [ ("qwen3_5", 16, 16), - ("qwen3_5", 64, 64), + ("qwen3_5", 64, 16), + ("qwen3_5", 128, 16), + ("qwen3_5", 256, 16), ("qwen3_5_moe", 16, 16), ("qwen3_5_moe", 128, 16), ], @@ -181,7 +183,7 @@ def run(feeds): } session = SimpleNamespace( - input_names=["cache_commit", "vision_embeds"], + input_names=["cache_commit", "vision_embeds", "num_logits_to_keep"], binding_index_map={"vision_embeds": 0, "input_ids": 1}, bindings=[SimpleNamespace(type=0, dims=(1, 1, 32)), SimpleNamespace(dims=(1, prompt_chunk_size))], allowed_shapes=[[(0, [1, 1, 32]), (8, [1, length])] for length in {prompt_chunk_size, target_seq_len}], @@ -205,6 +207,7 @@ def core(**kwargs): kwargs["tlm_prefill"](0, prompt_chunk_size) assert calls[-1]["input_ids"].shape[1] == prompt_chunk_size assert calls[-1]["cache_commit"].item() == 1 + assert calls[-1]["num_logits_to_keep"].shape == (1, 1) kwargs["prepare_decode"]() ids = np.arange(16, dtype=np.int64).reshape(1, 16) positions = np.arange(5, 21, dtype=np.int64).reshape(1, 16) @@ -216,9 +219,11 @@ def core(**kwargs): np.testing.assert_array_equal(verify["position_ids"][0, :, :16], positions) assert np.all(verify["position_ids"][..., 16:] == -1) assert verify["cache_commit"].item() == 0 + assert verify["num_logits_to_keep"].shape == (16, 1) kwargs["tlm_commit"](ids, positions, 2) commit = calls[-1] assert commit["cache_commit"].item() == 1 + assert commit["num_logits_to_keep"].shape == (16, 1) np.testing.assert_array_equal(commit["position_ids"][0, :, :2], positions[:, :2]) assert np.all(commit["position_ids"][..., 2:] == -1) return kwargs["metrics"] @@ -234,11 +239,11 @@ def core(**kwargs): mask_token_id=2, vocab_size=64, prompt_chunk_size=prompt_chunk_size, - ctx_len=256, + ctx_len=512, block_size=16, hidden_size=32, ) - assert [item["logits"].shape[1] for item in buffers if "logits" in item] == [prompt_chunk_size, target_seq_len] + assert [item["logits"].shape[1] for item in buffers if "logits" in item] == [1, target_seq_len] @pytest.mark.parametrize("with_image", [False, True]) @@ -261,7 +266,8 @@ def test_qwen35_moe_runtime_mrope_matches_hf(with_image): @pytest.mark.parametrize("block_size", [8, 16]) -def test_qwen35_moe_compile_preserves_dflash_metadata(monkeypatch, block_size): +@pytest.mark.parametrize("model_type", ["qwen3_5", "qwen3_5_moe"]) +def test_qwen35_compile_preserves_dflash_metadata(monkeypatch, block_size, model_type): import torch from examples.performance.dflash import utils @@ -281,7 +287,7 @@ def load_model(*args, **kwargs): lang_model=SimpleNamespace(model=SimpleNamespace(language_model=text_model)), compile=compile_model ) - monkeypatch.setattr(utils.AutoConfig, "from_pretrained", lambda *a, **k: SimpleNamespace(model_type="qwen3_5_moe")) + monkeypatch.setattr(utils.AutoConfig, "from_pretrained", lambda *a, **k: SimpleNamespace(model_type=model_type)) monkeypatch.setattr( utils, "read_dlm_meta", @@ -291,10 +297,12 @@ def load_model(*args, **kwargs): assert utils.compile_qwen3vl_vlm_qpcs( "target", "draft", ctx_len=8192, prefill_seq_len=128, num_cores=8, num_devices=4 ) == ("target.qpc", "vision.qpc") - assert calls[0]["qaic_config"] == qaic_config - assert calls[0]["prefill_only"] and calls[0]["enable_chunking"] + assert calls[0]["prefill_only"] assert calls[0]["prefill_seq_len"] == 128 and calls[0]["dflash_block_size"] == block_size - assert calls[0]["mxint8_kv_cache"] + if model_type == "qwen3_5_moe": + assert calls[0]["qaic_config"] == qaic_config + assert calls[0]["enable_chunking"] + assert calls[0]["mxint8_kv_cache"] class _FakeTokenizer: diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index 953c3a0572..b9967a94ea 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -969,6 +969,8 @@ def test_dflash_mixed_window_ort_parity(tmp_path): (False, 64, 64, 64), (False, 8, 16, 16), (False, 16, 16, 16), + (False, 128, 16, 16), + (False, 256, 16, 16), (True, 16, 16, 16), (True, 128, 16, 16), (True, 128, 8, 8), @@ -982,7 +984,8 @@ def test_qwen3_5_dflash_ort_recurrent_commit( torch.manual_seed(11) model = make_tiny_qwen3_5_vlm(moe=moe) - model.config.text_config.max_position_embeddings = 256 + ctx_len = max(256, prefill_length + verification_length) + model.config.text_config.max_position_embeddings = ctx_len qeff = QEFFAutoModelForImageTextToText(model, kv_offload=True, qaic_config={"target_layer_ids": [1, 3]}) qeff.lang_model.model.language_model.dflash_chunk_size = dflash_chunk_size for layer in qeff.lang_model.model.language_model.layers: @@ -995,7 +998,7 @@ def test_qwen3_5_dflash_ort_recurrent_commit( for i, kind in enumerate(qeff.model.config.text_config.layer_types): if kind == "full_attention": inputs["past_key_values"][i] = [ - tensor.new_zeros((tensor.shape[0], tensor.shape[1], 256, tensor.shape[3])) + tensor.new_zeros((tensor.shape[0], tensor.shape[1], ctx_len, tensor.shape[3])) for tensor in inputs["past_key_values"][i] ] qeff.lang_model.export( @@ -1023,7 +1026,8 @@ def test_qwen3_5_dflash_ort_recurrent_commit( (prompt_length + 2, verification_length, 1, 1), (prompt_length + 3, verification_length, verified_tokens, 0), ]: - inputs["input_ids"] = torch.full((batch_size, length), 10 + start, dtype=torch.int64) + token = (10 + start) % model.config.text_config.vocab_size + inputs["input_ids"] = torch.full((batch_size, length), token, dtype=torch.int64) positions = torch.arange(start, start + length).view(1, 1, -1).repeat(4, batch_size, 1) positions[:, :, valid:] = -1 inputs["position_ids"] = positions diff --git a/tests/unit_test/transforms/test_dflash_transform.py b/tests/unit_test/transforms/test_dflash_transform.py index 5ca55653cc..531093ddda 100644 --- a/tests/unit_test/transforms/test_dflash_transform.py +++ b/tests/unit_test/transforms/test_dflash_transform.py @@ -250,20 +250,21 @@ def test_qwen3_5_dflash_requires_explicit_commit(): qeff.lang_model.model(**inputs) -@pytest.mark.parametrize("prefill_seq_len", [1, 65, 128]) +@pytest.mark.parametrize("prefill_seq_len", [1, 15, 65]) def test_qwen3_5_dflash_rejects_unsupported_prefill(prefill_seq_len): from QEfficient import QEFFAutoModelForImageTextToText qeff = QEFFAutoModelForImageTextToText( make_tiny_qwen3_5_vlm(), kv_offload=True, qaic_config={"target_layer_ids": [1, 3]} ) - with pytest.raises(ValueError, match="prefill_seq_len=16"): + with pytest.raises(ValueError, match="positive multiple of 16"): qeff.model.get_specializations(batch_size=1, prefill_seq_len=prefill_seq_len, ctx_len=256, kv_offload=True) - with pytest.raises(ValueError, match="prefill_seq_len <= 64"): - qeff.model.get_dummy_inputs(kv_offload=True, prefill_seq_len=prefill_seq_len) + if prefill_seq_len < 2: + with pytest.raises(ValueError, match="prefill_seq_len >= 2"): + qeff.model.get_dummy_inputs(kv_offload=True, prefill_seq_len=prefill_seq_len) -@pytest.mark.parametrize("dummy_seq_len", [16, 32, 64]) +@pytest.mark.parametrize("dummy_seq_len", [16, 32, 128, 256, 1024]) def test_qwen3_5_dflash_export_traces_one_chunk(dummy_seq_len): from QEfficient import QEFFAutoModelForImageTextToText @@ -273,8 +274,28 @@ def test_qwen3_5_dflash_export_traces_one_chunk(dummy_seq_len): inputs = qeff.model.get_dummy_inputs(kv_offload=True, prefill_seq_len=dummy_seq_len)["lang"] assert inputs["input_ids"].shape[1] == 16 assert inputs["position_ids"].shape[-1] == 16 - with pytest.raises(ValueError, match="prefill_seq_len=16"): - qeff.model.get_specializations(batch_size=1, prefill_seq_len=64, ctx_len=256, kv_offload=True) + specs, _ = qeff.model.get_specializations( + batch_size=1, prefill_seq_len=dummy_seq_len, ctx_len=2048, kv_offload=True + ) + assert [(spec["seq_len"], spec["num_logits_to_keep"]) for spec in specs["lang"]] == [ + (dummy_seq_len, 1), + (16, 16), + ] + specs, _ = qeff.model.get_specializations( + batch_size=1, + prefill_seq_len=dummy_seq_len, + ctx_len=2048, + kv_offload=True, + comp_ctx_lengths_prefill=[1024, 2048], + comp_ctx_lengths_decode=[512, 1024, 2048], + ) + assert [(spec["seq_len"], spec["num_logits_to_keep"]) for spec in specs["lang"]] == [ + (dummy_seq_len, 1), + (dummy_seq_len, 1), + (16, 16), + (16, 16), + (16, 16), + ] @pytest.mark.parametrize("chunk_size", [0, 1, 3, 24]) From eac83d4a51c9873764049a7ea58c227fb52fe9b6 Mon Sep 17 00:00:00 2001 From: Vahid Janfaza Date: Wed, 14 Oct 2026 18:25:57 -0700 Subject: [PATCH 6/9] Enable DFlash VLM integration with vLLM Return vocabulary logits for Qwen3-VL, Qwen3.5, and Qwen3.5-MoE DFlash targets, and reserve Gemma4 sliding-cache headroom for block verification. Load draft embeddings from the VLM target's nested language model and distinguish these exports in the artifact cache. These changes are limited to DFlash paths; existing non-DFlash outputs and cache behavior are unchanged. Co-authored-by: FNU Annanya Signed-off-by: Vahid Janfaza --- .../models/gemma4/modeling_gemma4.py | 16 +++++++++++++--- QEfficient/transformers/models/modeling_auto.py | 8 ++++++++ .../transformers/models/pytorch_transforms.py | 7 ++++--- .../models/qwen3_5/modeling_qwen3_5.py | 8 +++++++- .../models/qwen3_5_moe/modeling_qwen3_5_moe.py | 8 +++++++- .../models/qwen3_vl/modeling_qwen3_vl.py | 5 +++-- 6 files changed, 42 insertions(+), 10 deletions(-) diff --git a/QEfficient/transformers/models/gemma4/modeling_gemma4.py b/QEfficient/transformers/models/gemma4/modeling_gemma4.py index 5f2220e632..3eb724e20b 100644 --- a/QEfficient/transformers/models/gemma4/modeling_gemma4.py +++ b/QEfficient/transformers/models/gemma4/modeling_gemma4.py @@ -743,6 +743,12 @@ def forward( elif use_cache and past_key_values is None: past_key_values = QEffGemma4DynamicCache(config=self.config) + if getattr(self, "target_layer_ids", None) and past_key_values is not None: + # DFlash reserves block headroom; the attention window stays unchanged. + for layer in past_key_values.layers: + if layer.is_sliding and layer.keys is not None: + layer.sliding_window = layer.keys.shape[2] + if position_ids is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens @@ -1395,6 +1401,9 @@ def get_specializations( dflash_block_size = None if getattr(self.model.language_model, "target_layer_ids", None): dflash_block_size = compiler_options.pop("dflash_block_size", None) + sliding_cache_len = self.model.language_model.config.sliding_window + if dflash_block_size is not None: + sliding_cache_len = min(ctx_len, sliding_cache_len + max(prefill_seq_len, dflash_block_size) - 1) user_vision_size = compiler_options.pop("vision_size", None) if user_vision_size: if user_vision_size >= ctx_len: @@ -1410,7 +1419,7 @@ def build_lang_prefill_spec(comp_ctx_lengths: int | None = None): "batch_size": 1 if continuous_batching else batch_size, "seq_len": prefill_seq_len, "ctx_len": ctx_len, - "sliding_window": self.model.language_model.config.sliding_window, + "sliding_window": sliding_cache_len, "vision_batch_size": vision_batch_size, "vision_size": vision_size, } @@ -1431,7 +1440,7 @@ def build_lang_decode_spec(comp_ctx_lengths: int | None = None): "batch_size": full_batch_size if continuous_batching else batch_size, "seq_len": dflash_block_size or "1", "ctx_len": ctx_len, - "sliding_window": self.model.language_model.config.sliding_window, + "sliding_window": sliding_cache_len, "vision_batch_size": vision_batch_size, "vision_size": vision_size, } @@ -1499,12 +1508,13 @@ def get_output_names(self, kv_offload: bool = False): return lang_output_names def get_dummy_pkv_cache(self, config, batch_size, seq_len): + dflash_tlm = bool(getattr(self.model.language_model, "target_layer_ids", None)) past_key_values = [] for i, layer_type in enumerate(config.layer_types): if layer_type == "sliding_attention": n_heads = config.num_key_value_heads d_head = config.head_dim - layer_seq_len = min(config.sliding_window, seq_len) + layer_seq_len = seq_len if dflash_tlm else min(config.sliding_window, seq_len) else: use_alternative_attention = getattr(config, "attention_k_eq_v", False) n_heads = ( diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 5982a7a09b..ac3b98fcbb 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -1594,12 +1594,17 @@ def __init__( self.model, _ = DFlashTLMTransform.apply(self.model, qaic_config) # ---Sampling--- + if self.dflash_tlm and qaic_config.get("speculative_model_type") == "target": + self.lang_model.hash_params["dflash_output_logits"] = True + self.lang_model.hash_params["dflash_dlm_repo"] = qaic_config.get("dflash_dlm_repo") if self.dflash_tlm and self.model.config.model_type == "qwen3_5_moe": self.lang_model.hash_params["dflash_cache_protocol"] = "qwen35_moe_dynamic_replay_v1" self.lang_model.hash_params["dflash_chunk_size"] = 16 if self.dflash_tlm and self.model.config.model_type == "qwen3_5": self.lang_model.hash_params["dflash_cache_protocol"] = "qwen35_dynamic_replay_v1" self.lang_model.hash_params["dflash_chunk_size"] = 16 + if self.dflash_tlm and self.model.config.model_type == "gemma4": + self.lang_model.hash_params["dflash_cache_protocol"] = "gemma4_block_window_v1" # Note: SamplerTransform should be applied after all other transforms # are done. The role of the sampler is to just add nodes at the output of the @@ -3836,6 +3841,9 @@ def __init__( self.model, _ = DFlashTransform.apply(self.model, qaic_config) self.model, _ = DFlashDLMTransform.apply(self.model, qaic_config) self.hash_params["dflash_attention"] = "per_layer_window_v1" + if qaic_config.get("dflash_tlm_embedding_key"): + self.hash_params["dflash_tlm_embedding_key"] = qaic_config["dflash_tlm_embedding_key"] + self.hash_params["dflash_tlm_repo"] = qaic_config["dflash_tlm_repo"] if self.dflash_tlm: self.model, _ = DFlashTLMTransform.apply(self.model, qaic_config) diff --git a/QEfficient/transformers/models/pytorch_transforms.py b/QEfficient/transformers/models/pytorch_transforms.py index 6457964d7e..73a9376628 100644 --- a/QEfficient/transformers/models/pytorch_transforms.py +++ b/QEfficient/transformers/models/pytorch_transforms.py @@ -1274,11 +1274,12 @@ def apply(cls, model: nn.Module, qaic_config: dict | None = None, **kwargs) -> t if not tlm_repo: return model, True - w = load_checkpoint_weights(tlm_repo, {"lm_head.weight", "lm_head.bias", "model.embed_tokens.weight"}) - embed_w = w.get("model.embed_tokens.weight") + embed_key = qaic_config.get("dflash_tlm_embedding_key", "model.embed_tokens.weight") + w = load_checkpoint_weights(tlm_repo, {"lm_head.weight", "lm_head.bias", embed_key}) + embed_w = w.get(embed_key) lm_head_w = w.get("lm_head.weight", embed_w) # tie_word_embeddings: lm_head is a view of embed_tokens if lm_head_w is None or embed_w is None: - return model, True + raise ValueError(f"DFlash could not load target embeddings from {tlm_repo}: missing {embed_key}.") with torch.no_grad(): model.lm_head.weight.data.copy_(lm_head_w.float()) diff --git a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py index 276507b303..7f9f3029cd 100644 --- a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1975,8 +1975,14 @@ def forward( if num_logits_to_keep is not None: hidden_states = filter_hidden_states(hidden_states, position_ids[0], num_logits_to_keep) logits = self.model.lm_head(hidden_states).float() + qaic_config = getattr(self, "qaic_config", None) or {} + predictions = ( + logits + if qaic_config.get("speculative_model_type") == "target" + else logits.argmax(dim=-1).to(torch.int32) + ) return ( - logits.argmax(dim=-1).to(torch.int32), + predictions, vision_embeds, image_idx, outputs.hidden_states, diff --git a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index 88fec29940..0c12fc9ee7 100644 --- a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -2029,8 +2029,14 @@ def forward( hidden_states = filter_hidden_states(hidden_states, position_ids[0], num_logits_to_keep) logits = self.model.lm_head(hidden_states).float() image_idx = (indices1.max() + 1).unsqueeze(0).unsqueeze(0) + qaic_config = getattr(self, "qaic_config", None) or {} + predictions = ( + logits + if qaic_config.get("speculative_model_type") == "target" + else logits.argmax(dim=-1).to(torch.int32) + ) return ( - logits.argmax(dim=-1).to(torch.int32), + predictions, vision_embeds, image_idx, outputs.hidden_states, diff --git a/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py index 2c34605f8c..fbf84f405e 100644 --- a/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -786,9 +786,10 @@ def forward( if num_logits_to_keep is not None: hidden_states = filter_hidden_states(hidden_states, position_ids[0], num_logits_to_keep) logits = self.model.lm_head(hidden_states).float() - predicted_token_ids = logits.argmax(dim=-1).to(torch.int32) + output_logits = (getattr(self, "qaic_config", None) or {}).get("speculative_model_type") == "target" + predictions = logits if output_logits else logits.argmax(dim=-1).to(torch.int32) return ( - predicted_token_ids, + predictions, vision_embeds, deepstack_features, image_idx, From f2c9447285c30bd5a53bc79e177582994ccf7135 Mon Sep 17 00:00:00 2001 From: Vahid Janfaza Date: Mon, 19 Oct 2026 23:37:29 -0700 Subject: [PATCH 7/9] Address review comments and add GPT-OSS DFlash support Use baseline expert parallelism for GPT-OSS DFlash, with last-token prefill logits and block-sized verification. Fix target hidden-state normalization and keep sliding-window cache changes scoped to DFlash. Co-authored-by: FNU Annanya Signed-off-by: Vahid Janfaza --- QEfficient/generation/dflash_generation.py | 71 ++++- QEfficient/generation/dflash_tlm_handoff.py | 45 +++ QEfficient/transformers/cache_utils.py | 13 +- .../models/gpt_oss/modeling_gpt_oss.py | 44 ++- .../transformers/models/modeling_auto.py | 5 + .../transformers/models/pytorch_transforms.py | 8 + examples/performance/dflash/README.md | 44 ++- .../dflash/basic_inference_text.py | 37 ++- .../dflash/basic_inference_vision.py | 38 ++- examples/performance/dflash/utils.py | 101 ++++-- .../unit_test/models/test_model_quickcheck.py | 292 ++++++++++++++++++ 11 files changed, 655 insertions(+), 43 deletions(-) create mode 100644 QEfficient/generation/dflash_tlm_handoff.py diff --git a/QEfficient/generation/dflash_generation.py b/QEfficient/generation/dflash_generation.py index eee44f28a8..826e3f8835 100644 --- a/QEfficient/generation/dflash_generation.py +++ b/QEfficient/generation/dflash_generation.py @@ -22,6 +22,7 @@ from qwen_vl_utils import process_vision_info from QEfficient.generation.cloud_infer import QAICInferenceSession +from QEfficient.generation.dflash_tlm_handoff import DFlashTlmHandoff from QEfficient.transformers.models.qwen3_5.modeling_qwen3_5 import ( QEffQwen3_5ForConditionalGeneration, ) @@ -312,25 +313,35 @@ def run_spd_inference_single( hidden_size: int = 4096, generation_len: int = 256, batch_size: int = 1, + tlm_prefill_session: QAICInferenceSession | None = None, ) -> SpecDecodingMetrics: input_ids, position_ids = _prepare_text_inputs(tokenizer, prompt_text, prompt_chunk_size) generated_ids = _initialize_generated_ids(batch_size, ctx_len, input_ids.shape[1], tokenizer.pad_token_id) metrics = SpecDecodingMetrics(block_size=block_size) - prefill_selection = _dflash_logit_selection_inputs(tlm_session, 1) + prefill_session = tlm_prefill_session or tlm_session + handoff = DFlashTlmHandoff(prefill_session, tlm_session) if tlm_prefill_session is not None else None + prefill_selection = _dflash_logit_selection_inputs(prefill_session, 1) decode_selection = _dflash_logit_selection_inputs(tlm_session, block_size) - prefill_logits = 1 if prefill_selection else prompt_chunk_size - tlm_session.set_buffers({"logits": np.zeros((batch_size, prefill_logits), dtype=np.int32)}) - tlm_session.set_buffers({"hidden_states": np.zeros((batch_size, prompt_chunk_size, hidden_size), dtype=np.float32)}) + prefill_logits = 1 if prefill_selection or tlm_prefill_session is not None else prompt_chunk_size + prefill_session.set_buffers({"logits": np.zeros((batch_size, prefill_logits), dtype=np.int32)}) + prefill_session.set_buffers( + {"hidden_states": np.zeros((batch_size, prompt_chunk_size, hidden_size), dtype=np.float32)} + ) def tlm_prefill(chunk_start, chunk_end): - outputs = tlm_session.run( + last_chunk = chunk_end == input_ids.shape[1] + outputs = prefill_session.run( { "input_ids": input_ids[:, chunk_start:chunk_end], "position_ids": position_ids[:, chunk_start:chunk_end], **prefill_selection, + **(handoff.outputs if handoff is not None and last_chunk else {}), } ) + if handoff is not None and last_chunk: + handoff.capture(outputs) + prefill_session.skip_buffers(handoff.outputs) return outputs["logits"], outputs["hidden_states"] def prepare_decode(): @@ -338,7 +349,18 @@ def prepare_decode(): tlm_session.set_buffers({"hidden_states": np.zeros((batch_size, block_size, hidden_size), dtype=np.float32)}) def tlm_decode(draft_ids, decode_position_ids): - outputs = tlm_session.run({"input_ids": draft_ids, "position_ids": decode_position_ids, **decode_selection}) + nonlocal handoff + outputs = tlm_session.run( + { + "input_ids": draft_ids, + "position_ids": decode_position_ids, + **decode_selection, + **(handoff.inputs if handoff is not None else {}), + } + ) + if handoff is not None: + tlm_session.skip_buffers(handoff.inputs) + handoff = None return outputs["logits"], outputs["hidden_states"] eos_token_ids = {tokenizer.eos_token_id} if tokenizer.eos_token_id is not None else set() @@ -715,6 +737,7 @@ def run_spd_inference_qwen3_vl( compiled_height: int = 354, compiled_width: int = 536, batch_size: int = 1, + tlm_prefill_session: QAICInferenceSession | None = None, ) -> SpecDecodingMetrics: """Run Qwen3-VL, Qwen3.5 or Qwen3.5-MoE DFlash for a text or image prompt.""" is_qwen35 = tlm_config.model_type in {"qwen3_5", "qwen3_5_moe"} @@ -726,6 +749,8 @@ def run_spd_inference_qwen3_vl( raise ValueError("Qwen3.5 DFlash requires prefill length to be a positive multiple of block_size >= 2.") target_seq_len = get_qwen35_target_seq_len(tlm_session, block_size) metrics = SpecDecodingMetrics(block_size=block_size) + prefill_session = tlm_prefill_session or tlm_session + handoff = DFlashTlmHandoff(prefill_session, tlm_session) if tlm_prefill_session is not None else None image_processing_start = time.time() if image is not None: @@ -789,25 +814,34 @@ def run_spd_inference_qwen3_vl( } if has_deepstack: vision_outputs["deepstack_features"] = deepstack_features.astype(deepstack_dtype) - tlm_session.set_buffers(vision_outputs) + prefill_session.set_buffers(vision_outputs) + if handoff is not None: + tlm_session.set_buffers(vision_outputs) lang_extra = {"image_idx": np.array([[0]], dtype=np.int64)} if is_qwen35: lang_extra["cache_commit"] = np.ones(1, dtype=np.int64) - prefill_selection = _dflash_logit_selection_inputs(tlm_session, 1) + prefill_selection = _dflash_logit_selection_inputs(prefill_session, 1) decode_selection = _dflash_logit_selection_inputs(tlm_session, target_seq_len) - prefill_logits = 1 if prefill_selection else prompt_chunk_size - tlm_session.set_buffers({"logits": np.zeros((batch_size, prefill_logits), dtype=np.int32)}) - tlm_session.set_buffers({"hidden_states": np.zeros((batch_size, prompt_chunk_size, hidden_size), dtype=np.float32)}) + prefill_logits = 1 if prefill_selection or tlm_prefill_session is not None else prompt_chunk_size + prefill_session.set_buffers({"logits": np.zeros((batch_size, prefill_logits), dtype=np.int32)}) + prefill_session.set_buffers( + {"hidden_states": np.zeros((batch_size, prompt_chunk_size, hidden_size), dtype=np.float32)} + ) def tlm_prefill(chunk_start, chunk_end): - outputs = tlm_session.run( + last_chunk = chunk_end == padded_len + outputs = prefill_session.run( { "input_ids": input_ids_np[:, chunk_start:chunk_end], "position_ids": position_ids_np[:, :, chunk_start:chunk_end], **prefill_selection, **lang_extra, + **(handoff.outputs if handoff is not None and last_chunk else {}), } ) + if handoff is not None and last_chunk: + handoff.capture(outputs) + prefill_session.skip_buffers(handoff.outputs) lang_extra["image_idx"] = outputs["image_idx_output"] return outputs["logits"], outputs["hidden_states"] @@ -816,7 +850,8 @@ def prepare_decode(): tlm_session.set_buffers( {"hidden_states": np.zeros((batch_size, target_seq_len, hidden_size), dtype=np.float32)} ) - tlm_session.skip_buffers(vision_outputs.keys()) + if handoff is None: + tlm_session.skip_buffers(vision_outputs.keys()) if is_qwen35: lang_extra["cache_commit"] = np.zeros(1, dtype=np.int64) @@ -830,7 +865,15 @@ def decode_inputs(draft_ids, decode_position_ids, valid_tokens): return {"input_ids": draft_ids, "position_ids": positions, **decode_selection, **lang_extra} def tlm_decode(draft_ids, decode_position_ids): - outputs = tlm_session.run(decode_inputs(draft_ids, decode_position_ids, block_size)) + nonlocal handoff + inputs = decode_inputs(draft_ids, decode_position_ids, block_size) + if handoff is not None: + inputs.update(handoff.inputs) + outputs = tlm_session.run(inputs) + if handoff is not None: + tlm_session.skip_buffers(handoff.inputs) + tlm_session.skip_buffers(vision_outputs.keys()) + handoff = None return outputs["logits"][:, :block_size], outputs["hidden_states"][:, :block_size] def tlm_commit(draft_ids, decode_position_ids, count): diff --git a/QEfficient/generation/dflash_tlm_handoff.py b/QEfficient/generation/dflash_tlm_handoff.py new file mode 100644 index 0000000000..cad850f5d6 --- /dev/null +++ b/QEfficient/generation/dflash_tlm_handoff.py @@ -0,0 +1,45 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import numpy as np + +from QEfficient.generation.cloud_infer import is_retained_state_name + + +class DFlashTlmHandoff: + """Host buffers for a single-request prefill-to-verification state transfer.""" + + def __init__(self, prefill_session, decode_session): + self.inputs = {} + self.outputs = {} + dtype_mapping = decode_session.aic_to_np_dtype_mapping + if any(len(getattr(session, "allowed_shapes", ())) > 1 for session in (prefill_session, decode_session)): + raise ValueError("Split DFlash TLM requires one specialization in each QPC.") + input_names = {name for name in decode_session.input_names if is_retained_state_name(name)} + output_names = { + name.removesuffix("_RetainedState") + for name in prefill_session.output_names + if name.endswith("_RetainedState") and is_retained_state_name(name) + } + if not input_names or input_names != output_names: + raise ValueError("DFlash TLM prefill/decode QPCs must expose matching retained-state bindings.") + for name in sorted(input_names): + output_name = name + "_RetainedState" + source = prefill_session.bindings[prefill_session.binding_index_map[output_name]] + target = decode_session.bindings[decode_session.binding_index_map[name]] + if tuple(source.dims) != tuple(target.dims) or source.type != target.type or source.size != target.size: + raise ValueError(f"DFlash TLM state shape or precision mismatch for {name}.") + if target.dims[0] != 1: + raise ValueError("Split DFlash TLM currently supports batch size 1 only.") + buffer = np.empty(tuple(target.dims), dtype=dtype_mapping[target.type]) + self.inputs[name] = buffer + self.outputs[output_name] = buffer + + def capture(self, outputs): + """Copy outputs from the standalone session, whose run() returns new arrays.""" + for name, buffer in self.outputs.items(): + np.copyto(buffer, outputs[name]) diff --git a/QEfficient/transformers/cache_utils.py b/QEfficient/transformers/cache_utils.py index f00fdff5b0..9075dc8790 100755 --- a/QEfficient/transformers/cache_utils.py +++ b/QEfficient/transformers/cache_utils.py @@ -1935,9 +1935,12 @@ def sliding_window_update_chunked( # Gather ctx_len = position_ids.shape[1] + sliding_window_len ctx_indices = torch.arange(ctx_len, dtype=position_ids.dtype, device=position_ids.device)[None, None, ...] - first_pos_idx = position_ids[0][0] - add_idx = torch.where(first_pos_idx >= sliding_window_len, first_pos_idx - sliding_window_len, 0) - ctx_indices += add_idx + if cache_kwargs.get("dflash", False): + ctx_indices = ctx_indices + torch.clamp(position_ids[:, :1] - sliding_window_len, min=0).unsqueeze(1) + else: + first_pos_idx = position_ids[0][0] + add_idx = torch.where(first_pos_idx >= sliding_window_len, first_pos_idx - sliding_window_len, 0) + ctx_indices += add_idx gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1).to(torch.int32) invalid_mask = ctx_indices > gather_limit ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices) @@ -1947,6 +1950,10 @@ def sliding_window_update_chunked( else: k_out = ctx_gather(k_out, ctx_indices, ctx_len) v_out = ctx_gather(v_out, ctx_indices, ctx_len) + if cache_kwargs.get("dflash", False): + # CtxGather inherits the full-cache shape; DFlash uses a variable-length window. + k_out = k_out.reshape(key_states.shape[0], key_states.shape[1], ctx_len, key_states.shape[-1]) + v_out = v_out.reshape(value_states.shape[0], value_states.shape[1], ctx_len, value_states.shape[-1]) v_out = torch.where(invalid_mask.unsqueeze(-1), torch.zeros_like(v_out, dtype=v_out.dtype), v_out) return k_out, v_out diff --git a/QEfficient/transformers/models/gpt_oss/modeling_gpt_oss.py b/QEfficient/transformers/models/gpt_oss/modeling_gpt_oss.py index e272616f10..b7ab5a7cff 100644 --- a/QEfficient/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/QEfficient/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -48,6 +48,8 @@ delete_module_attrs, gptoss_clamped_glu_mlp, ) +from QEfficient.transformers.spd.dflash import compute_dflash_target_hidden_states +from QEfficient.transformers.spd.spd_transform_forward import filter_hidden_states from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE from QEfficient.utils.logging_utils import QEFFLogger @@ -613,9 +615,16 @@ def _prepare_gpt_oss_sliding_chunked_attention_mask( sliding_mask: torch.Tensor, position_ids: torch.LongTensor, sliding_window: int, + dflash: bool = False, ) -> torch.Tensor: ctx_len = position_ids.shape[1] + sliding_window ctx_indices = torch.arange(ctx_len, device=position_ids.device) + if dflash: + start = torch.clamp(position_ids[:, :1] - sliding_window, min=0) + indices = ctx_indices.view(1, 1, 1, -1) + start[:, None, None, :] + indices = indices.expand(-1, sliding_mask.shape[1], position_ids.shape[1], -1) + outside_cache = indices >= sliding_mask.shape[-1] + return torch.gather(sliding_mask, -1, indices.clamp(max=sliding_mask.shape[-1] - 1)) | outside_cache first_pos_idx = position_ids[0][0] add_idx = torch.where(first_pos_idx >= sliding_window, first_pos_idx - sliding_window, 0) ctx_indices += add_idx @@ -659,6 +668,7 @@ def forward( "config": self.config, "is_sliding": self.sliding_window is not None, "sliding_window": self.sliding_window, + "dflash": kwargs.pop("dflash", False), } if self.sliding_window is not None: key_states, value_states = past_key_values.sliding_window_update_chunked( @@ -989,17 +999,23 @@ def forward( sin = self.sin_cached[position_ids].unsqueeze(1).to(device=hidden_states.device) cos = self.cos_cached[position_ids].unsqueeze(1).to(device=hidden_states.device) layer_sliding_mask = sliding_mask + target_layer_ids = getattr(self, "target_layer_ids", None) + target_hidden_list = [] if isinstance(self.layers[0].self_attn, QEffPrefillOnlyChunkedGptOssAttention): layer_sliding_mask = _prepare_gpt_oss_sliding_chunked_attention_mask( sliding_mask=sliding_mask, position_ids=position_ids, sliding_window=self.config.sliding_window, + dflash=bool(target_layer_ids), ) - for decoder_layer in self.layers: + for layer_idx, decoder_layer in enumerate(self.layers): if output_hidden_states: all_hidden_states += (hidden_states,) + if target_layer_ids and layer_idx in target_layer_ids: + target_hidden_list.append(hidden_states) + self_attn = decoder_layer.self_attn if self_attn.sliding_window is None: layer_attention_mask = causal_mask @@ -1019,6 +1035,7 @@ def forward( cache_position=cache_position, sin_cached=sin, cos_cached=cos, + dflash=bool(target_layer_ids), **kwargs, ) hidden_states = layer_outputs[0] @@ -1037,6 +1054,11 @@ def forward( return MoeModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None, + hidden_states=( + compute_dflash_target_hidden_states(target_hidden_list, self.fc, self.hidden_norm) + if target_layer_ids + else None + ), ) @@ -1180,6 +1202,7 @@ def forward( return_dict: Optional[bool] = None, output_router_logits: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs: Unpack[TransformersKwargs], ) -> MoeCausalLMOutputWithPast: @@ -1216,6 +1239,12 @@ def forward( output_router_logits if output_router_logits is not None else self.config.output_router_logits ) + is_dflash = bool(getattr(self.model, "target_layer_ids", None)) + if is_dflash: + if not isinstance(self.model.layers[0].self_attn, QEffPrefillOnlyChunkedGptOssAttention): + raise ValueError("GPT-OSS DFlash requires prefill_only=True with chunked, full-context KV caches.") + output_hidden_states = False + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -1238,6 +1267,19 @@ def forward( hidden_states = outputs.last_hidden_state + if is_dflash: + if num_logits_to_keep is not None: + hidden_states = filter_hidden_states(hidden_states, position_ids, num_logits_to_keep) + logits = self.lm_head(hidden_states).float() + qaic_config = getattr(self, "qaic_config", None) or {} + return MoeCausalLMOutputWithPast( + logits=logits + if qaic_config.get("speculative_model_type") == "target" + else logits.argmax(-1).to(torch.int32), + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + ) + logit_index = position_ids.to(torch.int32).argmax(1, keepdim=True) hidden_states = outputs[0][torch.arange(position_ids.shape[0]).view(-1, 1), logit_index] logits = self.lm_head(hidden_states) diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 291ffbda9a..74922994b7 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -2322,6 +2322,7 @@ def compile( compiler_options.pop("kv_cache_batch_size", None) compiler_options.pop("full_batch_size", None) self.qpc_paths = {} + kw_spec = compiler_options.pop("specializations", None) if not skip_vision: compiler_options_vision = compiler_options.copy() compiler_options_vision["node_precision_info"] = False @@ -2391,6 +2392,8 @@ def compile( specializations = lang_specs qpc_key = "lang_qpc_path" + if kw_spec: + specializations = kw_spec lang_qpc_path = self.lang_model._compile( onnx_path=self.lang_model.onnx_path, compile_dir=compile_dir, @@ -3853,6 +3856,8 @@ def __init__( self.hash_params["dflash_tlm_repo"] = qaic_config["dflash_tlm_repo"] if self.dflash_tlm: self.model, _ = DFlashTLMTransform.apply(self.model, qaic_config) + if self.model.config.model_type == "gpt_oss": + self.hash_params["dflash_hidden_norm"] = "custom_rms_norm_v1" self.model, transformed = SpDTransform.apply(self.model, qaic_config=qaic_config, **kwargs) self.is_tlm = transformed diff --git a/QEfficient/transformers/models/pytorch_transforms.py b/QEfficient/transformers/models/pytorch_transforms.py index 73a9376628..3ad68296fa 100644 --- a/QEfficient/transformers/models/pytorch_transforms.py +++ b/QEfficient/transformers/models/pytorch_transforms.py @@ -1166,6 +1166,12 @@ def apply(cls, model: nn.Module, qaic_config: dict | None = None, **kwargs) -> t raise ValueError( f"Speculative model type {speculative_model_type} is not supported. we currently only support {supported_spd_model_types}" ) + elif ( + speculative_model_type == SPD_TARGET + and qaic_config.get("target_layer_ids") + and isinstance(model, QEffGptOssForCausalLM) + ): + return model, True elif (model_class := model.__class__) in cls._module_mapping: model.forward = MethodType(tlm_forward, model) if speculative_model_type != SPD_TARGET: @@ -1325,6 +1331,8 @@ def apply(cls, model: nn.Module, qaic_config: dict | None = None, **kwargs) -> t RMSNorm = Qwen3RMSNorm elif "llama" in model_type: RMSNorm = LlamaRMSNorm + elif model_type == "gpt_oss": + RMSNorm = CustomRMSNormAIC else: RMSNorm = nn.RMSNorm inner.fc = nn.Linear(n * hidden_size, hidden_size, bias=False).to(device=module_device, dtype=module_dtype) diff --git a/examples/performance/dflash/README.md b/examples/performance/dflash/README.md index 30427e4d04..ae5a9b1c1a 100644 --- a/examples/performance/dflash/README.md +++ b/examples/performance/dflash/README.md @@ -12,9 +12,9 @@ single prompt in-process via `QEfficient.generation.dflash_generation`. The thre --- -## Text (dense language models) +## Text models -**Supported models:** Llama-3.1-8B-Instruct, Qwen3-4B, Qwen3-8B (see `--help` for the full list). +**Supported models:** Llama-3.1-8B-Instruct, Qwen3-4B, Qwen3-8B, GPT-OSS-20B, GPT-OSS-120B (see `--help` for the full list). ### Single prompt @@ -25,6 +25,43 @@ python basic_inference_text.py --model_name Qwen3-4B \ --- +### GPT-OSS + +GPT-OSS uses the existing expert-parallel implementation with chunked, full-context +KV caches. Prefill returns one prediction, while verification returns the draft +block's predictions (8 tokens for GPT-OSS-20B, 10 for GPT-OSS-120B). + +```bash +python basic_inference_text.py --model_name gpt-oss-20b \ + --tlm_devices 0,1,2,3 --dlm_devices 4,5,6,7 \ + --ctx_len 4096 --prefill_seq_len 128 --generation_len 256 \ + --prompt "Explain why the sky appears blue in three sentences." +``` + +Standalone QEfficient target QPCs return token IDs; vLLM target QPCs return +vocabulary logits. Use QPCs compiled for the appropriate runtime. + +### Separate MoE TLM QPCs + +For GPT-OSS and Qwen3.5-MoE, `--disagg-tlm` is enabled by default in the +single-prompt examples. The TLM has separate prefill and verification QPCs, +each with one specialization. Prefill uses `--prefill_seq_len` and keeps one +prediction; verification uses the draft block size and keeps every prediction. +Both retain the existing expert-parallel computation. The DLM is unchanged. + +Use `--no-disagg-tlm` to compile or run the previous combined TLM QPC. +Dense models keep the combined-QPC default. Split TLM currently supports BS=1. +After the last prefill chunk, KV state (plus convolution and recurrent state +for Qwen3.5-MoE) is transferred through host buffers to the verification QPC +on its first call. Later verification/replay calls keep state on-device. + +To reuse a split pair, pass `--tlm-prefill-qpc /path/to/prefill/qpc` and +`--tlm_qpc /path/to/verification/qpc`. Prefill uses `--tlm_devices` unless +`--tlm-prefill-devices` is supplied; both groups must have the same device count. +When they share devices, both TLM QPCs are resident, so their combined core +and memory requirements must fit. For example, two 8-core QPCs use all 16 cores +per TLM device. Existing combined QPCs remain usable with `--no-disagg-tlm`. + ## Vision-language models (VLM, multimodal) **Supported models:** `gemma-4-31B-it`, `Qwen3-VL-32B-Instruct` (more VLMs are added to @@ -33,9 +70,10 @@ qwen3-vl) is auto-detected from the TLM's `config.model_type`, so `basic_inferen works for either — gemma4 entries have no default TLM HF path (pass `--tlm_hf_path`); `Qwen3-VL-32B-Instruct` does, so `--tlm_hf_path` is optional for it. -The vision path compiles **three** QPCs: the language decoder (TLM), the **vision +The combined-TLM vision path compiles **three** QPCs: the language decoder (TLM), the **vision encoder** (`pixel_values -> vision_embeds`), and the DFlash draft (DLM). Because the TLM and DLM fill their own cards, the vision encoder usually needs its **own** devices. +Qwen3.5-MoE's default split TLM adds a fourth QPC for language prefill. > Set `QEFF_HOME` to a filesystem with free space before compiling — the VLM TLM QPC > is large (tens of GB): diff --git a/examples/performance/dflash/basic_inference_text.py b/examples/performance/dflash/basic_inference_text.py index a394bda2cb..4c2b752006 100644 --- a/examples/performance/dflash/basic_inference_text.py +++ b/examples/performance/dflash/basic_inference_text.py @@ -44,6 +44,7 @@ compile_tlm_qpc, format_prompt, get_spd_prompt_chunk_size, + load_spd_session, load_spd_sessions, resolve_model_name, validate_spd_decode_specialization, @@ -96,6 +97,15 @@ def parse_args(): # Optional pre-built QPCs (skip compilation) p.add_argument("--tlm_qpc", default=None, help="Pre-compiled TLM qpc dir (skip TLM compile).") p.add_argument("--dlm_qpc", default=None, help="Pre-compiled DLM qpc dir (skip DLM compile).") + p.add_argument( + "--disagg-tlm", + "--disagg_tlm", + action=argparse.BooleanOptionalAction, + default=None, + help="Separate TLM prefill/verification QPCs (default: enabled for MoE models).", + ) + p.add_argument("--tlm-prefill-qpc", "--tlm_prefill_qpc") + p.add_argument("--tlm-prefill-devices", "--tlm_prefill_devices", type=parse_device_list) # Devices / cores p.add_argument( @@ -150,6 +160,20 @@ def main(): if tlm_repo is None: raise SystemExit(f"No default TLM HF path for '{args.model_name}'. Pass --tlm_hf_path.") + tlm_config = AutoConfig.from_pretrained(tlm_repo, token=args.hf_token) + if args.disagg_tlm is None: + args.disagg_tlm = tlm_config.model_type == "gpt_oss" + if args.disagg_tlm and tlm_config.model_type != "gpt_oss": + raise ValueError("Split text DFlash TLM is currently supported for GPT-OSS only.") + if args.tlm_qpc and args.disagg_tlm and not args.tlm_prefill_qpc: + raise ValueError("Split TLM requires --tlm-prefill-qpc as well as --tlm_qpc, or use --no-disagg-tlm.") + if args.tlm_prefill_qpc and not args.disagg_tlm: + raise ValueError("--tlm-prefill-qpc requires --disagg-tlm.") + if args.tlm_prefill_qpc and not args.tlm_qpc: + raise ValueError("Provide both --tlm-prefill-qpc and --tlm_qpc to reuse a split TLM pair.") + if args.tlm_prefill_devices and len(args.tlm_prefill_devices) != len(args.tlm_devices): + raise ValueError("Split TLM prefill and verification device groups must have the same device count.") + if args.tlm_qpc: logger.info(f"[skip compile] using provided TLM qpc: {args.tlm_qpc}") tlm_qpc = args.tlm_qpc @@ -162,7 +186,10 @@ def main(): num_cores=args.tlm_cores, num_devices=len(args.tlm_devices), hf_token=args.hf_token, + disagg_tlm=args.disagg_tlm, ) + if args.disagg_tlm: + args.tlm_prefill_qpc, tlm_qpc = tlm_qpc["prefill"], tlm_qpc["decode"] if args.dlm_qpc: logger.info(f"[skip compile] using provided DLM qpc: {args.dlm_qpc}") @@ -177,6 +204,8 @@ def main(): hf_token=args.hf_token, ) logger.info(f"TLM qpc : {tlm_qpc}") + if args.disagg_tlm: + logger.info(f"TLM prefill qpc: {args.tlm_prefill_qpc}") logger.info(f"DLM qpc : {dlm_qpc}") prompt_text = format_prompt(args.prompt, args.category) if args.format_prompt else args.prompt @@ -186,7 +215,12 @@ def main(): tokenizer.pad_token_id = tokenizer.eos_token_id dlm_session, tlm_session = load_spd_sessions(tlm_qpc, dlm_qpc, args.tlm_devices, args.dlm_devices) - prompt_chunk_size = get_spd_prompt_chunk_size(tlm_session) + tlm_prefill_session = ( + load_spd_session(args.tlm_prefill_qpc, args.tlm_prefill_devices or args.tlm_devices) + if args.disagg_tlm + else None + ) + prompt_chunk_size = get_spd_prompt_chunk_size(tlm_prefill_session or tlm_session) validate_spd_decode_specialization(tlm_session, config.block_size) dflash_config = getattr(config, "dflash_config", None) or config.to_dict().get("dflash_config", {}) @@ -209,6 +243,7 @@ def main(): hidden_size=config.hidden_size, generation_len=args.generation_len, batch_size=batch_size, + tlm_prefill_session=tlm_prefill_session, ) output_parts = ["Output: "] diff --git a/examples/performance/dflash/basic_inference_vision.py b/examples/performance/dflash/basic_inference_vision.py index b5ab87e942..a3df241090 100644 --- a/examples/performance/dflash/basic_inference_vision.py +++ b/examples/performance/dflash/basic_inference_vision.py @@ -25,6 +25,7 @@ compile_qwen3vl_vlm_qpcs, get_dflash_block_size, get_spd_prompt_chunk_size, + load_spd_session, load_spd_sessions, resolve_model_name, validate_spd_decode_specialization, @@ -72,6 +73,15 @@ def parse_args(): parser.add_argument("--tlm_qpc", default=None) parser.add_argument("--vision_qpc", default=None) parser.add_argument("--dlm_qpc", default=None) + parser.add_argument( + "--disagg-tlm", + "--disagg_tlm", + action=argparse.BooleanOptionalAction, + default=None, + help="Separate TLM prefill/verification QPCs (default: enabled for MoE models).", + ) + parser.add_argument("--tlm-prefill-qpc", "--tlm_prefill_qpc") + parser.add_argument("--tlm-prefill-devices", "--tlm_prefill_devices", type=parse_device_list) parser.add_argument("--tlm_devices", type=parse_device_list, default=[0, 1, 2, 3]) parser.add_argument("--dlm_devices", type=parse_device_list, default=[0, 1, 2, 3]) parser.add_argument("--vision_devices", type=parse_device_list, default=[0, 1, 2, 3]) @@ -194,6 +204,7 @@ def _run_qwen3_vl( prompt_chunk_size, vision_qpc, batch_size, + tlm_prefill_session=None, ): if processor is None: raise RuntimeError("Qwen3-VL requires a processor, which failed to load.") @@ -233,6 +244,7 @@ def _run_qwen3_vl( compiled_height=args.height if args.height is not None else 354, compiled_width=args.width if args.width is not None else 536, batch_size=batch_size, + tlm_prefill_session=tlm_prefill_session, ) return metrics, {"vision_encode_time_s": metrics.vision_prefill_time if args.image else None} @@ -250,6 +262,18 @@ def main(): tlm_config = AutoConfig.from_pretrained(tlm_repo, trust_remote_code=True, token=args.hf_token) is_qwen3vl = tlm_config.model_type in {"qwen3_vl", "qwen3_5", "qwen3_5_moe"} + if args.disagg_tlm is None: + args.disagg_tlm = tlm_config.model_type == "qwen3_5_moe" + if args.disagg_tlm and tlm_config.model_type != "qwen3_5_moe": + raise ValueError("Split VLM TLM is currently supported for Qwen3.5-MoE only.") + if args.tlm_qpc and args.disagg_tlm and not args.tlm_prefill_qpc: + raise ValueError("Split TLM requires --tlm-prefill-qpc as well as --tlm_qpc, or use --no-disagg-tlm.") + if args.tlm_prefill_qpc and not args.disagg_tlm: + raise ValueError("--tlm-prefill-qpc requires --disagg-tlm.") + if args.tlm_prefill_qpc and not args.tlm_qpc: + raise ValueError("Provide both --tlm-prefill-qpc and --tlm_qpc to reuse a split TLM pair.") + if args.tlm_prefill_devices and len(args.tlm_prefill_devices) != len(args.tlm_devices): + raise ValueError("Split TLM prefill and verification device groups must have the same device count.") if args.tlm_qpc and args.vision_qpc: tlm_qpc, vision_qpc = args.tlm_qpc, args.vision_qpc @@ -265,8 +289,10 @@ def main(): "hf_token": args.hf_token, } if is_qwen3vl: - compile_kwargs.update({"height": args.height, "width": args.width}) + compile_kwargs.update({"height": args.height, "width": args.width, "disagg_tlm": args.disagg_tlm}) tlm_qpc, vision_qpc = compile_vlm(tlm_repo, dlm_repo, **compile_kwargs) + if args.disagg_tlm: + args.tlm_prefill_qpc, tlm_qpc = tlm_qpc["prefill"], tlm_qpc["decode"] if args.dlm_qpc: dlm_qpc = args.dlm_qpc @@ -282,6 +308,8 @@ def main(): ) logger.info(f"TLM lang qpc: {tlm_qpc}") + if args.disagg_tlm: + logger.info(f"TLM prefill qpc: {args.tlm_prefill_qpc}") logger.info(f"Vision qpc: {vision_qpc}") logger.info(f"DLM qpc: {dlm_qpc}") @@ -296,7 +324,12 @@ def main(): logger.warning(f"No processor for {tlm_repo} ({error}); using tokenizer template.") dlm_session, tlm_session = load_spd_sessions(tlm_qpc, dlm_qpc, args.tlm_devices, args.dlm_devices) - prompt_chunk_size = get_spd_prompt_chunk_size(tlm_session) + tlm_prefill_session = ( + load_spd_session(args.tlm_prefill_qpc, args.tlm_prefill_devices or args.tlm_devices) + if args.disagg_tlm + else None + ) + prompt_chunk_size = get_spd_prompt_chunk_size(tlm_prefill_session or tlm_session) target_decode_length = get_dflash_block_size(config) validate_spd_decode_specialization(tlm_session, target_decode_length) @@ -312,6 +345,7 @@ def main(): prompt_chunk_size, vision_qpc, batch_size, + tlm_prefill_session, ) else: metrics, output_extra = _run_gemma( diff --git a/examples/performance/dflash/utils.py b/examples/performance/dflash/utils.py index 365485d3a2..c1b63aedf4 100644 --- a/examples/performance/dflash/utils.py +++ b/examples/performance/dflash/utils.py @@ -30,6 +30,8 @@ # when known; otherwise None and must be supplied via --tlm_hf_path. # ───────────────────────────────────────────────────────────────────────────── MODEL_MAP = { + "gpt-oss-20b": ("openai/gpt-oss-20b", "z-lab/gpt-oss-20b-DFlash"), + "gpt-oss-120b": ("openai/gpt-oss-120b", "z-lab/gpt-oss-120b-DFlash"), "gemma-4-31B-it": ("google/gemma-4-31B-it", "z-lab/gemma-4-31B-it-DFlash"), "Qwen3-4B": ("Qwen/Qwen3-4B", "z-lab/Qwen3-4B-DFlash-b16"), "Qwen3.5-4B": ("Qwen/Qwen3.5-4B", "z-lab/Qwen3.5-4B-DFlash"), @@ -57,19 +59,21 @@ def _build_aliases(model_map): MODEL_ALIASES = _build_aliases(MODEL_MAP) +def load_spd_session(qpc: str, devices: list[int]): + session = QAICInferenceSession(qpc, devices) + session.skip_buffers( + [ + name + for name in session.input_names + session.output_names + if is_retained_state_name(name) or name.endswith("_RetainedState") + ] + ) + return session + + def load_spd_sessions(tlm_qpc: str, dlm_qpc: str, tlm_devices: list[int], dlm_devices: list[int]): """Create DFlash sessions and enable retained-state cache handling.""" - dlm_session = QAICInferenceSession(dlm_qpc, dlm_devices) - tlm_session = QAICInferenceSession(tlm_qpc, tlm_devices) - for session in (dlm_session, tlm_session): - session.skip_buffers( - [ - name - for name in session.input_names + session.output_names - if is_retained_state_name(name) or name.endswith("_RetainedState") - ] - ) - return dlm_session, tlm_session + return load_spd_session(dlm_qpc, dlm_devices), load_spd_session(tlm_qpc, tlm_devices) def get_spd_prompt_chunk_size(tlm_session: QAICInferenceSession) -> int: @@ -182,6 +186,20 @@ def read_dlm_meta(dlm_repo: str, hf_token: str | None = None): return state_dict, target_layer_ids, block_size +def _compile_tlm_phases(model, compile_kwargs, disagg_tlm): + if not disagg_tlm: + return model.compile(**compile_kwargs) + compile_kwargs = compile_kwargs.copy() + specializations = compile_kwargs.pop("specializations") + qpcs = {} + for phase, spec in zip(("prefill", "decode"), specializations): + qpc = model.compile( + **compile_kwargs, specializations=[spec], split_retained_state_io=True, offload_pt_weights=False + ) + qpcs[phase] = dict(qpc) if isinstance(qpc, dict) else str(qpc) + return qpcs + + def compile_tlm_qpc( tlm_repo: str, dlm_repo: str, @@ -191,7 +209,8 @@ def compile_tlm_qpc( num_cores: int, num_devices: int, hf_token: str | None = None, -) -> str: + disagg_tlm: bool = False, +) -> str | dict[str, str]: """Build the TLM and compile it to a QPC. fc/hidden_norm are injected by QEfficient's DFlashTLMTransform (weights from dflash_dlm_repo).""" from QEfficient import QEFFAutoModelForCausalLM @@ -200,13 +219,36 @@ def compile_tlm_qpc( tlm_target_ids = [i + 1 for i in target_layer_ids] logger.info(f"[compile_tlm] base={tlm_repo} dlm={dlm_repo} block_size={block_size}") + config = AutoConfig.from_pretrained(tlm_repo, token=hf_token) + if disagg_tlm and config.model_type != "gpt_oss": + raise ValueError("Split text DFlash TLM compilation currently supports GPT-OSS only.") + qaic_config = {"target_layer_ids": tlm_target_ids, "dflash_dlm_repo": dlm_repo} + compile_kwargs = {} + if config.model_type == "gpt_oss": + compile_kwargs = { + "qaic_config": qaic_config, + "prefill_only": True, + "enable_chunking": True, + "use_onnx_subfunctions": False, + "specializations": [ + { + "batch_size": 1, + "seq_len": length, + "ctx_len": ctx_len, + "sliding_window": config.sliding_window, + "num_logits_to_keep": num_logits, + } + for length, num_logits in ((prefill_seq_len, 1), (block_size, block_size)) + ], + } tlm_qeff = QEFFAutoModelForCausalLM.from_pretrained( tlm_repo, + config=config, torch_dtype=torch.float32, token=hf_token, - qaic_config={"target_layer_ids": tlm_target_ids, "dflash_dlm_repo": dlm_repo}, + qaic_config=qaic_config, ) - qpc = tlm_qeff.compile( + compile_kwargs.update( prefill_seq_len=prefill_seq_len, ctx_len=ctx_len, num_cores=num_cores, @@ -216,7 +258,9 @@ def compile_tlm_qpc( mos=1, dflash_block_size=block_size, ) - qpc = str(qpc) + qpc = _compile_tlm_phases(tlm_qeff, compile_kwargs, disagg_tlm) + if not disagg_tlm: + qpc = str(qpc) logger.info(f"[compile_tlm] qpc={qpc}") return qpc @@ -350,7 +394,8 @@ def compile_qwen3vl_vlm_qpcs( height: int | None = None, width: int | None = None, hf_token: str | None = None, -) -> tuple[str, str]: + disagg_tlm: bool = False, +) -> tuple[str | dict[str, str], str]: """Build the Qwen3-VL or Qwen3.5 TLM vision encoder + language decoder for SPD. Returns ``(lang_qpc, vision_qpc)``. @@ -374,6 +419,8 @@ def compile_qwen3vl_vlm_qpcs( config = AutoConfig.from_pretrained(tlm_repo, trust_remote_code=True, token=hf_token) is_qwen35 = config.model_type == "qwen3_5" is_qwen35_moe = config.model_type == "qwen3_5_moe" + if disagg_tlm and not is_qwen35_moe: + raise ValueError("Split VLM DFlash TLM compilation currently supports Qwen3.5-MoE only.") qaic_config = {"target_layer_ids": tlm_target_ids, "dflash_dlm_repo": dlm_repo} if is_qwen35 or is_qwen35_moe: if block_size < 2 or block_size & (block_size - 1): @@ -426,9 +473,25 @@ def compile_qwen3vl_vlm_qpcs( if width is not None: compile_kwargs["width"] = width - qpc = tlm_qeff.compile(**compile_kwargs) - lang_qpc = str(qpc["lang_prefill_qpc_path"] if is_qwen35 or is_qwen35_moe else qpc["lang_qpc_path"]) - vision_qpc = str(qpc["vision_qpc_path"]) + if disagg_tlm: + text_model.dflash_chunk_size = block_size + specializations, _ = tlm_qeff.model.get_specializations( + batch_size=1, + prefill_seq_len=prefill_seq_len, + ctx_len=ctx_len, + height=height, + width=width, + kv_offload=True, + ) + compile_kwargs["specializations"] = specializations["lang"] + + qpc = _compile_tlm_phases(tlm_qeff, compile_kwargs, disagg_tlm) + if disagg_tlm: + lang_qpc = {phase: str(paths["lang_prefill_qpc_path"]) for phase, paths in qpc.items()} + vision_qpc = str(qpc["prefill"]["vision_qpc_path"]) + else: + lang_qpc = str(qpc["lang_prefill_qpc_path"] if is_qwen35 or is_qwen35_moe else qpc["lang_qpc_path"]) + vision_qpc = str(qpc["vision_qpc_path"]) logger.info(f"[compile_qwen3vl_tlm] lang_qpc={lang_qpc} vision_qpc={vision_qpc}") return lang_qpc, vision_qpc diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index 775481f530..5a5863827f 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -861,6 +861,298 @@ def _ort_qwen_logits(onnx_path: Path, qeff_model, inputs): return outputs["logits"].astype(np.float32) +@pytest.mark.parametrize("model_name,block_size", [("gpt-oss-20b", 8), ("gpt-oss-120b", 10)]) +@pytest.mark.parametrize("disagg_tlm", [False, True]) +def test_gpt_oss_dflash_compile_contract(monkeypatch, tmp_path, model_name, block_size, disagg_tlm): + from examples.performance.dflash import utils + + model = MagicMock() + model.compile.return_value = tmp_path / "qpc" + load_model = MagicMock(return_value=model) + monkeypatch.setattr(QEFFAutoModelForCausalLM, "from_pretrained", load_model) + monkeypatch.setattr(utils.AutoConfig, "from_pretrained", lambda *a, **kw: GptOssConfig()) + monkeypatch.setattr(utils, "read_dlm_meta", lambda *a: ({}, [1, 3], block_size)) + tlm_repo, dlm_repo = utils.MODEL_MAP[utils.resolve_model_name(f"openai/{model_name}")] + result = utils.compile_tlm_qpc( + tlm_repo, dlm_repo, prefill_seq_len=128, ctx_len=4096, num_cores=8, num_devices=4, disagg_tlm=disagg_tlm + ) + if disagg_tlm: + assert set(result) == {"prefill", "decode"} + assert model.compile.call_count == 2 + assert all(len(call.kwargs["specializations"]) == 1 for call in model.compile.call_args_list) + assert all(call.kwargs["offload_pt_weights"] is False for call in model.compile.call_args_list) + assert all(call.kwargs["split_retained_state_io"] is True for call in model.compile.call_args_list) + else: + assert model.compile.call_count == 1 + specs = [] + for call in model.compile.call_args_list: + options = call.kwargs + assert options["qaic_config"] == load_model.call_args.kwargs["qaic_config"] + assert "moe_config" not in options["qaic_config"] + assert options["prefill_only"] and options["enable_chunking"] + assert options["mxint8_kv_cache"] and options["mxfp6_matmul"] + specs.extend(options["specializations"]) + assert [(s["seq_len"], s["num_logits_to_keep"]) for s in specs] == [ + (128, 1), + (block_size, block_size), + ] + + +@pytest.mark.parametrize("prefill_length", [16, 128]) +def test_dflash_split_tlm_phase_selection(prefill_length): + from examples.performance.dflash.utils import _compile_tlm_phases + + specs = [{"seq_len": prefill_length, "num_logits_to_keep": 1}, {"seq_len": 16, "num_logits_to_keep": 16}] + model = MagicMock() + model.compile.side_effect = ["prefill.qpc", "decode.qpc"] + options = {"prefill_only": True, "specializations": specs} + assert _compile_tlm_phases(model, options, True) == {"prefill": "prefill.qpc", "decode": "decode.qpc"} + assert options == {"prefill_only": True, "specializations": specs} + for call, spec in zip(model.compile.call_args_list, specs): + assert call.kwargs == { + "prefill_only": True, + "specializations": [spec], + "split_retained_state_io": True, + "offload_pt_weights": False, + } + + +def test_dflash_split_tlm_handoff_includes_hybrid_state(): + from types import SimpleNamespace + + from QEfficient.generation.dflash_tlm_handoff import DFlashTlmHandoff + + names = ["past_key.3", "past_value.3", "conv_state.0", "recurrent_state.0"] + shapes = [(1, 2, 32, 8), (1, 2, 32, 8), (1, 8, 4), (1, 8, 8)] + bindings = [SimpleNamespace(dims=shape, type=1, size=int(np.prod(shape)) * 4) for shape in shapes] + source = SimpleNamespace(output_names=[name + "_RetainedState" for name in names], bindings=bindings) + source.binding_index_map = {name: index for index, name in enumerate(source.output_names)} + target = SimpleNamespace(input_names=names, bindings=bindings, aic_to_np_dtype_mapping={1: np.dtype(np.float32)}) + target.binding_index_map = {name: index for index, name in enumerate(names)} + handoff = DFlashTlmHandoff(source, target) + outputs = { + name: np.full(buffer.shape, index + 1, dtype=np.float32) + for index, (name, buffer) in enumerate(handoff.outputs.items()) + } + handoff.capture(outputs) + for name, buffer in handoff.inputs.items(): + np.testing.assert_array_equal(buffer, outputs[name + "_RetainedState"]) + assert buffer is handoff.outputs[name + "_RetainedState"] + source.output_names.pop() + with pytest.raises(ValueError, match="matching retained-state"): + DFlashTlmHandoff(source, target) + + +@pytest.mark.parametrize("prefill_length", [16, 128]) +@pytest.mark.parametrize("block_size", [8, 16]) +def test_dflash_split_tlm_qwen_moe_compile(monkeypatch, prefill_length, block_size): + from examples.performance.dflash import utils + from QEfficient.transformers.models import modeling_auto + from tests.unit_test.transforms.test_dflash_transform import make_tiny_qwen3_5_vlm + + model = QEFFAutoModelForImageTextToText( + make_tiny_qwen3_5_vlm(moe=True), kv_offload=True, qaic_config={"target_layer_ids": [1, 3]} + ) + compile_mock = MagicMock(return_value="test.qpc") + monkeypatch.setattr(model.lang_model, "_compile", compile_mock) + vision_compile = MagicMock(return_value="vision.qpc") + monkeypatch.setattr(model.vision_model, "_compile", vision_compile) + monkeypatch.setattr(model, "export", MagicMock()) + monkeypatch.setattr(modeling_auto, "_filter_custom_io_for_onnx", lambda custom_io, onnx_path: custom_io) + monkeypatch.setattr(QEFFAutoModelForImageTextToText, "from_pretrained", lambda *a, **kw: model) + monkeypatch.setattr(utils.AutoConfig, "from_pretrained", lambda *a, **kw: model.model.config) + text_model = model.lang_model.model.language_model + state_dict = { + "fc.weight": text_model.fc.weight.detach().clone(), + "hidden_norm.weight": text_model.hidden_norm.weight, + } + monkeypatch.setattr(utils, "read_dlm_meta", lambda *a: (state_dict, [0, 2], block_size)) + lang_qpcs, vision_qpc = utils.compile_qwen3vl_vlm_qpcs( + "tiny-qwen3.5-moe", + "tiny-dflash", + num_cores=4, + num_devices=1, + prefill_seq_len=prefill_length, + ctx_len=256, + height=32, + width=32, + disagg_tlm=True, + ) + assert lang_qpcs == {"prefill": "test.qpc", "decode": "test.qpc"} + assert vision_qpc == "vision.qpc" + assert compile_mock.call_count == 2 + for call, (length, num_logits) in zip(compile_mock.call_args_list, [(prefill_length, 1), (block_size, block_size)]): + specs = call.kwargs["specializations"] + assert len(specs) == 1 + assert specs[0]["num_logits_to_keep"] == num_logits + assert specs[0]["seq_len"] == length + assert "vision_size" in specs[0] + assert call.kwargs["split_retained_state_io"] is True + for call in vision_compile.call_args_list: + assert "grid_height" in call.kwargs["specializations"][0] + assert "num_logits_to_keep" not in call.kwargs["specializations"][0] + + +@pytest.mark.parametrize( + "mode,continuous_batching", [("baseline", False), ("standalone", False), ("target", False), ("target", True)] +) +@pytest.mark.parametrize("prefill_length", [8, 16]) +def test_gpt_oss_dflash_expert_parallel_contract(tmp_path, mode, prefill_length, continuous_batching): + from QEfficient.transformers.moe import MoEFlavour + + torch.manual_seed(37) + config = GptOssConfig( + vocab_size=127, + hidden_size=32, + intermediate_size=32, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=16, + num_local_experts=4, + num_experts_per_tok=2, + sliding_window=16, + max_position_embeddings=128, + ) + hf_model = AutoModelForCausalLM.from_config(config, attn_implementation="eager").eval() + dflash = mode != "baseline" + qaic_config = {"target_layer_ids": [1, 3]} if dflash else {} + if mode == "target": + qaic_config["speculative_model_type"] = "target" + qeff = QEFFAutoModelForCausalLM( + deepcopy(hf_model), qaic_config=qaic_config, continuous_batching=continuous_batching + ) + qeff.transform( + ctx_len=64, + seq_len=prefill_length, + bs=1, + num_devices=1, + num_cores=4, + qaic_config=qaic_config, + prefill_only=True, + enable_chunking=True, + prefill_seq_len=prefill_length, + ) + assert all(layer.mlp._moe_flavour == MoEFlavour.EXPERT_PARALLEL for layer in qeff.model.model.layers) + path = qeff.export( + tmp_path / "gpt-oss-dflash", + prefill_only=True, + enable_chunking=True, + prefill_seq_len=prefill_length, + offload_pt_weights=False, + ) + graph = onnx.load(_exported_onnx_path(path), load_external_data=False).graph + assert ("num_logits_to_keep" in [value.name for value in graph.input]) == dflash + if dflash: + assert qeff.hash_params["dflash_hidden_norm"] == "custom_rms_norm_v1" + assert any(node.op_type == "CustomRMSNorm" and "hidden_states" in node.output for node in graph.node) + else: + assert "dflash_hidden_norm" not in qeff.hash_params + cache = _qwen_causal_inputs(config, seq_len=8, ctx_len=64)["past_key_values"] + if continuous_batching: + cache = tuple(tuple(value.repeat(3, 1, 1, 1) for value in layer) for layer in cache) + committed = torch.randint(1, config.vocab_size, (1, 13)) + + def run(ids, start, valid, keep): + nonlocal cache + positions = torch.arange(start, start + ids.shape[1]).unsqueeze(0) + positions[:, valid:] = -1 + inputs = {"input_ids": ids, "position_ids": positions, "past_key_values": cache} + if continuous_batching: + inputs["batch_index"] = torch.tensor([[2]]) + if dflash: + inputs["num_logits_to_keep"] = torch.arange(keep).view(keep, 1) + with torch.no_grad(): + outputs = qeff.model(**inputs) + cache = outputs.past_key_values + assert outputs.logits.shape[1] == (keep if dflash else 1) + if dflash: + assert outputs.hidden_states.shape == (1, ids.shape[1], config.hidden_size) + return outputs + + for start in range(0, committed.shape[1], prefill_length): + chunk = committed[:, start : start + prefill_length] + actual = run(torch.nn.functional.pad(chunk, (0, prefill_length - chunk.shape[1])), start, chunk.shape[1], 1) + if dflash: + with torch.no_grad(): + expected = hf_model( + input_ids=committed[:, : start + chunk.shape[1]], use_cache=False, output_hidden_states=True + ) + features = qeff.model.model.hidden_norm( + qeff.model.model.fc(torch.cat([expected.hidden_states[i][:, start:] for i in [1, 3]], dim=-1)) + ) + torch.testing.assert_close(actual.hidden_states[:, : chunk.shape[1]], features, atol=3e-4, rtol=3e-4) + for accepted in (2, 0, 7): + candidates = torch.randint(1, config.vocab_size, (1, 8)) + actual = run(candidates, committed.shape[1], 8, 8) + with torch.no_grad(): + expected = hf_model( + input_ids=torch.cat((committed, candidates), dim=1), use_cache=False, output_hidden_states=True + ) + logits = expected.logits[:, -8:] if dflash else expected.logits[:, -1:] + if mode == "standalone": + logits = logits.argmax(-1).to(torch.int32) + torch.testing.assert_close(actual.logits, logits, atol=3e-4, rtol=3e-4) + if dflash: + features = qeff.model.model.hidden_norm( + qeff.model.model.fc(torch.cat([expected.hidden_states[i][:, -8:] for i in [1, 3]], dim=-1)) + ) + torch.testing.assert_close(actual.hidden_states, features, atol=3e-4, rtol=3e-4) + committed = torch.cat((committed, candidates[:, : accepted + 1]), dim=1) + + +def test_gpt_oss_dflash_ctx_gather_dynamic_length(tmp_path, monkeypatch): + from QEfficient.transformers import cache_utils + from QEfficient.transformers.cache_utils import QEffGPTOSSDynamicLayer + + # Isolate gather shape propagation from hardware scatter operations. + monkeypatch.setattr(cache_utils, "ctx_scatter", lambda data, *args: data) + + class GatherWindow(torch.nn.Module): + def forward(self, data, updates, positions): + layer = QEffGPTOSSDynamicLayer.from_tensors(data, data, is_sliding=True) + gathered, values = layer.sliding_window_update_chunked( + updates, + updates, + { + "position_ids": positions, + "sliding_window": 16, + "dflash": True, + }, + ) + return gathered.reshape(updates.shape[0], -1, gathered.shape[2] * data.shape[3]), values + + data = torch.randn(2, 2, 64, 8) + updates = torch.randn(2, 2, 8, 8) + positions = torch.tensor([[4], [24]]) + torch.arange(8) + path = tmp_path / "gather-window.onnx" + model = GatherWindow() + torch.onnx.export( + model, + (data, updates, positions), + path, + input_names=["data", "updates", "positions"], + output_names=["gathered", "values"], + dynamic_axes={ + "updates": {2: "seq_len"}, + "positions": {1: "seq_len"}, + "gathered": {2: "window_features"}, + "values": {2: "window_len"}, + }, + opset_version=17, + dynamo=False, + ) + session = _ort_session(path) + for length in (8, 16, 32): + updates = torch.randn(2, 2, length, 8) + positions = torch.tensor([[4], [24]]) + torch.arange(length) + expected = model(data, updates, positions) + feeds = {"data": data.numpy(), "updates": updates.numpy(), "positions": positions.numpy()} + actual = session.run(None, {item.name: feeds[item.name] for item in session.get_inputs()}) + for result, reference in zip(actual, expected): + np.testing.assert_array_equal(result, reference.numpy()) + + @pytest.mark.parametrize("model_type", ["qwen3_5", "qwen3_5_moe"]) def test_qwen3_5_without_dflash_hf_qeff_ort_parity(tmp_path, model_type): _assert_qwen_hf_qeff_ort_parity(model_type, tmp_path) From a47e58dc8ee12f7c3cb9f85e8243f379d8cd0ba8 Mon Sep 17 00:00:00 2001 From: Vahid Janfaza Date: Fri, 23 Oct 2026 11:23:28 -0700 Subject: [PATCH 8/9] Address DFlash text example and artifact-cache review comments Use the shared block-size helper once in the text example for both validation and inference. Document required re-export and recompilation of legacy Qwen3.5 DFlash TLM artifacts. Co-authored-by: FNU Annanya Signed-off-by: Vahid Janfaza --- examples/performance/dflash/README.md | 10 ++++++++-- examples/performance/dflash/basic_inference_text.py | 6 ++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/examples/performance/dflash/README.md b/examples/performance/dflash/README.md index ae5a9b1c1a..a273ca33b1 100644 --- a/examples/performance/dflash/README.md +++ b/examples/performance/dflash/README.md @@ -128,6 +128,11 @@ python basic_inference_vision.py --model_name Qwen3.5-4B \ Use `--prompt "..."` instead of `--image --image_prompt ...` for text-only input. Prefill length can be any positive multiple of the draft block size; verification and replay remain block-sized. +**Recompilation required:** Qwen3.5 DFlash TLM ONNX/QPC artifacts created before +the variable-length prefill and block-sized verification/replay update are incompatible. +Re-export and recompile the TLM; do not reuse those old paths with `--tlm_qpc`. +To let this example rebuild the VLM artifacts, omit both `--tlm_qpc` and `--vision_qpc`. + --- ## `--model_name` @@ -164,8 +169,9 @@ DFlash targets use the existing SpD `num_logits_to_keep` input to project only the last valid position during prompt prefill. Verification still predicts every block position, and all hidden states needed by the DLM are preserved. This applies automatically to Qwen3, Llama, Gemma4, Qwen3-VL, Qwen3.5, and Qwen3.5-MoE. -Existing TLM QPCs remain usable but need re-export/compilation for this change; -DLM and vision QPCs are unchanged. +Apart from the incompatible Qwen3.5 artifacts noted above, existing TLM QPCs remain +usable; re-export/recompile to benefit from this optimization. DLM and vision QPCs +are unchanged. | Flag | Default | Notes | |---|---|---| diff --git a/examples/performance/dflash/basic_inference_text.py b/examples/performance/dflash/basic_inference_text.py index 4c2b752006..4706baadce 100644 --- a/examples/performance/dflash/basic_inference_text.py +++ b/examples/performance/dflash/basic_inference_text.py @@ -43,6 +43,7 @@ compile_dlm_qpc, compile_tlm_qpc, format_prompt, + get_dflash_block_size, get_spd_prompt_chunk_size, load_spd_session, load_spd_sessions, @@ -211,6 +212,7 @@ def main(): prompt_text = format_prompt(args.prompt, args.category) if args.format_prompt else args.prompt tokenizer = AutoTokenizer.from_pretrained(tlm_repo, token=args.hf_token, trust_remote_code=True) config = AutoConfig.from_pretrained(dlm_repo, token=args.hf_token, trust_remote_code=True) + block_size = get_dflash_block_size(config) if tokenizer.pad_token_id is None: tokenizer.pad_token_id = tokenizer.eos_token_id @@ -221,7 +223,7 @@ def main(): else None ) prompt_chunk_size = get_spd_prompt_chunk_size(tlm_prefill_session or tlm_session) - validate_spd_decode_specialization(tlm_session, config.block_size) + validate_spd_decode_specialization(tlm_session, block_size) dflash_config = getattr(config, "dflash_config", None) or config.to_dict().get("dflash_config", {}) mask_token_id = dflash_config["mask_token_id"] if isinstance(dflash_config, dict) else dflash_config.mask_token_id @@ -238,7 +240,7 @@ def main(): vocab_size=config.vocab_size, prompt_chunk_size=prompt_chunk_size, ctx_len=args.ctx_len, - block_size=config.block_size, + block_size=block_size, max_iterations=args.iteration, hidden_size=config.hidden_size, generation_len=args.generation_len, From 560c02a8989503cf33430c5b544cdb54d6de216c Mon Sep 17 00:00:00 2001 From: Vahid Janfaza Date: Wed, 28 Oct 2026 16:33:38 -0700 Subject: [PATCH 9/9] Optimize Qwen3.5 DFlash with concatenated target decoding Combine accepted-prefix commitment and verification in one target call for Qwen3.5 dense and MoE models. Use twice the draft block width for target inputs while keeping block-sized logits and draft features. Update the export cache protocol, validate fixed-shape QPCs with folded logits selectors, and align the DFlash examples, timing labels, and regression coverage with the concatenated path. Preserve the existing BS=1 VLM restrictions. Validation: 175 CPU tests passed; hardware compilation and execution were not run for this push. Signed-off-by: Vahid Janfaza --- QEfficient/generation/dflash_generation.py | 94 ++++---- .../transformers/models/modeling_auto.py | 8 +- .../models/qwen3_5/modeling_qwen3_5.py | 90 ++++---- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 89 ++++---- QEfficient/transformers/spd/dflash.py | 7 +- examples/performance/dflash/README.md | 33 ++- .../dflash/basic_inference_vision.py | 16 +- examples/performance/dflash/utils.py | 20 +- .../spd/test_dflash_spd_inference.py | 64 +++--- .../transforms/test_dflash_transform.py | 203 +++++++++++++----- 10 files changed, 396 insertions(+), 228 deletions(-) diff --git a/QEfficient/generation/dflash_generation.py b/QEfficient/generation/dflash_generation.py index 5713751983..d233498364 100644 --- a/QEfficient/generation/dflash_generation.py +++ b/QEfficient/generation/dflash_generation.py @@ -53,9 +53,8 @@ def __init__(self, block_size: int = 10): self.block_size = block_size self.total_prefill_time = 0.0 self.tlm_decode_time = 0.0 - self.tlm_verify_time = 0.0 - self.tlm_commit_time = 0.0 self.dlm_decode_time = 0.0 + self.decode_wall_time = 0.0 self.total_accepted_tokens = 0 self.total_rejected_tokens = 0 self.total_generated_tokens = 0 @@ -131,7 +130,7 @@ def _run_spd_core( generated_ids: np.ndarray, metrics: SpecDecodingMetrics, prefill_start_time: float | None = None, - tlm_commit: Callable[[np.ndarray, np.ndarray, int], None] | None = None, + queue_accepted_prefix: Callable[[np.ndarray, np.ndarray, int], None] | None = None, ) -> SpecDecodingMetrics: """Run model-independent DFlash prefill, verification, and cache advancement. @@ -225,6 +224,7 @@ def run_dlm_prefill_block(hidden_states, position_ids, sub_start, sub_len): gen_idx = 0 iteration_count = 0 continue_generation = True + decode_wall_start = time.perf_counter() while gen_idx < generation_limit and iteration_count < max_iterations and continue_generation: iteration_count += 1 @@ -233,9 +233,7 @@ def run_dlm_prefill_block(hidden_states, position_ids, sub_start, sub_len): tlm_decode_start = time.time() tlm_logits, target_hidden = tlm_decode(dlm_candidates, dlm_inputs["position_ids"]) - verify_time = time.time() - tlm_decode_start - metrics.tlm_verify_time += verify_time - metrics.tlm_decode_time += verify_time + metrics.tlm_decode_time += time.time() - tlm_decode_start accepted_length = 0 rejected_flag = False @@ -272,15 +270,9 @@ def run_dlm_prefill_block(hidden_states, position_ids, sub_start, sub_len): if any(tok_id in eos_token_ids for tok_id in this_iter_gen_ids): continue_generation = False - if tlm_commit is not None: - # A commit call reuses the session's output buffers. Preserve the - # verification features that the draft model consumes next. - target_hidden = target_hidden.copy() - commit_start = time.time() - tlm_commit(dlm_candidates, dlm_inputs["position_ids"], accepted_length + 1) - commit_time = time.time() - commit_start - metrics.tlm_commit_time += commit_time - metrics.tlm_decode_time += commit_time + if queue_accepted_prefix is not None: + # Queue the accepted prefix; the next verification call commits it. + queue_accepted_prefix(dlm_candidates, dlm_inputs["position_ids"], accepted_length + 1) if not continue_generation or gen_idx >= generation_limit: break @@ -307,6 +299,7 @@ def run_dlm_prefill_block(hidden_states, position_ids, sub_start, sub_len): dlm_candidates = dlm_outputs["logits"].argmax(axis=-1) metrics.num_total_iters = iteration_count + metrics.decode_wall_time += time.perf_counter() - decode_wall_start return metrics @@ -715,16 +708,29 @@ def build_inputs_qwen3_vl(processor, user_prompt, image=None, enable_thinking=No def get_qwen35_target_seq_len(tlm_session: QAICInferenceSession, block_size: int) -> int: - """Require block-sized verification, independently of prompt prefill.""" + """Require concatenated verification independently of prompt prefill.""" index = tlm_session.binding_index_map["input_ids"] - lengths = {shape[index][1][1] for shape in tlm_session.allowed_shapes} - if not lengths: - lengths = {tlm_session.bindings[index].dims[1]} - if block_size in lengths: - return block_size + selection_index = tlm_session.binding_index_map.get("num_logits_to_keep") + logits_index = tlm_session.binding_index_map.get("logits") + shapes = tlm_session.allowed_shapes or [[(None, binding.dims) for binding in tlm_session.bindings]] + # Single-specialization QPCs can fold away num_logits_to_keep. + lengths = { + shape[index][1][1] + for shape in shapes + if (selection_index is not None and shape[selection_index][1][0] == block_size) + or ( + selection_index is None + and logits_index is not None + and len(shape[logits_index][1]) >= 2 + and shape[logits_index][1][1] == block_size + ) + } + if 2 * block_size in lengths: + return 2 * block_size raise ValueError( - f"Qwen3.5 DFlash requires a {block_size}-token verification specialization " - f"(found {sorted(lengths)}). Recompile the TLM with dflash_block_size={block_size}." + f"Qwen3.5 DFlash requires a {2 * block_size}-token concatenated " + f"verification specialization with num_logits_to_keep={block_size} (found {sorted(lengths)}). " + f"Older two-call QPCs are unsupported. Re-export and recompile the TLM with dflash_block_size={block_size}." ) @@ -757,9 +763,13 @@ def run_spd_inference_qwen3_vl( raise ValueError("Qwen3.5 DFlash requires batch size 1 and a TLM QPC with cache_commit input.") target_seq_len = block_size if is_qwen35: - if block_size < 2 or prompt_chunk_size < block_size or prompt_chunk_size % block_size: - raise ValueError("Qwen3.5 DFlash requires prefill length to be a positive multiple of block_size >= 2.") + if block_size < 2 or prompt_chunk_size < 2 * block_size or prompt_chunk_size % block_size: + raise ValueError( + "Qwen3.5 DFlash requires prefill length to be a multiple of block_size and >= 2 * block_size." + ) target_seq_len = get_qwen35_target_seq_len(tlm_session, block_size) + replay_ids = np.full((batch_size, block_size), tokenizer.pad_token_id, dtype=np.int64) if is_qwen35 else None + replay_positions = np.full((4, batch_size, block_size), -1, dtype=np.int64) if is_qwen35 else None metrics = SpecDecodingMetrics(block_size=block_size) prefill_session = tlm_prefill_session or tlm_session handoff = DFlashTlmHandoff(prefill_session, tlm_session) if tlm_prefill_session is not None else None @@ -833,7 +843,7 @@ def run_spd_inference_qwen3_vl( if is_qwen35: lang_extra["cache_commit"] = np.ones(1, dtype=np.int64) prefill_selection = _dflash_logit_selection_inputs(prefill_session, 1) - decode_selection = _dflash_logit_selection_inputs(tlm_session, target_seq_len) + decode_selection = _dflash_logit_selection_inputs(tlm_session, block_size) prefill_logits = 1 if prefill_selection or tlm_prefill_session is not None else prompt_chunk_size prefill_session.set_buffers({"logits": np.zeros((batch_size, prefill_logits), dtype=np.int32)}) prefill_session.set_buffers( @@ -858,7 +868,7 @@ def tlm_prefill(chunk_start, chunk_end): return outputs["logits"], outputs["hidden_states"] def prepare_decode(): - tlm_session.set_buffers({"logits": np.zeros((batch_size, target_seq_len), dtype=np.int32)}) + tlm_session.set_buffers({"logits": np.zeros((batch_size, block_size), dtype=np.int32)}) tlm_session.set_buffers( {"hidden_states": np.zeros((batch_size, target_seq_len, hidden_size), dtype=np.float32)} ) @@ -867,18 +877,13 @@ def prepare_decode(): if is_qwen35: lang_extra["cache_commit"] = np.zeros(1, dtype=np.int64) - def decode_inputs(draft_ids, decode_position_ids, valid_tokens): - positions = build_decode_position_ids_qwen3_vl(decode_position_ids, rope_delta) - if is_qwen35: - padding = target_seq_len - block_size - draft_ids = np.pad(draft_ids, ((0, 0), (0, padding)), constant_values=tokenizer.pad_token_id) - positions = np.pad(positions, ((0, 0), (0, 0), (0, padding)), constant_values=-1) - positions[:, :, valid_tokens:] = -1 - return {"input_ids": draft_ids, "position_ids": positions, **decode_selection, **lang_extra} - def tlm_decode(draft_ids, decode_position_ids): nonlocal handoff - inputs = decode_inputs(draft_ids, decode_position_ids, block_size) + positions = build_decode_position_ids_qwen3_vl(decode_position_ids, rope_delta) + if is_qwen35: + draft_ids = np.concatenate([replay_ids, draft_ids], axis=-1) + positions = np.concatenate([replay_positions, positions], axis=-1) + inputs = {"input_ids": draft_ids, "position_ids": positions, **decode_selection, **lang_extra} if handoff is not None: inputs.update(handoff.inputs) outputs = tlm_session.run(inputs) @@ -886,15 +891,12 @@ def tlm_decode(draft_ids, decode_position_ids): tlm_session.skip_buffers(handoff.inputs) tlm_session.skip_buffers(vision_outputs.keys()) handoff = None - return outputs["logits"][:, :block_size], outputs["hidden_states"][:, :block_size] + return outputs["logits"][:, :block_size], outputs["hidden_states"][:, -block_size:] - def tlm_commit(draft_ids, decode_position_ids, count): - tlm_session.run( - { - **decode_inputs(draft_ids, decode_position_ids, count), - "cache_commit": np.ones(1, dtype=np.int64), - } - ) + def queue_accepted_prefix(draft_ids, decode_position_ids, count): + replay_ids[:] = draft_ids + replay_positions[:] = build_decode_position_ids_qwen3_vl(decode_position_ids, rope_delta) + replay_positions[:, :, count:] = -1 eos_token_ids = {tokenizer.eos_token_id} if tokenizer.eos_token_id is not None else set() return _run_spd_core( @@ -915,5 +917,5 @@ def tlm_commit(draft_ids, decode_position_ids, count): generated_ids=generated_ids, metrics=metrics, prefill_start_time=prefill_start, - tlm_commit=tlm_commit if is_qwen35 else None, + queue_accepted_prefix=queue_accepted_prefix if is_qwen35 else None, ) diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index d5049b4475..b58bce96eb 100644 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -1620,11 +1620,9 @@ def __init__( # ---Sampling--- if self.dflash_tlm and qaic_config.get("speculative_model_type") == "target": self.lang_model.hash_params["dflash_output_logits"] = True - if self.dflash_tlm and self.model.config.model_type == "qwen3_5_moe": - self.lang_model.hash_params["dflash_cache_protocol"] = "qwen35_moe_dynamic_replay_v1" - self.lang_model.hash_params["dflash_chunk_size"] = 16 - if self.dflash_tlm and self.model.config.model_type == "qwen3_5": - self.lang_model.hash_params["dflash_cache_protocol"] = "qwen35_dynamic_replay_v1" + if self.dflash_tlm and self.model.config.model_type in {"qwen3_5", "qwen3_5_moe"}: + prefix = "qwen35_moe" if self.model.config.model_type == "qwen3_5_moe" else "qwen35" + self.lang_model.hash_params["dflash_cache_protocol"] = f"{prefix}_fused_replay_v1" self.lang_model.hash_params["dflash_chunk_size"] = 16 if self.dflash_tlm and self.model.config.model_type == "gemma4": self.lang_model.hash_params["dflash_cache_protocol"] = "gemma4_block_window_v1" diff --git a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py index c710d116ac..ccc0c86125 100644 --- a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -861,7 +861,7 @@ def torch_chunk_gated_delta_rule_qeff( ones_lower=None, eye=None, use_matrix_solve=False, - use_dynamic_chunk_loop=False, + dflash_commit_final=None, ): tensor_device = query.device if mask_causal is not None: @@ -962,9 +962,17 @@ def torch_chunk_gated_delta_rule_qeff( if initial_state is None else initial_state.to(value) ) - if use_dynamic_chunk_loop: + if dflash_commit_final is not None: core_attn_out, last_recurrent_state = dflash_chunk_gated_delta_scan( - query, key, value, k_cumdecay, g, decay_mask, mask_strict, last_recurrent_state + query, + key, + value, + k_cumdecay, + g, + decay_mask, + mask_strict, + last_recurrent_state, + dflash_commit_final, ) else: core_attn_out = torch.zeros_like(value) @@ -1021,6 +1029,7 @@ def forward( gdn_num_head_blocks: int = 1, dflash_chunk_size: int | None = None, return_cache_states: bool = False, + dflash_commit_final: torch.Tensor | None = None, ): batch_size, seq_len, _ = hidden_states.shape @@ -1084,14 +1093,35 @@ def forward( recurrent_keep = (~recurrent_reset_mask).to(recurrent_state.dtype) recurrent_state = recurrent_state * recurrent_keep - mixed_qkv, new_conv_state = qeff_torch_causal_conv1d_update( - mixed_qkv, - conv_state, - self.conv1d.weight.squeeze(1), - position_ids, - self.conv1d.bias, - dflash=dflash_chunk_size is not None, - ) + if dflash_commit_final is not None: + # Replay padding must not become history for the verification half. + prefix_qkv, prefix_conv_state = qeff_torch_causal_conv1d_update( + mixed_qkv[:, :, :-dflash_chunk_size], + conv_state, + self.conv1d.weight.squeeze(1), + position_ids[:, :, :-dflash_chunk_size], + self.conv1d.bias, + dflash=True, + ) + suffix_qkv, final_conv_state = qeff_torch_causal_conv1d_update( + mixed_qkv[:, :, -dflash_chunk_size:], + prefix_conv_state, + self.conv1d.weight.squeeze(1), + position_ids[:, :, -dflash_chunk_size:], + self.conv1d.bias, + dflash=True, + ) + mixed_qkv = torch.cat([prefix_qkv, suffix_qkv], dim=-1) + new_conv_state = torch.where(dflash_commit_final.bool(), final_conv_state, prefix_conv_state) + else: + mixed_qkv, new_conv_state = qeff_torch_causal_conv1d_update( + mixed_qkv, + conv_state, + self.conv1d.weight.squeeze(1), + position_ids, + self.conv1d.bias, + dflash=dflash_chunk_size is not None, + ) if batch_index is not None: if conv_state_all.ndim == 4: conv_state_all_flat = conv_state_all.reshape( @@ -1168,7 +1198,7 @@ def forward( ones_lower=self._ones_lower[:dflash_chunk_size, :dflash_chunk_size], eye=self._eye[:dflash_chunk_size, :dflash_chunk_size], use_matrix_solve=getattr(self, "dflash_matrix_solve", True), - use_dynamic_chunk_loop=True, + dflash_commit_final=dflash_commit_final, ) elif batch_fold: # Folded decode is exported with a static one-token sequence. @@ -1293,6 +1323,7 @@ def forward( ) -> torch.FloatTensor: del use_cache return_cache_states = kwargs.pop("return_cache_states", False) + dflash_commit_final = kwargs.pop("dflash_commit_final", None) residual = hidden_states hidden_states = self.input_layernorm(hidden_states) @@ -1308,6 +1339,7 @@ def forward( gdn_num_head_blocks=gdn_num_head_blocks, dflash_chunk_size=dflash_chunk_size, return_cache_states=return_cache_states, + dflash_commit_final=dflash_commit_final, ) if return_cache_states: hidden_states, conv_state, recurrent_state = hidden_states @@ -1454,7 +1486,8 @@ def forward( raise ValueError("Qwen3.5 DFlash requires the full target model; layerwise execution is unsupported.") if cache_commit is None: raise ValueError( - "Qwen3.5 DFlash requires cache_commit: 0 for verification, 1 for accepted-prefix replay." + "Qwen3.5 DFlash requires cache_commit: 1 commits all valid positions; " + "0 commits only the accepted-prefix half of the concatenated decode input." ) if batch_index is not None or batch_fold: raise ValueError("Qwen3.5 DFlash currently supports non-continuous batch size 1 only.") @@ -1492,8 +1525,6 @@ def forward( conv_reset = reset if conv.ndim == 4 else reset.squeeze(-1) past_key_values.conv_states[i] = torch.where(conv_reset, torch.zeros_like(conv), conv) past_key_values.recurrent_states[i] = torch.where(reset, torch.zeros_like(recurrent), recurrent) - committed_conv = list(past_key_values.conv_states) - committed_recurrent = list(past_key_values.recurrent_states) rotary_position_ids = ( position_ids[1:] if position_ids.ndim == 3 and position_ids.shape[0] == 4 else position_ids ) @@ -1545,6 +1576,7 @@ def forward( if target_layer_ids else None, return_cache_states=export_cache_states, + dflash_commit_final=cache_commit if target_layer_ids else None, **kwargs, ) if export_cache_states: @@ -1560,24 +1592,6 @@ def forward( if output_hidden_states: all_hidden_states += (hidden_states,) - if target_layer_ids: - # Verification may overwrite attention KV slots, but recurrent states - # cannot be rolled back by positions. Only an accepted-prefix pass commits them. - for i, kind in enumerate(self.config.layer_types): - if kind == "linear_attention": - conv_state, recurrent_state = ( - exported_cache_states[i] - if export_cache_states - else (past_key_values.conv_states[i], past_key_values.recurrent_states[i]) - ) - conv_state = torch.where(cache_commit.bool(), conv_state, committed_conv[i]) - recurrent_state = torch.where(cache_commit.bool(), recurrent_state, committed_recurrent[i]) - if export_cache_states: - exported_cache_states[i] = (conv_state, recurrent_state) - else: - past_key_values.conv_states[i] = conv_state - past_key_values.recurrent_states[i] = recurrent_state - if export_cache_states: if is_layerwise_active(): past_key_values = (exported_cache_states[QEffQwen3_5TextModel._start],) @@ -2371,7 +2385,9 @@ def _build_lang_spec(seq_len_val, comp_ctx_len=None): spec["full_batch_exec_size"] = full_batch_size return spec - decode_seq_len = chunk_size if getattr(language_model, "target_layer_ids", None) else 1 + decode_seq_len = 2 * chunk_size if getattr(language_model, "target_layer_ids", None) else 1 + if getattr(language_model, "target_layer_ids", None) and prefill_seq_len < decode_seq_len: + raise ValueError("Qwen3.5 DFlash requires prefill_seq_len >= 2 * block_size.") lang = [] if comp_ctx_lengths_prefill is not None or comp_ctx_lengths_decode is not None: for comp_ctx in comp_ctx_lengths_prefill or []: @@ -2382,14 +2398,14 @@ def _build_lang_spec(seq_len_val, comp_ctx_len=None): for comp_ctx in comp_ctx_lengths_decode or []: decode_spec = _build_lang_spec(decode_seq_len, comp_ctx_len=comp_ctx) if getattr(language_model, "target_layer_ids", None): - decode_spec["num_logits_to_keep"] = decode_seq_len + decode_spec["num_logits_to_keep"] = chunk_size lang.append(decode_spec) else: prefill_spec = _build_lang_spec(prefill_seq_len) decode_spec = _build_lang_spec(decode_seq_len) if getattr(language_model, "target_layer_ids", None): prefill_spec["num_logits_to_keep"] = 1 - decode_spec["num_logits_to_keep"] = decode_seq_len + decode_spec["num_logits_to_keep"] = chunk_size lang = [prefill_spec, decode_spec] if kv_offload: @@ -2467,7 +2483,7 @@ def get_dummy_inputs( raise ValueError("Qwen3.5 DFlash requires prefill_seq_len >= 2 for the exported gated-delta graph.") # Keep tracing small; the scripted scan preserves the dynamic chunk count. chunk_size = getattr(self.model.language_model, "dflash_chunk_size", QWEN3_5_DFLASH_CHUNK_SIZE) - dummy_seq_len = min(dummy_seq_len, chunk_size) + dummy_seq_len = 2 * chunk_size bs = kwargs.get("batch_size", constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE) if kwargs.get("weight_free", False): bs = max(2, bs) diff --git a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index c9aab0a8a9..40041b0420 100644 --- a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -880,7 +880,7 @@ def torch_chunk_gated_delta_rule_qeff( ones_lower=None, eye=None, use_matrix_solve=False, - use_dynamic_chunk_loop=False, + dflash_commit_final=None, ): tensor_device = query.device if mask_causal is not None: @@ -985,9 +985,17 @@ def torch_chunk_gated_delta_rule_qeff( if initial_state is None else initial_state.to(value) ) - if use_dynamic_chunk_loop: + if dflash_commit_final is not None: core_attn_out, last_recurrent_state = dflash_chunk_gated_delta_scan( - query, key, value, k_cumdecay, g, decay_mask, mask_strict, last_recurrent_state + query, + key, + value, + k_cumdecay, + g, + decay_mask, + mask_strict, + last_recurrent_state, + dflash_commit_final, ) else: core_attn_out = torch.zeros_like(value) @@ -1043,6 +1051,7 @@ def forward( gdn_num_head_blocks: int = 1, dflash_chunk_size: int | None = None, return_cache_states: bool = False, + dflash_commit_final: torch.Tensor | None = None, ): batch_size, seq_len, _ = hidden_states.shape @@ -1107,14 +1116,35 @@ def forward( recurrent_keep = (~recurrent_reset_mask).to(recurrent_state.dtype) recurrent_state = recurrent_state * recurrent_keep - mixed_qkv, new_conv_state = qeff_torch_causal_conv1d_update( - mixed_qkv, - conv_state, - self.conv1d.weight.squeeze(1), - position_ids, - self.conv1d.bias, - dflash=dflash_chunk_size is not None, - ) + if dflash_commit_final is not None: + # Replay padding must not become history for the verification half. + prefix_qkv, prefix_conv_state = qeff_torch_causal_conv1d_update( + mixed_qkv[:, :, :-dflash_chunk_size], + conv_state, + self.conv1d.weight.squeeze(1), + position_ids[:, :, :-dflash_chunk_size], + self.conv1d.bias, + dflash=True, + ) + suffix_qkv, final_conv_state = qeff_torch_causal_conv1d_update( + mixed_qkv[:, :, -dflash_chunk_size:], + prefix_conv_state, + self.conv1d.weight.squeeze(1), + position_ids[:, :, -dflash_chunk_size:], + self.conv1d.bias, + dflash=True, + ) + mixed_qkv = torch.cat([prefix_qkv, suffix_qkv], dim=-1) + new_conv_state = torch.where(dflash_commit_final.bool(), final_conv_state, prefix_conv_state) + else: + mixed_qkv, new_conv_state = qeff_torch_causal_conv1d_update( + mixed_qkv, + conv_state, + self.conv1d.weight.squeeze(1), + position_ids, + self.conv1d.bias, + dflash=dflash_chunk_size is not None, + ) if state_batch_index is not None: if conv_state_all.ndim == 4: conv_state_all_flat = conv_state_all.reshape( @@ -1190,7 +1220,7 @@ def forward( ones_lower=self._ones_lower[:dflash_chunk_size, :dflash_chunk_size], eye=self._eye[:dflash_chunk_size, :dflash_chunk_size], use_matrix_solve=getattr(self, "dflash_matrix_solve", True), - use_dynamic_chunk_loop=True, + dflash_commit_final=dflash_commit_final, ) elif batch_fold: # Folded decode is exported with a static one-token sequence. @@ -1317,6 +1347,7 @@ def forward( ) -> torch.FloatTensor: del use_cache return_cache_states = kwargs.pop("return_cache_states", False) + dflash_commit_final = kwargs.pop("dflash_commit_final", None) residual = hidden_states hidden_states = self.input_layernorm(hidden_states) @@ -1332,6 +1363,7 @@ def forward( gdn_num_head_blocks=gdn_num_head_blocks, dflash_chunk_size=dflash_chunk_size, return_cache_states=return_cache_states, + dflash_commit_final=dflash_commit_final, ) if return_cache_states: hidden_states, conv_state, recurrent_state = hidden_states @@ -1483,7 +1515,8 @@ def forward( ) if cache_commit is None: raise ValueError( - "Qwen3.5-MoE DFlash requires cache_commit: 0 for verification, 1 for accepted-prefix replay." + "Qwen3.5-MoE DFlash requires cache_commit: 1 commits all valid positions; " + "0 commits only the accepted-prefix half of the concatenated decode input." ) if batch_index is not None or batch_fold: raise ValueError("Qwen3.5-MoE DFlash currently supports non-continuous batch size 1 only.") @@ -1521,8 +1554,6 @@ def forward( conv_reset = reset if conv.ndim == 4 else reset.squeeze(-1) past_key_values.conv_states[i] = torch.where(conv_reset, torch.zeros_like(conv), conv) past_key_values.recurrent_states[i] = torch.where(reset, torch.zeros_like(recurrent), recurrent) - committed_conv = list(past_key_values.conv_states) - committed_recurrent = list(past_key_values.recurrent_states) rotary_position_ids = ( position_ids[1:] if position_ids.ndim == 3 and position_ids.shape[0] == 4 else position_ids ) @@ -1583,6 +1614,7 @@ def forward( if target_layer_ids else None, return_cache_states=export_cache_states, + dflash_commit_final=cache_commit if target_layer_ids else None, **kwargs, ) if export_cache_states: @@ -1598,23 +1630,6 @@ def forward( if output_hidden_states: all_hidden_states += (hidden_states,) - if target_layer_ids: - # Attention KV uses positional overwrite; recurrent states require replay. - for i, kind in enumerate(self.config.layer_types): - if kind == "linear_attention": - conv_state, recurrent_state = ( - exported_cache_states[i] - if export_cache_states - else (past_key_values.conv_states[i], past_key_values.recurrent_states[i]) - ) - conv_state = torch.where(cache_commit.bool(), conv_state, committed_conv[i]) - recurrent_state = torch.where(cache_commit.bool(), recurrent_state, committed_recurrent[i]) - if export_cache_states: - exported_cache_states[i] = (conv_state, recurrent_state) - else: - past_key_values.conv_states[i] = conv_state - past_key_values.recurrent_states[i] = recurrent_state - if export_cache_states: if is_layerwise_active(): past_key_values = (exported_cache_states[QEffQwen3_5MoeTextModel._start],) @@ -2540,7 +2555,9 @@ def _build_lang_spec(seq_len_val, comp_ctx_len=None): spec["full_batch_exec_size"] = full_batch_size return spec - decode_seq_len = chunk_size if getattr(language_model, "target_layer_ids", None) else 1 + decode_seq_len = 2 * chunk_size if getattr(language_model, "target_layer_ids", None) else 1 + if getattr(language_model, "target_layer_ids", None) and prefill_seq_len < decode_seq_len: + raise ValueError("Qwen3.5 DFlash requires prefill_seq_len >= 2 * block_size.") lang = [] if comp_ctx_lengths_prefill is not None or comp_ctx_lengths_decode is not None: for comp_ctx in comp_ctx_lengths_prefill or []: @@ -2551,14 +2568,14 @@ def _build_lang_spec(seq_len_val, comp_ctx_len=None): for comp_ctx in comp_ctx_lengths_decode or []: decode_spec = _build_lang_spec(decode_seq_len, comp_ctx_len=comp_ctx) if getattr(language_model, "target_layer_ids", None): - decode_spec["num_logits_to_keep"] = decode_seq_len + decode_spec["num_logits_to_keep"] = chunk_size lang.append(decode_spec) else: prefill_spec = _build_lang_spec(prefill_seq_len) decode_spec = _build_lang_spec(decode_seq_len) if getattr(language_model, "target_layer_ids", None): prefill_spec["num_logits_to_keep"] = 1 - decode_spec["num_logits_to_keep"] = decode_seq_len + decode_spec["num_logits_to_keep"] = chunk_size lang = [prefill_spec, decode_spec] if kv_offload: @@ -2640,7 +2657,7 @@ def get_dummy_inputs( raise ValueError("Qwen3.5-MoE DFlash requires prefill_seq_len >= 2 for the exported gated-delta graph.") # A small dummy still exports the scripted loop's dynamic trip count. chunk_size = getattr(self.model.language_model, "dflash_chunk_size", QWEN3_5_MOE_DFLASH_CHUNK_SIZE) - dummy_seq_len = min(dummy_seq_len, chunk_size) + dummy_seq_len = 2 * chunk_size bs = kwargs.get("batch_size", constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE) if kwargs.get("weight_free", False): bs = max(2, bs) diff --git a/QEfficient/transformers/spd/dflash.py b/QEfficient/transformers/spd/dflash.py index 617735d02d..7a7a385552 100644 --- a/QEfficient/transformers/spd/dflash.py +++ b/QEfficient/transformers/spd/dflash.py @@ -18,11 +18,14 @@ def dflash_chunk_gated_delta_scan( decay_mask: torch.Tensor, mask_strict: torch.Tensor, state: torch.Tensor, + commit_final: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """Preserve the recurrent chunk loop in ONNX for sequence specialization.""" + """Commit all prefill chunks, or retain the state before the verification chunk.""" output = torch.zeros_like(value) batch, heads, chunk = query.shape[0], query.shape[1], query.shape[3] + prefix_state = state for i in range(query.shape[2]): + prefix_state = state q_i = query[:, :, i : i + 1].reshape(batch, heads, chunk, query.shape[4]) k_i = key[:, :, i : i + 1].reshape(batch, heads, chunk, key.shape[4]) v_i = value[:, :, i : i + 1].reshape(batch, heads, chunk, value.shape[4]) @@ -38,7 +41,7 @@ def dflash_chunk_gated_delta_scan( state = ( state * last_g.exp().unsqueeze(-1) + (k_i * (last_g - g_i).exp().unsqueeze(-1)).transpose(-1, -2) @ v_new ) - return output, state + return output, torch.where(commit_final.to(torch.bool), state, prefix_state) def compute_dflash_target_hidden_states( diff --git a/examples/performance/dflash/README.md b/examples/performance/dflash/README.md index acf54e8f47..649fcfefd1 100644 --- a/examples/performance/dflash/README.md +++ b/examples/performance/dflash/README.md @@ -53,7 +53,8 @@ Use `--no-disagg-tlm` to compile or run the previous combined TLM QPC. Dense models keep the combined-QPC default. Split TLM currently supports BS=1. After the last prefill chunk, KV state (plus convolution and recurrent state for Qwen3.5-MoE) is transferred through host buffers to the verification QPC -on its first call. Later verification/replay calls keep state on-device. +on its first call. Later decode calls keep state on-device; Qwen3.5 dense/MoE +targets concatenate the accepted-prefix replay and new verification block in one call. To reuse a split pair, pass `--tlm-prefill-qpc /path/to/prefill/qpc` and `--tlm_qpc /path/to/verification/qpc`. Prefill uses `--tlm_devices` unless @@ -126,12 +127,30 @@ python basic_inference_vision.py --model_name Qwen3.5-4B \ ``` Use `--prompt "..."` instead of `--image --image_prompt ...` for text-only input. -Prefill length can be any positive multiple of the draft block size; verification and replay remain block-sized. - -**Recompilation required:** Qwen3.5 DFlash TLM ONNX/QPC artifacts created before -the variable-length prefill and block-sized verification/replay update are incompatible. -Re-export and recompile the TLM; do not reuse those old paths with `--tlm_qpc`. -To let this example rebuild the VLM artifacts, omit both `--tlm_qpc` and `--vision_qpc`. +Qwen3.5 dense and MoE targets use concatenated replay and verification. Each decode call takes +`2 * block_size` positions: the previous accepted prefix (padded with position IDs +of `-1`), followed by the next draft block. It commits the prefix's convolution +and recurrent states while verifying the new block without committing it. +Only the new block's token predictions and projected features are used by DFlash. +The first decode call has an entirely masked replay half. After each verification, +the runtime queues the seed plus accepted draft tokens for the next call; the new +replacement/bonus seed belongs to the verification half, not the replay half. +The example reports this combined call as `TLM Replay+Verify`, with no separate +device replay time. SPD wall-time throughput includes DLM calls and host orchestration; +the rate based on time spent in model calls is shown separately, not as device-only profiling. +Prefill length must be a multiple of the draft block size and at least twice that size. +For a block size of 16, decode takes 32 positions and returns 16 token predictions. + +**Recompilation required to enable fused replay:** Re-export and recompile the TLM; +older TLM ONNX/QPC artifacts cannot gain fused replay without rebuilding. The language +export hash distinguishes the new cache protocol; existing artifacts are not deleted. +Omit both `--tlm_qpc` and `--vision_qpc` to let this example select/compile compatible artifacts. +The existing vision and DLM QPCs do not need to change for fused replay. + +The previous two-call TLM QPCs are no longer supported; the example rejects them +with a recompilation message. There is no two-call fallback or mode switch. +MXINT8/MXFP6 settings are unchanged. External runtimes, including vLLM, must support +the new fused input/state contract before using these new TLM QPCs. --- diff --git a/examples/performance/dflash/basic_inference_vision.py b/examples/performance/dflash/basic_inference_vision.py index fb6ca1a1af..255c1f5836 100644 --- a/examples/performance/dflash/basic_inference_vision.py +++ b/examples/performance/dflash/basic_inference_vision.py @@ -401,12 +401,18 @@ def main(): print(f" {'TLM Rounds (round/s)':<30} {round_rate:>6.1f}") else: print(f" {'TLM Throughput (tok/s)':<30} {metrics.tlm_tok_rate():>6.1f}") - if metrics.tlm_commit_time > 0 and metrics.tlm_verify_time > 0 and metrics.num_total_iters > 0: - print(f" {'TLM Verify (ms/call)':<30} {1000 * metrics.tlm_verify_time / metrics.num_total_iters:>6.1f}") - print(f" {'TLM Replay (ms/call)':<30} {1000 * metrics.tlm_commit_time / metrics.num_total_iters:>6.1f}") - verified_rate = metrics.block_size * metrics.num_total_iters / metrics.tlm_verify_time + if metrics.tlm_decode_time > 0 and metrics.num_total_iters > 0: + label = ( + "TLM Replay+Verify (s/call)" + if tlm_config.model_type in {"qwen3_5", "qwen3_5_moe"} + else "TLM Verify (s/call)" + ) + print(f" {label:<30} {metrics.tlm_decode_time / metrics.num_total_iters:>8.4f}") + verified_rate = metrics.block_size * metrics.num_total_iters / metrics.tlm_decode_time print(f" {'TLM Verified positions/s':<30} {verified_rate:>6.1f}") - print(f" {'SPD Decode Speed (tok/s)':<30} {metrics.spd_tok_rate():>6.1f}") + wall_speed = metrics.total_generated_tokens / metrics.decode_wall_time if metrics.decode_wall_time else 0.0 + print(f" {'SPD Wall Speed (tok/s)':<30} {wall_speed:>6.1f}") + print(f" {'SPD Model-call Speed (tok/s)':<30} {metrics.spd_tok_rate():>6.1f}") print(f" {'Generated tokens':<30} {metrics.total_generated_tokens:>6}") print(f" {'Iterations':<30} {metrics.num_total_iters:>6}") print(f" {'Prefill time (s)':<30} {metrics.total_prefill_time:>6.3f}") diff --git a/examples/performance/dflash/utils.py b/examples/performance/dflash/utils.py index 2503378aac..df603370a5 100644 --- a/examples/performance/dflash/utils.py +++ b/examples/performance/dflash/utils.py @@ -21,6 +21,7 @@ QAICInferenceSession, is_retained_state_name, ) +from QEfficient.generation.dflash_generation import get_qwen35_target_seq_len from QEfficient.utils.logging_utils import QEFFLogger logger = QEFFLogger.get_logger("INFRA") @@ -86,6 +87,12 @@ def get_spd_prompt_chunk_size(tlm_session: QAICInferenceSession) -> int: def validate_spd_decode_specialization(tlm_session: QAICInferenceSession, block_size: int) -> None: """Require a TLM specialization that verifies one complete DFlash block.""" + if "cache_commit" in getattr(tlm_session, "input_names", ()): + seq_len = get_qwen35_target_seq_len(tlm_session, block_size) + logger.info( + f"TLM uses concatenated replay + verification: {seq_len} input positions, {block_size} predictions." + ) + return input_index = tlm_session.binding_index_map["input_ids"] sequence_lengths = sorted({shape[input_index][1][1] for shape in tlm_session.allowed_shapes}) if not sequence_lengths: @@ -455,9 +462,11 @@ def compile_qwen3vl_vlm_qpcs( if is_qwen35 or is_qwen35_moe: if block_size < 2 or block_size & (block_size - 1): raise ValueError("Qwen3.5 DFlash requires a power-of-two draft block_size >= 2.") - if prefill_seq_len < block_size or prefill_seq_len % block_size: - raise ValueError(f"Qwen3.5 DFlash requires prefill_seq_len to be a positive multiple of {block_size}.") - logger.warning("Qwen3.5 DFlash uses accepted-prefix replay; target timing includes the extra commit pass.") + if prefill_seq_len < 2 * block_size or prefill_seq_len % block_size: + raise ValueError( + f"Qwen3.5 DFlash requires prefill_seq_len to be a multiple of {block_size} and at least {2 * block_size}." + ) + logger.info("Qwen3.5 DFlash combines accepted-prefix replay and verification in one TLM call.") tlm_qeff = QEFFAutoModelForImageTextToText.from_pretrained( tlm_repo, config=config, @@ -469,8 +478,7 @@ def compile_qwen3vl_vlm_qpcs( qaic_config=qaic_config, ) - # DFlashTLMTransform attached fc/hidden_norm but left them random; copy in the real - # DFlash weights and scale fc into fp16 range. + # Apply the DLM projection weights with the same fp16-range scaling as the transform. text_model = tlm_qeff.lang_model.model.language_model text_model.fc.weight.data.copy_(state_dict["fc.weight"].to(torch.float32)) with torch.no_grad(): @@ -493,7 +501,7 @@ def compile_qwen3vl_vlm_qpcs( "dflash_block_size": block_size, } if is_qwen35 or is_qwen35_moe: - # Hybrid verification and accepted-prefix replay use chunk attention. + # Concatenated replay and verification use chunk attention. compile_kwargs["prefill_only"] = True if is_qwen35_moe: compile_kwargs["qaic_config"] = qaic_config diff --git a/tests/transformers/spd/test_dflash_spd_inference.py b/tests/transformers/spd/test_dflash_spd_inference.py index b991809836..1ba7bd9a2b 100644 --- a/tests/transformers/spd/test_dflash_spd_inference.py +++ b/tests/transformers/spd/test_dflash_spd_inference.py @@ -85,7 +85,7 @@ def stop_loading(*args, **kwargs): "Qwen/Qwen3.6-27B", "z-lab/Qwen3.6-27B-DFlash", ctx_len=ctx_len, - prefill_seq_len=16, + prefill_seq_len=32, num_cores=8, num_devices=4, ) @@ -143,7 +143,7 @@ def test_dflash_decode_specialization_accepts_fixed_and_dynamic_shapes(sequence_ @pytest.mark.parametrize("prompt_chunk_size", [8, 31, 127]) def test_qwen3_5_runtime_rejects_unsupported_prefill_qpc(prompt_chunk_size): - with pytest.raises(ValueError, match="positive multiple of block_size"): + with pytest.raises(ValueError, match="multiple of block_size"): run_spd_inference_qwen3_vl( prompt_text="test", tokenizer=None, @@ -160,15 +160,17 @@ def test_qwen3_5_runtime_rejects_unsupported_prefill_qpc(prompt_chunk_size): @pytest.mark.parametrize( "model_type,prompt_chunk_size,target_seq_len", [ - ("qwen3_5", 16, 16), - ("qwen3_5", 64, 16), - ("qwen3_5", 128, 16), - ("qwen3_5", 256, 16), - ("qwen3_5_moe", 16, 16), - ("qwen3_5_moe", 128, 16), + ("qwen3_5", 32, 32), + ("qwen3_5", 64, 32), + ("qwen3_5", 128, 32), + ("qwen3_5", 256, 32), + ("qwen3_5_moe", 32, 32), + ("qwen3_5_moe", 128, 32), ], ) -def test_qwen3_5_runtime_pads_verification_and_commit(monkeypatch, model_type, prompt_chunk_size, target_seq_len): +def test_qwen3_5_runtime_concatenates_prefix_without_second_call( + monkeypatch, model_type, prompt_chunk_size, target_seq_len +): import torch from QEfficient.generation import dflash_generation as generation @@ -179,16 +181,19 @@ def test_qwen3_5_runtime_pads_verification_and_commit(monkeypatch, model_type, p def run(feeds): calls.append({key: value.copy() for key, value in feeds.items()}) return { - "logits": np.zeros((1, target_seq_len), dtype=np.int32), - "hidden_states": np.zeros((1, target_seq_len, 32), dtype=np.float32), + "logits": np.zeros((1, feeds["num_logits_to_keep"].shape[0]), dtype=np.int32), + "hidden_states": np.broadcast_to(np.arange(target_seq_len)[None, :, None], (1, target_seq_len, 32)), "image_idx_output": np.zeros((1, 1), dtype=np.int64), } session = SimpleNamespace( input_names=["cache_commit", "vision_embeds", "num_logits_to_keep"], - binding_index_map={"vision_embeds": 0, "input_ids": 1}, + binding_index_map={"vision_embeds": 0, "input_ids": 1, "num_logits_to_keep": 2}, bindings=[SimpleNamespace(type=0, dims=(1, 1, 32)), SimpleNamespace(dims=(1, prompt_chunk_size))], - allowed_shapes=[[(0, [1, 1, 32]), (8, [1, length])] for length in {prompt_chunk_size, target_seq_len}], + allowed_shapes=[ + [(0, [1, 1, 32]), (8, [1, length]), (8, [keep, 1])] + for length, keep in [(prompt_chunk_size, 1), (target_seq_len, 16)] + ], aic_to_np_dtype_mapping={0: np.dtype(np.float32)}, run=run, set_buffers=buffers.append, @@ -218,16 +223,22 @@ def core(**kwargs): assert hidden.shape == (1, 16, 32) verify = calls[-1] assert verify["input_ids"].shape == (1, target_seq_len) - np.testing.assert_array_equal(verify["position_ids"][0, :, :16], positions) - assert np.all(verify["position_ids"][..., 16:] == -1) + np.testing.assert_array_equal(verify["position_ids"][0, :, 16:], positions) + assert np.all(verify["position_ids"][..., :16] == -1) + np.testing.assert_array_equal(hidden[0, :, 0], np.arange(16, 32)) assert verify["cache_commit"].item() == 0 assert verify["num_logits_to_keep"].shape == (16, 1) - kwargs["tlm_commit"](ids, positions, 2) - commit = calls[-1] - assert commit["cache_commit"].item() == 1 - assert commit["num_logits_to_keep"].shape == (16, 1) - np.testing.assert_array_equal(commit["position_ids"][0, :, :2], positions[:, :2]) - assert np.all(commit["position_ids"][..., 2:] == -1) + call_count = len(calls) + kwargs["queue_accepted_prefix"](ids, positions, 2) + assert len(calls) == call_count + kwargs["tlm_decode"](ids + 1, positions + 2) + combined = calls[-1] + assert len(calls) == call_count + 1 + assert combined["cache_commit"].item() == 0 + np.testing.assert_array_equal(combined["input_ids"], np.concatenate([ids, ids + 1], axis=1)) + np.testing.assert_array_equal(combined["position_ids"][0, :, :2], positions[:, :2]) + assert np.all(combined["position_ids"][..., 2:16] == -1) + np.testing.assert_array_equal(combined["position_ids"][0, :, 16:], positions + 2) return kwargs["metrics"] monkeypatch.setattr(generation, "_run_spd_core", core) @@ -245,7 +256,7 @@ def core(**kwargs): block_size=16, hidden_size=32, ) - assert [item["logits"].shape[1] for item in buffers if "logits" in item] == [1, target_seq_len] + assert [item["logits"].shape[1] for item in buffers if "logits" in item] == [1, 16] @pytest.mark.parametrize("with_image", [False, True]) @@ -423,7 +434,7 @@ def test_dflash_acceptance_rate_uses_accepted_tokens(): @pytest.mark.parametrize("accepted", [0, 1, 3]) -def test_dflash_commit_counts_seed_and_preserves_verification_features(accepted): +def test_dflash_queued_prefix_counts_seed_and_preserves_verification_features(accepted): commits = [] hidden = np.ones((1, BLOCK_SIZE, HIDDEN_SIZE), dtype=np.float32) @@ -440,9 +451,8 @@ def verify(ids, positions): logits[:, accepted] = 10 return logits, hidden - def commit(ids, positions, count): + def queue_prefix(ids, positions, count): commits.append(count) - hidden.fill(99) # Simulate session output-buffer reuse by the replay call. metrics = _run_spd_core( input_ids=np.array([[3, 4, 0, 0]]), @@ -450,7 +460,7 @@ def commit(ids, positions, count): tlm_prefill=lambda start, end: (np.array([[0, 5, 0, 0]]), hidden), prepare_decode=lambda: None, tlm_decode=verify, - tlm_commit=commit, + queue_accepted_prefix=queue_prefix, dlm_session=Draft(), mask_token_id=MASK_TOKEN_ID, vocab_size=VOCAB_SIZE, @@ -464,7 +474,7 @@ def commit(ids, positions, count): metrics=SpecDecodingMetrics(block_size=BLOCK_SIZE), ) assert commits == [accepted + 1, accepted + 1] - assert metrics.tlm_decode_time == pytest.approx(metrics.tlm_verify_time + metrics.tlm_commit_time) + assert metrics.tlm_decode_time > 0 def test_dflash_generation_len_caps_partial_block(): diff --git a/tests/unit_test/transforms/test_dflash_transform.py b/tests/unit_test/transforms/test_dflash_transform.py index 531093ddda..d9fd6d3b81 100644 --- a/tests/unit_test/transforms/test_dflash_transform.py +++ b/tests/unit_test/transforms/test_dflash_transform.py @@ -95,38 +95,85 @@ def make_tiny_qwen3_5_vlm(moe=False): return model_cls(config).eval() +@pytest.mark.parametrize("moe", [False, True]) +@pytest.mark.parametrize("seq_len,num_logits", [(32, 16), (128, 1)]) +def test_qwen3_5_dflash_vllm_returns_vocabulary_logits(moe, seq_len, num_logits): + from QEfficient import QEFFAutoModelForImageTextToText + + torch.manual_seed(23) + hf_model = make_tiny_qwen3_5_vlm(moe=moe) + hf_model.config.text_config.max_position_embeddings = 256 + qaic_config = {"target_layer_ids": [1, 3], "speculative_model_type": "target"} + qeff = QEFFAutoModelForImageTextToText(hf_model, kv_offload=True, qaic_config=qaic_config) + if moe: + qeff.transform( + ctx_len=256, + seq_len=128, + bs=1, + num_devices=1, + aic_num_cores=4, + qaic_config=qaic_config, + prefill_only=True, + prefill_seq_len=128, + enable_chunking=True, + ) + wrapper = qeff.lang_model.model + inputs = qeff.model.get_dummy_inputs(kv_offload=True, batch_size=1)["lang"] + inputs["input_ids"] = torch.randint(10, VOCAB_SIZE, (1, seq_len)) + inputs["input_ids"][:, :2] = hf_model.config.image_token_id + inputs["vision_embeds"].normal_() + inputs["position_ids"] = torch.arange(seq_len).view(1, 1, -1).repeat(4, 1, 1) + inputs["num_logits_to_keep"] = torch.zeros((num_logits, 1), dtype=torch.int64) + for index, kind in enumerate(hf_model.config.text_config.layer_types): + if kind == "full_attention": + inputs["past_key_values"][index] = [ + torch.zeros((*state.shape[:2], 256, state.shape[-1])) for state in inputs["past_key_values"][index] + ] + with torch.no_grad(): + actual = wrapper(**copy.deepcopy(inputs)) + wrapper.qaic_config = {"target_layer_ids": [1, 3]} + greedy = wrapper(**copy.deepcopy(inputs)) + assert actual[0].shape == (1, num_logits, VOCAB_SIZE) + assert torch.isfinite(actual[0]).all() + torch.testing.assert_close(greedy[0], actual[0].argmax(-1).to(torch.int32)) + assert actual[3].shape == (1, seq_len, 32) + torch.testing.assert_close(actual[1:], greedy[1:]) + assert qeff.lang_model.hash_params["dflash_output_logits"] + + @pytest.mark.parametrize("accepted", [0, 1, 3]) @pytest.mark.parametrize("prompt_len", [1, 5]) @pytest.mark.parametrize("chunk_size", [16, 64]) @pytest.mark.parametrize("matrix_solve", [False, True]) @pytest.mark.parametrize("moe_flavour", [None, "decode_bmm", "expert_parallel"]) -def test_qwen3_5_dflash_replay_matches_autoregressive( +def test_qwen3_5_dflash_concatenated_matches_autoregressive( accepted, prompt_len, chunk_size, matrix_solve, moe_flavour, monkeypatch ): - """Verify then replay a prefix; next-round results must exclude rejected tokens.""" + """Commit the prior prefix and verify the next block in one target call.""" from QEfficient import QEFFAutoModelForImageTextToText torch.manual_seed(7) - moe = moe_flavour is not None - hf_model = make_tiny_qwen3_5_vlm(moe=moe) - reference = copy.deepcopy(hf_model) + hf_model = make_tiny_qwen3_5_vlm(moe=moe_flavour is not None) + hf_model.config.text_config.max_position_embeddings = 512 + reference = copy.deepcopy(hf_model).model.language_model qaic_config = {"target_layer_ids": [1, 3]} - if moe: + if moe_flavour: qaic_config["moe_config"] = {"flavour": moe_flavour} qeff = QEFFAutoModelForImageTextToText(hf_model, kv_offload=True, qaic_config=qaic_config) - if moe: + if moe_flavour: qeff.transform( - ctx_len=CTX_LEN, - seq_len=16, + ctx_len=512, + seq_len=2 * chunk_size, bs=1, num_devices=1, aic_num_cores=4, qaic_config=qaic_config, prefill_only=True, - prefill_seq_len=16, + prefill_seq_len=2 * chunk_size, ) model = qeff.lang_model.model.language_model model.dflash_chunk_size = chunk_size + assert qeff.lang_model.hash_params["dflash_cache_protocol"].endswith("fused_replay_v1") for layer in model.layers: if layer.layer_type == "linear_attention": layer.linear_attn.dflash_matrix_solve = matrix_solve @@ -134,53 +181,66 @@ def test_qwen3_5_dflash_replay_matches_autoregressive( def unexpected_decode(*args, **kwargs): raise AssertionError("DFlash must not compute the one-token recurrent branch") - model_type = "qwen3_5_moe" if moe else "qwen3_5" + model_type = "qwen3_5_moe" if moe_flavour else "qwen3_5" monkeypatch.setattr( f"QEfficient.transformers.models.{model_type}.modeling_{model_type}.recurrent_gdn_decode_forward", unexpected_decode, ) - initial = qeff.model.get_dummy_inputs(kv_offload=True)["lang"]["past_key_values"] - tokens = torch.randint(10, VOCAB_SIZE, (1, 16)) - # This standalone runtime uses a single request. - cache = tuple(tuple(t[:1].clone() for t in layer) for layer in initial) - - def run(ids, start, state, commit, valid=None): - positions = torch.arange(start, start + ids.shape[1]).view(1, -1) - if valid is not None: - positions[:, valid:] = -1 - return model( - input_ids=ids, + cache = qeff.model.get_dummy_inputs(kv_offload=True, ctx_len=512)["lang"]["past_key_values"] + prompt = torch.randint(10, VOCAB_SIZE, (1, 2 * chunk_size)) + positions = torch.arange(2 * chunk_size).view(1, -1) + positions[:, prompt_len:] = -1 + with torch.no_grad(): + prefill = model( + input_ids=prompt, position_ids=positions, - past_key_values=state, - cache_commit=torch.tensor([commit]), + past_key_values=cache, + cache_commit=torch.ones(1, dtype=torch.long), use_cache=True, ) - - with torch.no_grad(): - prefill = run(tokens[:, :8], 0, cache, 1, valid=prompt_len) - verify = run(tokens[:, prompt_len : prompt_len + 4], prompt_len, prefill.past_key_values, 0) - for i in (0, 2): - for actual, expected in zip(verify.past_key_values[i], prefill.past_key_values[i]): - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - commit_count = accepted + 1 # seed token plus accepted DLM tokens - committed = run( - tokens[:, prompt_len : prompt_len + 4], prompt_len, verify.past_key_values, 1, valid=commit_count + cache = prefill.past_key_values + replay_ids = torch.zeros((1, chunk_size), dtype=torch.long) + replay_positions = torch.full_like(replay_ids, -1) + history = prompt[:, :prompt_len] + for count in [accepted + 1, 1, chunk_size, 3]: + ids = torch.randint(10, VOCAB_SIZE, (1, chunk_size)) + pos = torch.arange(history.shape[1], history.shape[1] + chunk_size).view(1, -1) + output = model( + input_ids=torch.cat([replay_ids, ids], dim=1), + position_ids=torch.cat([replay_positions, pos], dim=1), + past_key_values=cache, + cache_commit=torch.zeros(1, dtype=torch.long), + use_cache=True, + ) + expected = reference(input_ids=torch.cat([history, ids], dim=1), use_cache=False, output_hidden_states=True) + torch.testing.assert_close( + output.last_hidden_state[:, -chunk_size:], + expected.last_hidden_state[:, -chunk_size:], + rtol=2e-4, + atol=2e-4, + ) + projection = model.hidden_norm( + model.fc(torch.cat([expected.hidden_states[i][:, -chunk_size:] for i in [1, 3]], dim=-1)) + ) + torch.testing.assert_close(output.hidden_states[:, -chunk_size:], projection, rtol=2e-4, atol=2e-4) + if torch.all(replay_positions == -1): + for layer in (0, 2): + torch.testing.assert_close(output.past_key_values[layer], cache[layer], rtol=0, atol=0) + cache = output.past_key_values + replay_ids, replay_positions = ids, pos.clone() + replay_positions[:, count:] = -1 + history = torch.cat([history, ids[:, :count]], dim=1) + restarted = model( + input_ids=prompt, + position_ids=positions, + past_key_values=cache, + cache_commit=torch.ones(1, dtype=torch.long), + use_cache=True, ) - next_pos = prompt_len + commit_count - followup_ids = torch.tensor([[48, 49, 50, 51]]) - followup = run(followup_ids, next_pos, committed.past_key_values, 0) - expected_ids = torch.cat([tokens[:, :next_pos], followup_ids], dim=1) - expected = reference.model.language_model(input_ids=expected_ids, use_cache=False, output_hidden_states=True) - - torch.testing.assert_close(followup.last_hidden_state, expected.last_hidden_state[:, -4:], rtol=2e-4, atol=2e-4) - projection = model.hidden_norm(model.fc(torch.cat([expected.hidden_states[i][:, -4:] for i in [1, 3]], dim=-1))) - torch.testing.assert_close(followup.hidden_states, projection, rtol=2e-4, atol=2e-4) - assert followup.hidden_states.shape == (1, 4, 32) - with torch.no_grad(): - restarted = run(tokens[:, :8], 0, followup.past_key_values, 1, valid=prompt_len) - for i in (0, 2): - for actual, expected in zip(restarted.past_key_values[i], prefill.past_key_values[i]): - torch.testing.assert_close(actual, expected, rtol=2e-4, atol=2e-4) + for layer in (0, 2): + torch.testing.assert_close( + restarted.past_key_values[layer], prefill.past_key_values[layer], rtol=2e-4, atol=2e-4 + ) @pytest.mark.parametrize( @@ -265,21 +325,29 @@ def test_qwen3_5_dflash_rejects_unsupported_prefill(prefill_seq_len): @pytest.mark.parametrize("dummy_seq_len", [16, 32, 128, 256, 1024]) -def test_qwen3_5_dflash_export_traces_one_chunk(dummy_seq_len): +@pytest.mark.parametrize("moe", [False, True]) +def test_qwen3_5_dflash_export_traces_decode_width(dummy_seq_len, moe): from QEfficient import QEFFAutoModelForImageTextToText qeff = QEFFAutoModelForImageTextToText( - make_tiny_qwen3_5_vlm(), kv_offload=True, qaic_config={"target_layer_ids": [1, 3]} + make_tiny_qwen3_5_vlm(moe=moe), + kv_offload=True, + qaic_config={"target_layer_ids": [1, 3]}, ) + decode_width = 32 inputs = qeff.model.get_dummy_inputs(kv_offload=True, prefill_seq_len=dummy_seq_len)["lang"] - assert inputs["input_ids"].shape[1] == 16 - assert inputs["position_ids"].shape[-1] == 16 + assert inputs["input_ids"].shape[1] == decode_width + assert inputs["position_ids"].shape[-1] == decode_width + if dummy_seq_len < 32: + with pytest.raises(ValueError, match="prefill_seq_len >= 2"): + qeff.model.get_specializations(batch_size=1, prefill_seq_len=dummy_seq_len, ctx_len=2048, kv_offload=True) + return specs, _ = qeff.model.get_specializations( batch_size=1, prefill_seq_len=dummy_seq_len, ctx_len=2048, kv_offload=True ) assert [(spec["seq_len"], spec["num_logits_to_keep"]) for spec in specs["lang"]] == [ (dummy_seq_len, 1), - (16, 16), + (decode_width, 16), ] specs, _ = qeff.model.get_specializations( batch_size=1, @@ -292,12 +360,33 @@ def test_qwen3_5_dflash_export_traces_one_chunk(dummy_seq_len): assert [(spec["seq_len"], spec["num_logits_to_keep"]) for spec in specs["lang"]] == [ (dummy_seq_len, 1), (dummy_seq_len, 1), - (16, 16), - (16, 16), - (16, 16), + (decode_width, 16), + (decode_width, 16), + (decode_width, 16), ] +@pytest.mark.parametrize("decode_width", [16, 32]) +def test_qwen3_5_dflash_runtime_distinguishes_prefill_from_verify(decode_width): + from types import SimpleNamespace + + from QEfficient.generation.dflash_generation import get_qwen35_target_seq_len + + shapes = [[(8, (1, 32)), (8, (1, 1))], [(8, (1, decode_width)), (8, (16, 1))]] + session = SimpleNamespace(binding_index_map={"input_ids": 0, "num_logits_to_keep": 1}, allowed_shapes=shapes) + if decode_width == 32: + assert get_qwen35_target_seq_len(session, 16) == 32 + session.allowed_shapes = [] + session.bindings = [SimpleNamespace(dims=(1, 32)), SimpleNamespace(dims=(16, 1))] + assert get_qwen35_target_seq_len(session, 16) == 32 + else: + with pytest.raises(ValueError, match="Older two-call QPCs are unsupported"): + get_qwen35_target_seq_len(session, 16) + session.allowed_shapes = shapes[:1] + with pytest.raises(ValueError, match="Re-export and recompile"): + get_qwen35_target_seq_len(session, 16) + + @pytest.mark.parametrize("chunk_size", [0, 1, 3, 24]) def test_qwen3_5_dflash_matrix_solve_requires_power_of_two(chunk_size): from QEfficient.transformers.models.qwen3_5.modeling_qwen3_5 import QEffQwen3_5GatedDeltaNet