Skip to content

Fix webgpu shared buffer issue #2191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
apsonawane merged 4 commits into from
Jun 2, 2026
Merged

Fix webgpu shared buffer issue #2191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5af1b69
Fix webgpu shared buffer issue
apsonawane Jun 1, 2026
ba2a39e
Add devicetype check for recurrent_state
apsonawane Jun 1, 2026
b1e827e
fix pipeline
apsonawane Jun 1, 2026
6cf765b
Add todo comment
apsonawane Jun 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 49 additions & 19 deletions src/models/recurrent_state.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,61 +88,91 @@ RecurrentState::RecurrentState(State& state)

const int num_layers = static_cast<int>(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;
Comment thread
apsonawane marked this conversation as resolved.

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<int>(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());
}
}

void RecurrentState::Update() {
if (layer_indices_.empty() || share_buffers_) return;

const int num_layers = static_cast<int>(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<int>(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<std::unique_ptr<OrtValue>>& states) {
Expand Down
6 changes: 6 additions & 0 deletions src/models/recurrent_state.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,14 @@ struct RecurrentState {
std::vector<int> layer_indices_;

// Interleaved as [conv_0, recurrent_0, conv_1, recurrent_1, ...]
std::vector<std::unique_ptr<OrtValue>> pasts_;
std::vector<std::unique_ptr<OrtValue>> 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<std::string> input_name_strings_;
std::vector<std::string> output_name_strings_;
Expand Down
153 changes: 130 additions & 23 deletions test/python/test_qwen_fara_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand All @@ -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."""
Expand All @@ -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"]],
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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."""
Expand All @@ -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."""
Expand All @@ -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):
"""
Expand Down Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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"])
Expand Down Expand Up @@ -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):
"""
Expand All @@ -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)."""
Expand Down Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading