[WebGPU] `Kernel "[GroupQueryAttention] /model/layers.0/attn/GroupQueryAttention" failed. Error: Input "key" is expected to have 3, 4, or 5 dimensions".`

[xenova]

Describe the issue

The following error occurs when trying to run https://huggingface.co/HuggingFaceTB/SmolVLM-Instruct on WebGPU.

Note that the CPU implementation operates correctly, so this is indeed a bug with the WebGPU EP. Moreover, the zero-dimension tensor is by design, and is used for the first generation step.

To reproduce

1. Install and build Transformers.js from source (https://github.com/huggingface/transformers.js)
2. Run the following code in-browser:

import {
  AutoProcessor,
  AutoModelForVision2Seq,
  load_image,
} from "@huggingface/transformers";

// Initialize processor and model
const model_id = "HuggingFaceTB/SmolVLM-Instruct";
const processor = await AutoProcessor.from_pretrained(model_id);
const model = await AutoModelForVision2Seq.from_pretrained(model_id, {
  dtype: {
    embed_tokens: "fp16", // "fp32", "fp16", "q8"
    vision_encoder: "q4", // "fp32", "fp16", "q8", "q4", "q4f16"
    decoder_model_merged: "q4", // "q8", "q4", "q4f16"
  },
  device: 'webgpu',
});

// Load images
const image1 = await load_image("https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg");
const image2 = await load_image("https://huggingface.co/spaces/merve/chameleon-7b/resolve/main/bee.jpg");

// Create input messages
const messages = [
  {
    role: "user",
    content: [
      { type: "image" },
      { type: "image" },
      { type: "text", text: "Can you describe the two images?" },
    ],
  },
];

// Prepare inputs
const text = processor.apply_chat_template(messages, { add_generation_prompt: true });
const inputs = await processor(text, [image1, image2], {
  // Set `do_image_splitting: true` to split images into multiple patches.
  // NOTE: This uses more memory, but can provide more accurate results.
  do_image_splitting: false,
});

// Generate outputs
const generated_ids = await model.generate({
  ...inputs,
  max_new_tokens: 500,
});
const generated_texts = processor.batch_decode(
  generated_ids.slice(null, [inputs.input_ids.dims.at(-1), null]),
  { skip_special_tokens: true },
);
console.log(generated_texts[0]);
// ' In the first image, there is a green statue of liberty on a pedestal in the middle of the water. The water is surrounded by trees and buildings in the background. In the second image, there are pink and red flowers with a bee on the pink flower.'

Urgency

This blocks SmolVLM usage in Transformers.js.

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.20.1

Execution Provider

'webgpu' (WebGPU)

[xenova]

This also happens for https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct and https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct, which are easier models to test with.

[guschmue]

looking

[xenova]

Thanks! I'd recommend testing with https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct (or other models with GQA and num_key_value_heads != num_attention_heads).

[github-actions]

This issue has been automatically marked as stale due to inactivity and will be closed in 30 days if no further activity occurs. If further support is needed, please provide an update and/or more details.

[xenova]

Bump since this is pretty important

[guschmue]

for me https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct does not crash and instead produces garbage.
llama3.2-1b with gqa is working fine.
I created my own gqa model, output is:

user: Tell me about Constantinople.
assistant: You are a helpful AI assistant. You are trained by a user named Hugging.

Just to be sure I tried MHA:

system
You are a helpful AI assistant named SmolLM, trained by Hugging Face
user
Tell me about Constantinople.
assistant
The city of Constantinople, the greatest metropolis in the world. It was the capital of the Byzantine Empire, the largest empire in history, and the largest city in the world.
As a historian, I can tell you that the city was the center of the Byzantine Empire, the largest empire in history, and the largest city in the world. The city was the seat of the emperor, the capital, the largest city, the largest city, the largest city, the largest city, the largest city, the largest city, the largest city, the largest city, the largest city, the largest city, the largest...

With MHA if I do something like summarize the model is doing a little better initially but then doesn't know how to stop and repeats the same sentence over and over.

looking

[xenova]

Thanks Guenther Schmuelling (@guschmue)! Let me know if there's anything I can do to assist in debugging 🫡

[guschmue]

using transformers.js-examples/smollm-webgpu with SmolLM2-135M-Instruct:
my MHA model is fine
my GQA model does not crash but has wrong output or hangs depending on the query

[xenova]

That's the same behaviour as I was seeing too 👍 Model output is fine in Node.js, so the garbage output is not due to severe quantisation.