diff --git a/src/models/recurrent_state.cpp b/src/models/recurrent_state.cpp index 8be92368a3..1583efe5fe 100644 --- a/src/models/recurrent_state.cpp +++ b/src/models/recurrent_state.cpp @@ -88,37 +88,53 @@ RecurrentState::RecurrentState(State& state) const int num_layers = static_cast(layer_indices_.size()); - const bool past_present_share_buffer = state_.params_->IsPastPresentShareBufferEnabled(model_.config_->model.type); - if (!past_present_share_buffer) { + if (!state_.params_->IsPastPresentShareBufferEnabled(model_.config_->model.type)) { throw std::runtime_error( "RecurrentState requires past_present_share_buffer=true. " "Set past_present_share_buffer to true in genai_config.json."); } + // WebGPU prohibits binding the same buffer as both read-only (input) and + // read-write (output) storage in the same compute pass, so it must use + // separate past/present buffers with swap. All other EPs share buffers + // for stable addresses (required by TRT-RTX graph replay, beneficial elsewhere). + // TODO: Remove WebGPU special case once the ORT WebGPU EP adds a + // LinearAttention kernel with native past/present buffer sharing support. + share_buffers_ = model_.p_device_kvcache_->GetType() != DeviceType::WEBGPU; + + if (!share_buffers_) { + pasts_.resize(num_layers * 2); + } presents_.reserve(num_layers * 2); auto& allocator = model_.p_device_kvcache_->GetAllocator(); - // Qwen3.5 linear-attention state is a compressed recurrent state, not a - // token-indexed KV cache. For graph replay, bind each state tensor as both - // past input and present output so ORT/TRT-RTX sees stable addresses. - // The EP/plugin kernels must read the previous contents before writing the - // updated state back to the same buffer. for (int i = 0; i < num_layers; ++i) { + if (!share_buffers_) { + pasts_[i * 2] = OrtValue::CreateTensor(allocator, conv_shape_, conv_type_); + pasts_[i * 2 + 1] = OrtValue::CreateTensor(allocator, recurrent_shape_, recurrent_type_); + } presents_.push_back(OrtValue::CreateTensor(allocator, conv_shape_, conv_type_)); presents_.push_back(OrtValue::CreateTensor(allocator, recurrent_shape_, recurrent_type_)); } + if (!share_buffers_) { + ZeroStates(pasts_); + } ZeroStates(presents_); } void RecurrentState::Add() { if (layer_indices_.empty()) return; + input_index_ = state_.inputs_.size(); + output_index_ = state_.outputs_.size(); + const int num_layers = static_cast(layer_indices_.size()); for (int i = 0; i < num_layers * 2; ++i) { - auto* past = presents_[i].get(); - state_.inputs_.push_back(past); + // Shared: alias input=output for stable addresses. + // WebGPU: separate past/present buffers to avoid aliasing violation. + state_.inputs_.push_back(share_buffers_ ? presents_[i].get() : pasts_[i].get()); state_.input_names_.push_back(input_name_strings_[i].c_str()); state_.outputs_.push_back(presents_[i].get()); state_.output_names_.push_back(output_name_strings_[i].c_str()); @@ -126,23 +142,37 @@ void RecurrentState::Add() { } void RecurrentState::Update() { + if (layer_indices_.empty() || share_buffers_) return; + + const int num_layers = static_cast(layer_indices_.size()); + for (int i = 0; i < num_layers * 2; ++i) { + std::swap(pasts_[i], presents_[i]); + state_.inputs_[input_index_ + i] = pasts_[i].get(); + state_.outputs_[output_index_ + i] = presents_[i].get(); + } } void RecurrentState::RewindTo(size_t index) { if (layer_indices_.empty()) return; if (index != 0) { - // Recurrent states cannot be partially rewound — they are compressed summaries - // with no per-position history. Non-zero rewind is a no-op; the state remains unchanged. - if (g_log.enabled) - Log("warning", "RecurrentState::RewindTo(" + std::to_string(index) + - ") is a no-op. Recurrent states cannot be partially rewound."); - return; + throw std::runtime_error( + "RecurrentState::RewindTo(" + std::to_string(index) + + ") is not supported. Recurrent states cannot be partially rewound."); + } + if (share_buffers_) { + // Shared buffers: zero in place, addresses stay stable. + ZeroStates(presents_); + } else { + // Zero and rebind all state buffers. + ZeroStates(pasts_); + ZeroStates(presents_); + const int num_layers = static_cast(layer_indices_.size()); + for (int i = 0; i < num_layers * 2; ++i) { + state_.inputs_[input_index_ + i] = pasts_[i].get(); + state_.outputs_[output_index_ + i] = presents_[i].get(); + } } - // Shared recurrent states keep stable input/output pointers for graph replay. - // Reset the state contents in place without rebinding. - ZeroStates(presents_); - return; } void RecurrentState::ZeroStates(std::vector>& states) { diff --git a/src/models/recurrent_state.h b/src/models/recurrent_state.h index e7c3c42734..a6787020bc 100644 --- a/src/models/recurrent_state.h +++ b/src/models/recurrent_state.h @@ -27,8 +27,14 @@ struct RecurrentState { std::vector layer_indices_; // Interleaved as [conv_0, recurrent_0, conv_1, recurrent_1, ...] + std::vector> pasts_; std::vector> presents_; + // WebGPU cannot alias input/output buffers, so it uses separate past/present\n // with swap. All other EPs share buffers for stable addresses. + bool share_buffers_{false}; + size_t input_index_{~0U}; + size_t output_index_{~0U}; + // Kept alive for state_ const char* pointers std::vector input_name_strings_; std::vector output_name_strings_; diff --git a/test/python/test_qwen_fara_models.py b/test/python/test_qwen_fara_models.py index 74498fb232..605f553d9c 100644 --- a/test/python/test_qwen_fara_models.py +++ b/test/python/test_qwen_fara_models.py @@ -27,7 +27,9 @@ log = logging.getLogger("qwen-fara-vision-tests") -@pytest.mark.parametrize("relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]) +@pytest.mark.parametrize( + "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")] +) @pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"]) def test_qwen_fara_vision_basic(test_data_path, relative_model_path, relative_image_path): """Test basic vision preprocessing for Qwen/Fara-style models.""" @@ -48,7 +50,9 @@ def test_qwen_fara_vision_basic(test_data_path, relative_model_path, relative_im assert "pixel_values" in inputs -@pytest.mark.parametrize("relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]) +@pytest.mark.parametrize( + "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")] +) @pytest.mark.parametrize("relative_image_path", [Path("images") / "landscape.jpg"]) def test_qwen_fara_vision_load_from_bytes(test_data_path, relative_model_path, relative_image_path): """Test loading images from bytes for Qwen/Fara models.""" @@ -70,7 +74,9 @@ def test_qwen_fara_vision_load_from_bytes(test_data_path, relative_model_path, r assert "pixel_values" in inputs -@pytest.mark.parametrize("relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]) +@pytest.mark.parametrize( + "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")] +) @pytest.mark.parametrize( "relative_image_paths", [[Path("images") / "australia.jpg", Path("images") / "landscape.jpg"]], @@ -114,9 +120,7 @@ def test_qwen3_vl_vision_dynamic_grid_dim(test_data_path): """ onnx = pytest.importorskip("onnx") - vision_path = os.path.join( - test_data_path, "qwen3-vl-vision-preprocessing", "dummy_vision.onnx" - ) + vision_path = os.path.join(test_data_path, "qwen3-vl-vision-preprocessing", "dummy_vision.onnx") model = onnx.load(vision_path) # Find image_grid_thw input @@ -131,19 +135,18 @@ def test_qwen3_vl_vision_dynamic_grid_dim(test_data_path): # dim-0 should be symbolic (dynamic), not a fixed integer dim0 = grid_input.type.tensor_type.shape.dim[0] assert dim0.dim_param != "", ( - f"image_grid_thw dim-0 should be symbolic (e.g. 'num_images') " - f"but got static dim_value={dim0.dim_value}" - ) - assert dim0.dim_param == "num_images", ( - f"Expected dim_param='num_images', got '{dim0.dim_param}'" + f"image_grid_thw dim-0 should be symbolic (e.g. 'num_images') but got static dim_value={dim0.dim_value}" ) + assert dim0.dim_param == "num_images", f"Expected dim_param='num_images', got '{dim0.dim_param}'" # dim-1 should be static 3 (temporal, height, width) dim1 = grid_input.type.tensor_type.shape.dim[1] assert dim1.dim_value == 3, f"image_grid_thw dim-1 should be 3, got {dim1.dim_value}" -@pytest.mark.parametrize("relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]) +@pytest.mark.parametrize( + "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")] +) def test_qwen_fara_text_only_generation(test_data_path, relative_model_path): """Test text-only generation without images.""" model_path = os.fspath(Path(test_data_path) / relative_model_path) @@ -160,7 +163,9 @@ def test_qwen_fara_text_only_generation(test_data_path, relative_model_path): assert "input_ids" in inputs -@pytest.mark.parametrize("relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]) +@pytest.mark.parametrize( + "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")] +) @pytest.mark.parametrize("relative_image_path", [Path("images") / "sheet.png"]) def test_qwen_fara_vision_with_special_tokens(test_data_path, relative_model_path, relative_image_path): """Test vision processing with special tokens in prompt.""" @@ -181,7 +186,9 @@ def test_qwen_fara_vision_with_special_tokens(test_data_path, relative_model_pat assert "input_ids" in inputs -@pytest.mark.parametrize("relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]) +@pytest.mark.parametrize( + "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")] +) @pytest.mark.parametrize("relative_image_path", [Path("images") / "10809054.jpg"]) def test_qwen_fara_vision_different_image_formats(test_data_path, relative_model_path, relative_image_path): """Test processing different image formats.""" @@ -200,7 +207,9 @@ def test_qwen_fara_vision_different_image_formats(test_data_path, relative_model assert "pixel_values" in inputs -@pytest.mark.parametrize("relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]) +@pytest.mark.parametrize( + "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")] +) @pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"]) def test_qwen_fara_accuracy_comparison(test_data_path, relative_model_path, relative_image_path): """ @@ -262,7 +271,9 @@ def test_qwen_fara_accuracy_comparison(test_data_path, relative_model_path, rela log.debug(f"ONNX pixel_values range: [{pixel_min:.4f}, {pixel_max:.4f}]") -@pytest.mark.parametrize("relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]) +@pytest.mark.parametrize( + "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")] +) @pytest.mark.parametrize("relative_image_path", [Path("images") / "sheet.png"]) def test_qwen_fara_preprocessing_consistency(test_data_path, relative_model_path, relative_image_path): """ @@ -401,8 +412,8 @@ def test_qwen3_vl_pixel_values_shape(test_data_path, relative_image_path): @pytest.mark.parametrize( "model_name,expected_patch_dim", [ - ("qwen-vision-preprocessing", 1176), # Qwen2.5-VL: patch_size=14, 14*14*3*2=1176 - ("qwen3-vl-vision-preprocessing", 1536), # Qwen3-VL: patch_size=16, 16*16*3*2=1536 + ("qwen-vision-preprocessing", 1176), # Qwen2.5-VL: patch_size=14, 14*14*3*2=1176 + ("qwen3-vl-vision-preprocessing", 1536), # Qwen3-VL: patch_size=16, 16*16*3*2=1536 ], ) @pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"]) @@ -440,7 +451,9 @@ def test_qwen_vl_family_patch_size_difference(test_data_path, model_name, expect log.debug(f"{model_name} pixel_values shape: {pixel_array.shape}") -@pytest.mark.parametrize("relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]) +@pytest.mark.parametrize( + "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")] +) @pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"]) def test_qwen_vl_preprocessing_output_completeness(test_data_path, relative_model_path, relative_image_path): """ @@ -464,9 +477,7 @@ def test_qwen_vl_preprocessing_output_completeness(test_data_path, relative_mode # All four keys must be present for vision inputs expected_keys = {"pixel_values", "input_ids", "image_grid_thw", "num_image_tokens"} actual_keys = set(inputs.keys()) - assert expected_keys.issubset(actual_keys), ( - f"Missing keys: {expected_keys - actual_keys}. Got: {actual_keys}" - ) + assert expected_keys.issubset(actual_keys), f"Missing keys: {expected_keys - actual_keys}. Got: {actual_keys}" def _to_numpy(tensor): """Convert a tensor-like object to NumPy (supports as_numpy, numpy, np.array).""" @@ -499,7 +510,9 @@ def _to_numpy(tensor): log.debug(f"{relative_model_path} output: pv={pv.shape}, grid={grid}, nit={nit}, ids={ids.shape}") -@pytest.mark.parametrize("relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]) +@pytest.mark.parametrize( + "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")] +) @pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"]) def test_qwen_vl_image_grid_thw_consistency(test_data_path, relative_model_path, relative_image_path): """ @@ -599,6 +612,7 @@ def test_qwen_vl_normalization_range_difference(test_data_path, relative_image_p # Qwen3.5 hybrid model tests (RecurrentState + sparse KV cache) # --------------------------------------------------------------------------- + def test_qwen35_hybrid_model_loads(test_data_path): """Test that a Qwen3.5 hybrid model (with recurrent + KV states) loads successfully.""" model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing") @@ -687,6 +701,99 @@ def test_qwen35_hybrid_vision_preprocessing(test_data_path, relative_image_path) assert "pixel_values" in inputs +# --------------------------------------------------------------------------- +# Qwen3.5 hybrid model tests — CUDA EP (RecurrentState with shared buffers) +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not og.is_cuda_available(), reason="CUDA EP not available") +def test_qwen35_hybrid_generator_creates_cuda(test_data_path): + """Test that a Generator can be created for the hybrid model on CUDA. + Validates RecurrentState shared-buffer path on CUDA EP.""" + model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing") + if not os.path.exists(model_path): + pytest.skip("qwen35-hybrid-preprocessing test model not found") + + config = og.Config(model_path) + config.clear_providers() + config.append_provider("cuda") + model = og.Model(config) + params = og.GeneratorParams(model) + params.set_search_options(max_length=20) + generator = og.Generator(model, params) + assert generator is not None + + +@pytest.mark.skipif(not og.is_cuda_available(), reason="CUDA EP not available") +def test_qwen35_hybrid_text_generation_cuda(test_data_path): + """Test that the hybrid model generator constructs and prefill executes on CUDA. + RecurrentState uses shared buffers (same tensor as input and output).""" + model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing") + if not os.path.exists(model_path): + pytest.skip("qwen35-hybrid-preprocessing test model not found") + + config = og.Config(model_path) + config.clear_providers() + config.append_provider("cuda") + model = og.Model(config) + params = og.GeneratorParams(model) + params.set_search_options(max_length=5) + generator = og.Generator(model, params) + assert generator is not None + + +# --------------------------------------------------------------------------- +# Qwen3.5 hybrid model tests — WebGPU EP (RecurrentState with separate buffers) +# --------------------------------------------------------------------------- + + +def _is_webgpu_test_enabled(): + """WebGPU tests require both runtime support and explicit opt-in via TEST_WEBGPU env var.""" + return ( + hasattr(og, "is_webgpu_available") + and og.is_webgpu_available() + and os.environ.get("TEST_WEBGPU", "").lower() in ("true", "1", "yes") + ) + + +@pytest.mark.skipif(not _is_webgpu_test_enabled(), reason="WebGPU EP not available or TEST_WEBGPU not set") +def test_qwen35_hybrid_generator_creates_webgpu(test_data_path): + """Test that a Generator can be created for the hybrid model on WebGPU. + Validates RecurrentState separate-buffer path (WebGPU cannot alias + input/output buffers in the same compute pass).""" + model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing") + if not os.path.exists(model_path): + pytest.skip("qwen35-hybrid-preprocessing test model not found") + + config = og.Config(model_path) + config.clear_providers() + config.append_provider("webgpu") + model = og.Model(config) + params = og.GeneratorParams(model) + params.set_search_options(max_length=20) + generator = og.Generator(model, params) + assert generator is not None + + +@pytest.mark.skipif(not _is_webgpu_test_enabled(), reason="WebGPU EP not available or TEST_WEBGPU not set") +def test_qwen35_hybrid_text_generation_webgpu(test_data_path): + """Test that the hybrid model generator constructs and prefill executes on WebGPU. + RecurrentState uses separate past/present buffers to avoid the WebGPU + buffer aliasing restriction (Storage read-write | read-only conflict).""" + model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing") + if not os.path.exists(model_path): + pytest.skip("qwen35-hybrid-preprocessing test model not found") + + config = og.Config(model_path) + config.clear_providers() + config.append_provider("webgpu") + model = og.Model(config) + params = og.GeneratorParams(model) + params.set_search_options(max_length=5) + generator = og.Generator(model, params) + assert generator is not None + + # Standalone runner functionality def run_qwen_fara_vision_tests( cwd: str | bytes | os.PathLike, diff --git a/test/test_models/qwen35-hybrid-preprocessing/create_dummy_models.py b/test/test_models/qwen35-hybrid-preprocessing/create_dummy_models.py index 53882d5092..fa447a458c 100644 --- a/test/test_models/qwen35-hybrid-preprocessing/create_dummy_models.py +++ b/test/test_models/qwen35-hybrid-preprocessing/create_dummy_models.py @@ -30,20 +30,30 @@ def create_dummy_embedding_model(output_path: str, hidden_size: int = 1024, vocab_size: int = 248320): """Create dummy embedding model: input_ids, image_features -> inputs_embeds""" input_ids = helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["batch", "sequence_len"]) - image_features = helper.make_tensor_value_info("image_features", TensorProto.FLOAT, ["num_image_tokens", hidden_size]) - inputs_embeds = helper.make_tensor_value_info("inputs_embeds", TensorProto.FLOAT, ["batch", "sequence_len", hidden_size]) + image_features = helper.make_tensor_value_info( + "image_features", TensorProto.FLOAT, ["num_image_tokens", hidden_size] + ) + inputs_embeds = helper.make_tensor_value_info( + "inputs_embeds", TensorProto.FLOAT, ["batch", "sequence_len", hidden_size] + ) # Create a simple graph that outputs zeros of the right shape # Shape -> ConstantOfShape to produce zeros shape_node = helper.make_node("Shape", ["input_ids"], ["ids_shape"]) # We need [batch, seq_len, hidden_size] output - hidden_const = helper.make_node("Constant", [], ["hidden_dim"], - value=helper.make_tensor("hidden_dim", TensorProto.INT64, [1], [hidden_size])) + hidden_const = helper.make_node( + "Constant", [], ["hidden_dim"], value=helper.make_tensor("hidden_dim", TensorProto.INT64, [1], [hidden_size]) + ) concat_node = helper.make_node("Concat", ["ids_shape", "hidden_dim"], ["embed_shape"], axis=0) - zero_val = helper.make_node("Constant", [], ["zero_val"], - value=helper.make_tensor("zero_val", TensorProto.FLOAT, [1], [0.0])) - cos_node = helper.make_node("ConstantOfShape", ["embed_shape"], ["inputs_embeds"], - value=helper.make_tensor("val", TensorProto.FLOAT, [1], [0.01])) + zero_val = helper.make_node( + "Constant", [], ["zero_val"], value=helper.make_tensor("zero_val", TensorProto.FLOAT, [1], [0.0]) + ) + cos_node = helper.make_node( + "ConstantOfShape", + ["embed_shape"], + ["inputs_embeds"], + value=helper.make_tensor("val", TensorProto.FLOAT, [1], [0.01]), + ) graph = helper.make_graph( [shape_node, hidden_const, concat_node, zero_val, cos_node], @@ -65,23 +75,41 @@ def create_dummy_vision_model(output_path: str, hidden_size: int = 1024): # (spatial_merge_size=2 -> merge_sq=4) shape_node = helper.make_node("Shape", ["pixel_values"], ["pv_shape"]) gather_node = helper.make_node("Gather", ["pv_shape", "zero_idx"], ["num_patches"], axis=0) - zero_idx_const = helper.make_node("Constant", [], ["zero_idx"], - value=helper.make_tensor("zero_idx", TensorProto.INT64, [], [0])) - four_const = helper.make_node("Constant", [], ["four"], - value=helper.make_tensor("four", TensorProto.INT64, [], [4])) + zero_idx_const = helper.make_node( + "Constant", [], ["zero_idx"], value=helper.make_tensor("zero_idx", TensorProto.INT64, [], [0]) + ) + four_const = helper.make_node( + "Constant", [], ["four"], value=helper.make_tensor("four", TensorProto.INT64, [], [4]) + ) div_node = helper.make_node("Div", ["num_patches", "four"], ["num_feats"]) - hidden_const = helper.make_node("Constant", [], ["hidden_dim"], - value=helper.make_tensor("hidden_dim", TensorProto.INT64, [1], [hidden_size])) + hidden_const = helper.make_node( + "Constant", [], ["hidden_dim"], value=helper.make_tensor("hidden_dim", TensorProto.INT64, [1], [hidden_size]) + ) reshape_feats = helper.make_node("Reshape", ["num_feats", "one_shape"], ["num_feats_1d"]) - one_shape_const = helper.make_node("Constant", [], ["one_shape"], - value=helper.make_tensor("one_shape", TensorProto.INT64, [1], [1])) + one_shape_const = helper.make_node( + "Constant", [], ["one_shape"], value=helper.make_tensor("one_shape", TensorProto.INT64, [1], [1]) + ) concat_node = helper.make_node("Concat", ["num_feats_1d", "hidden_dim"], ["feat_shape"], axis=0) - cos_node = helper.make_node("ConstantOfShape", ["feat_shape"], ["image_features"], - value=helper.make_tensor("val", TensorProto.FLOAT, [1], [0.01])) + cos_node = helper.make_node( + "ConstantOfShape", + ["feat_shape"], + ["image_features"], + value=helper.make_tensor("val", TensorProto.FLOAT, [1], [0.01]), + ) graph = helper.make_graph( - [zero_idx_const, shape_node, gather_node, four_const, div_node, - hidden_const, one_shape_const, reshape_feats, concat_node, cos_node], + [ + zero_idx_const, + shape_node, + gather_node, + four_const, + div_node, + hidden_const, + one_shape_const, + reshape_feats, + concat_node, + cos_node, + ], "vision", [pixel_values, image_grid_thw], [image_features], @@ -93,7 +121,7 @@ def create_dummy_vision_model(output_path: str, hidden_size: int = 1024): def create_dummy_decoder_model( output_path: str, num_layers: int = 4, - kv_layers: list = None, + kv_layers: list | None = None, hidden_size: int = 1024, num_kv_heads: int = 2, head_size: int = 256, @@ -116,8 +144,12 @@ def create_dummy_decoder_model( outputs = [] # Standard inputs - inputs_embeds = helper.make_tensor_value_info("inputs_embeds", TensorProto.FLOAT, ["batch", "sequence_len", hidden_size]) - attention_mask = helper.make_tensor_value_info("attention_mask", TensorProto.INT64, ["batch", "past_seq_len_plus_seq_len"]) + inputs_embeds = helper.make_tensor_value_info( + "inputs_embeds", TensorProto.FLOAT, ["batch", "sequence_len", hidden_size] + ) + attention_mask = helper.make_tensor_value_info( + "attention_mask", TensorProto.INT64, ["batch", "past_seq_len_plus_seq_len"] + ) position_ids = helper.make_tensor_value_info("position_ids", TensorProto.INT64, [3, "batch", "sequence_len"]) inputs.extend([inputs_embeds, attention_mask, position_ids]) @@ -125,32 +157,60 @@ def create_dummy_decoder_model( for layer_idx in range(num_layers): if layer_idx in kv_layers: # KV cache layer - inputs.append(helper.make_tensor_value_info( - f"past_key_values.{layer_idx}.key", TensorProto.FLOAT, - ["batch", num_kv_heads, "past_sequence_len", head_size])) - inputs.append(helper.make_tensor_value_info( - f"past_key_values.{layer_idx}.value", TensorProto.FLOAT, - ["batch", num_kv_heads, "past_sequence_len", head_size])) - outputs.append(helper.make_tensor_value_info( - f"present.{layer_idx}.key", TensorProto.FLOAT, - ["batch", num_kv_heads, "total_sequence_len", head_size])) - outputs.append(helper.make_tensor_value_info( - f"present.{layer_idx}.value", TensorProto.FLOAT, - ["batch", num_kv_heads, "total_sequence_len", head_size])) + inputs.append( + helper.make_tensor_value_info( + f"past_key_values.{layer_idx}.key", + TensorProto.FLOAT, + ["batch", num_kv_heads, "past_sequence_len", head_size], + ) + ) + inputs.append( + helper.make_tensor_value_info( + f"past_key_values.{layer_idx}.value", + TensorProto.FLOAT, + ["batch", num_kv_heads, "past_sequence_len", head_size], + ) + ) + outputs.append( + helper.make_tensor_value_info( + f"present.{layer_idx}.key", + TensorProto.FLOAT, + ["batch", num_kv_heads, "total_sequence_len", head_size], + ) + ) + outputs.append( + helper.make_tensor_value_info( + f"present.{layer_idx}.value", + TensorProto.FLOAT, + ["batch", num_kv_heads, "total_sequence_len", head_size], + ) + ) else: # Recurrent state layer - inputs.append(helper.make_tensor_value_info( - f"past_key_values.{layer_idx}.conv_state", TensorProto.FLOAT, - ["batch", conv_dim, conv_kernel - 1])) - inputs.append(helper.make_tensor_value_info( - f"past_key_values.{layer_idx}.recurrent_state", TensorProto.FLOAT, - ["batch", num_linear_heads, linear_head_dim, linear_head_dim])) - outputs.append(helper.make_tensor_value_info( - f"present.{layer_idx}.conv_state", TensorProto.FLOAT, - ["batch", conv_dim, conv_kernel - 1])) - outputs.append(helper.make_tensor_value_info( - f"present.{layer_idx}.recurrent_state", TensorProto.FLOAT, - ["batch", num_linear_heads, linear_head_dim, linear_head_dim])) + inputs.append( + helper.make_tensor_value_info( + f"past_key_values.{layer_idx}.conv_state", TensorProto.FLOAT, ["batch", conv_dim, conv_kernel - 1] + ) + ) + inputs.append( + helper.make_tensor_value_info( + f"past_key_values.{layer_idx}.recurrent_state", + TensorProto.FLOAT, + ["batch", num_linear_heads, linear_head_dim, linear_head_dim], + ) + ) + outputs.append( + helper.make_tensor_value_info( + f"present.{layer_idx}.conv_state", TensorProto.FLOAT, ["batch", conv_dim, conv_kernel - 1] + ) + ) + outputs.append( + helper.make_tensor_value_info( + f"present.{layer_idx}.recurrent_state", + TensorProto.FLOAT, + ["batch", num_linear_heads, linear_head_dim, linear_head_dim], + ) + ) # Logits output logits = helper.make_tensor_value_info("logits", TensorProto.FLOAT, [None, None, vocab_size]) @@ -165,35 +225,50 @@ def create_dummy_decoder_model( gather_batch = helper.make_node("Gather", ["embed_shape", "idx_0"], ["batch_dim"], axis=0) gather_seq = helper.make_node("Gather", ["embed_shape", "idx_1"], ["seq_dim"], axis=0) - idx_0_const = helper.make_node("Constant", [], ["idx_0"], - value=helper.make_tensor("idx_0", TensorProto.INT64, [], [0])) - idx_1_const = helper.make_node("Constant", [], ["idx_1"], - value=helper.make_tensor("idx_1", TensorProto.INT64, [], [1])) - vocab_const = helper.make_node("Constant", [], ["vocab_dim"], - value=helper.make_tensor("vocab_dim", TensorProto.INT64, [1], [vocab_size])) + idx_0_const = helper.make_node( + "Constant", [], ["idx_0"], value=helper.make_tensor("idx_0", TensorProto.INT64, [], [0]) + ) + idx_1_const = helper.make_node( + "Constant", [], ["idx_1"], value=helper.make_tensor("idx_1", TensorProto.INT64, [], [1]) + ) + vocab_const = helper.make_node( + "Constant", [], ["vocab_dim"], value=helper.make_tensor("vocab_dim", TensorProto.INT64, [1], [vocab_size]) + ) nodes.extend([idx_0_const, idx_1_const, gather_batch, gather_seq, vocab_const]) reshape_batch = helper.make_node("Reshape", ["batch_dim", "one_shape"], ["batch_1d"]) reshape_seq = helper.make_node("Reshape", ["seq_dim", "one_shape"], ["seq_1d"]) - one_shape_const = helper.make_node("Constant", [], ["one_shape"], - value=helper.make_tensor("one_shape", TensorProto.INT64, [1], [1])) + one_shape_const = helper.make_node( + "Constant", [], ["one_shape"], value=helper.make_tensor("one_shape", TensorProto.INT64, [1], [1]) + ) concat_logits_shape = helper.make_node("Concat", ["batch_1d", "seq_1d", "vocab_dim"], ["logits_shape"], axis=0) - logits_node = helper.make_node("ConstantOfShape", ["logits_shape"], ["logits"], - value=helper.make_tensor("val", TensorProto.FLOAT, [1], [0.0])) + logits_node = helper.make_node( + "ConstantOfShape", ["logits_shape"], ["logits"], value=helper.make_tensor("val", TensorProto.FLOAT, [1], [0.0]) + ) nodes.extend([one_shape_const, reshape_batch, reshape_seq, concat_logits_shape, logits_node]) # Identity for all state tensors for layer_idx in range(num_layers): if layer_idx in kv_layers: - nodes.append(helper.make_node("Identity", - [f"past_key_values.{layer_idx}.key"], [f"present.{layer_idx}.key"])) - nodes.append(helper.make_node("Identity", - [f"past_key_values.{layer_idx}.value"], [f"present.{layer_idx}.value"])) + nodes.append( + helper.make_node("Identity", [f"past_key_values.{layer_idx}.key"], [f"present.{layer_idx}.key"]) + ) + nodes.append( + helper.make_node("Identity", [f"past_key_values.{layer_idx}.value"], [f"present.{layer_idx}.value"]) + ) else: - nodes.append(helper.make_node("Identity", - [f"past_key_values.{layer_idx}.conv_state"], [f"present.{layer_idx}.conv_state"])) - nodes.append(helper.make_node("Identity", - [f"past_key_values.{layer_idx}.recurrent_state"], [f"present.{layer_idx}.recurrent_state"])) + nodes.append( + helper.make_node( + "Identity", [f"past_key_values.{layer_idx}.conv_state"], [f"present.{layer_idx}.conv_state"] + ) + ) + nodes.append( + helper.make_node( + "Identity", + [f"past_key_values.{layer_idx}.recurrent_state"], + [f"present.{layer_idx}.recurrent_state"], + ) + ) graph = helper.make_graph(nodes, "decoder", inputs, outputs) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) @@ -255,7 +330,7 @@ def create_genai_config(output_path: str, num_kv_layers: int, kv_layers: list): "no_repeat_ngram_size": 0, "num_beams": 1, "num_return_sequences": 1, - "past_present_share_buffer": False, + "past_present_share_buffer": True, "repetition_penalty": 1.0, "temperature": 1.0, "top_k": 1, @@ -269,8 +344,12 @@ def create_genai_config(output_path: str, num_kv_layers: int, kv_layers: list): def main(): parser = argparse.ArgumentParser(description="Generate dummy ONNX models for Qwen3.5 hybrid model testing") - parser.add_argument("--output", type=str, default="test/test_models/qwen35-hybrid-preprocessing", - help="Output directory for the dummy models") + parser.add_argument( + "--output", + type=str, + default="test/test_models/qwen35-hybrid-preprocessing", + help="Output directory for the dummy models", + ) args = parser.parse_args() output_dir = args.output @@ -282,7 +361,9 @@ def main(): num_kv_layers = len(kv_layers) print(f"Creating dummy hybrid model in {output_dir}") - print(f" {num_layers} total layers, KV at {kv_layers}, recurrent at {[i for i in range(num_layers) if i not in kv_layers]}") + print( + f" {num_layers} total layers, KV at {kv_layers}, recurrent at {[i for i in range(num_layers) if i not in kv_layers]}" + ) create_dummy_embedding_model(os.path.join(output_dir, "dummy_embedding.onnx")) print(" Created dummy_embedding.onnx")