Fix Llama4 shape mismatch for 32k+ context window (#842)

[afierka-intel]

Llama4 for max_model_len > 32k enable temperature adjustment https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/llama4.py#L719. Enabled adjustment causes tensor q shape modification from 2D to 3D: https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/llama4.py#L307. This tensor is passing to UnqnatizedFusedMoEMetod -> forward: https://github.com/vllm-project/vllm-gaudi/blob/main/vllm_gaudi/ops/hpu_fused_moe.py#L163 causing invalid reshaping - we trying to return a 3D output.view based on 2D output tensor.

Found that following PR introduced the bug: #680 and #684

Cherry-picked from releases/v0.13.0

---

[Copilot]

Pull request overview

This PR fixes a shape mismatch issue in the Llama4 model when using context windows larger than 32k tokens. The problem occurs because temperature adjustment enabled for long contexts modifies the input tensor shape from 2D to 3D, but the output reshaping logic didn't account for this variation based on data parallelism configuration.

Changes:

• Modified the output reshaping logic in UnquantizedFusedMoEMethod.forward_oot() to handle both 2D and 3D input tensors correctly based on the dp_size value

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The conditional logic based on dp_size > 1 appears to be a workaround for handling different input shapes rather than addressing the root cause. Consider explicitly checking the input tensor dimensionality (len(input_shape)) to make the intent clearer and more maintainable. This would better document why different reshaping strategies are needed and make the code less fragile if dp_size semantics change.

Suggested change

	if layer.dp_size > 1:
	return output.view(*(output.size(0), *input_shape[1:]))
	else:
	if len(input_shape) == 2:
	# Handle 2D inputs where the leading dimension may have been
	# modified (e.g. by data parallel dispatch); keep the trailing
	# dimension(s) from the original shape and infer the leading one
	# from the actual output tensor.
	return output.view(output.size(0), *input_shape[1:])
	else:
	# For higher-rank inputs, restore the original shape directly.

[github-actions]

✅ CI Passed

All checks passed successfully against the following vllm commit:
6218034dd7f9a56596e4fd8c8c8fc1d8011ed9c2