[Perf][Bugfix] Update dflash aux layer indexing

[benchislett]

Purpose

A discrepancy in indexing causes a slight gap in the acceptance rates for DFlash v.s. the reference.

See: https://github.com/z-lab/dflash/blob/main/dflash/model.py#L44

This will have implications for Speculators, not sure if they also have this issue.

Test Plan

Existing dflash AR test passes and acceptance rate goes up.

Test Result

Baseline vLLM:

DFlash acceptance_len for mt-bench: [3.92, 4.03] (expected at least 3.82)
DFlash acceptance_len for humaneval: [6.29, 6.30] (expected at least 5.85)
DFlash acceptance_len for gsm8k (subset of 200 prompts): 6.18 (expected at least 5.59)

Offset by +2:

DFlash acceptance_len for mt-bench: 3.95 (expected at least 3.82)
DFlash acceptance_len for humaneval: 6.06 (expected at least 5.85)
DFlash acceptance_len for gsm8k (subset of 200 prompts): 6.08 (expected at least 5.59)

Offset by +1:
DFlash acceptance_len for mt-bench: [4.11, 4.16] (expected at least 3.82)
DFlash acceptance_len for humaneval: [6.44, 6.48] (expected at least 5.85)
DFlash acceptance_len for gsm8k (subset of 200 prompts): [6.37, 6.42] (expected at least 5.59)

Reference data from the paper (https://arxiv.org/pdf/2602.06036):
"mt-bench": 4.24,
"humaneval": 6.50,
"gsm8k": 6.54

[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.

[gemini-code-assist]

Code Review

This pull request implements a +1 shift for DFlash auxiliary layer IDs in the GPU model runner to align with expected semantics. Feedback suggests applying this shift directly in the configuration utility to ensure consistency across execution paths and hardening the list comprehension against potential null values in the configuration dictionary.

[gemini-code-assist]

The TODO should be addressed by applying the +1 shift here. Since gpu_model_runner.py prioritizes the eagle_aux_hidden_state_layer_ids field (line 4937), the fix in the model runner's fallback logic is bypassed for DFlash models configured through this function. Applying the shift here ensures that the correct layer indices are used in the primary execution path.

Suggested change

	# TODO: does this need to be shifted by 1 like in gpu_model_runner?
	aux_layer_ids = config_dict["aux_hidden_state_layer_ids"]
	pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids
	# Add 1 to convert DFlash's aux layer id semantics
	aux_layer_ids = [i + 1 for i in config_dict["aux_hidden_state_layer_ids"]]
	pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids

[gemini-code-assist]

Using dflash_config.get("target_layer_ids", []) can lead to a TypeError if the key exists in the dictionary but its value is explicitly set to None. It is safer to use dflash_config.get("target_layer_ids") or [] to ensure the list comprehension always receives an iterable.

Suggested change

	# Add 1 to convert DFlash's aux layer id semantics
	layer_ids = [i + 1 for i in dflash_config.get("target_layer_ids", [])]
	# Add 1 to convert DFlash's aux layer id semantics
	layer_ids = [i + 1 for i in (dflash_config.get("target_layer_ids") or [])]