[Model] fix(dflash): dtype mismatch in combine_hidden_states

[ciphernaut]

When running AWQ-quantized models with DFlash speculative decoding,
attention layers can produce float32 activations while fc.weight is
float16, causing a dtype mismatch error in combine_hidden_states.

Add a dtype guard before the fc call, matching the existing pattern
used for cos_sin_cache in precompute_and_store_context_kv.

Reproduces with: QuantTrio/Qwen3.6-35B-A3B-AWQ +
z-lab/Qwen3.6-35B-A3B-DFlash, dtype=float16, cpu_offload_gb > 0.

Co-authored-by: Claude

Purpose

Fixes a dtype mismatch error in DFlashQwen3ForCausalLM.combine_hidden_states when using AWQ-quantized models with DFlash speculative decoding.

When cpu_offload_gb > 0 is set (required to fit large AWQ models on consumer GPUs), attention layers produce float32 activations. The draft head's fc layer weight is float16, causing a runtime error on the bare self.model.fc(hidden_states) call.

The fix adds a dtype guard matching the pattern already used in precompute_and_store_context_kv for cos_sin_cache.

Tested with: QuantTrio/Qwen3.6-35B-A3B-AWQ + z-lab/Qwen3.6-35B-A3B-DFlash, dtype=float16, cpu_offload_gb=12, single RTX 4090 (24 GB).

Note: AI assistance (Claude) was used to identify the fix and draft this PR. All changed lines have been reviewed and tested by the submitter.
This fix does not duplicate any open PR (checked: #39995 targets FlashInfer backends; no open PRs touch combine_hidden_states).

Test Plan

# Start vLLM with AWQ model + DFlash + CPU offload
vllm serve QuantTrio/Qwen3.6-35B-A3B-AWQ \
  --speculative-config '{"method":"dflash","model":"z-lab/Qwen3.6-35B-A3B-DFlash","num_speculative_tokens":8}' \
  --dtype float16 \
  --gpu-memory-utilization 0.90 \
  --max-model-len 65536 \
  --cpu-offload-gb 12

# Confirm generation works without dtype error
curl http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"QuantTrio/Qwen3.6-35B-A3B-AWQ","prompt":"Hello","max_tokens":32}'

Test Result

Before: RuntimeError: expected scalar type Float but found Half raised in combine_hidden_states on first generation request.

After: Generation completes without error. Quality validated against GSM8K (8-shot CoT, 1319 samples, 95.75% exact match) and HumanEval (0-shot, 164 samples, 95.1% pass@1) — no regression vs non-offload baseline.

---

• The purpose of the PR: fixes dtype mismatch in combine_hidden_states when using AWQ + DFlash with CPU offload (no existing issue to link).
• The test plan: manual server startup + curl generation request above.
• The test results: before/after comparison above; GSM8K and HumanEval quality confirmed unchanged.
• The necessary documentation update: not applicable (no new model or user-facing behaviour change).

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[gemini-code-assist]

Code Review

This pull request adds a type cast to the hidden states in the combine_hidden_states method of the Qwen3 DFlash model to ensure compatibility with mixed-precision targets. Feedback suggests using params_dtype instead of weight.dtype to correctly handle quantized models where weights may be stored in a packed integer format.

[gemini-code-assist]

Using self.model.fc.weight.dtype to determine the target dtype is potentially unsafe if the draft model is quantized (e.g., using AWQ or GPTQ). In such cases, the weight attribute often contains packed integers (like torch.int32), and casting activations to that dtype would cause a crash or incorrect results.

It is more robust to use self.model.fc.params_dtype, which vLLM uses to store the intended compute/parameter dtype for the layer, regardless of the underlying weight storage format.

Suggested change

	# Cast to fc weight dtype to handle mixed-precision targets (e.g. AWQ
	# with unquantized attention layers that output float32 activations).
	fc_weight = self.model.fc.weight
	if hidden_states.dtype != fc_weight.dtype:
	hidden_states = hidden_states.to(fc_weight.dtype)
	# Cast to fc compute dtype to handle mixed-precision targets (e.g. AWQ
	# with unquantized attention layers that output float32 activations).
	target_dtype = self.model.fc.params_dtype
	if hidden_states.dtype != target_dtype:
	hidden_states = hidden_states.to(target_dtype)

[ciphernaut]

full validation was done before this suggestion. basic validation done afterwards

[mergify]

Hi @ciphernaut, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @ciphernaut, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[ciphernaut]

I don't know how to fix the remaining items, I will just close this MR now.