A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 06:45 UTC
[Perf][Bugfix] Update dflash aux layer indexing (#40727)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
May 19, 2026 at 23:15 UTC
c628a93a64fb4929c3c11d8e2c7244c4826b4f76
3 files changed
+11
-4
tests/v1/e2e/spec_decode/test_spec_decode.py
+5
-2
@@ -1315,7 +1315,7 @@ def test_dflash_acceptance_rates(dflash_config):
1315
expected_acceptance_lengths = {
1316
"mt-bench": 4.24,
1317
"humaneval": 6.50,
1318
- "gsm8k": 6.54 * 0.95, # runs with a subset of prompts so extra wide tol here
1318
+ "gsm8k": 6.54 * 0.975, # runs with a subset of prompts so extra wide tol here
1319
}
1320
1321
tokenizer = spec_llm.get_tokenizer()
@@ -1347,7 +1347,10 @@ def test_dflash_acceptance_rates(dflash_config):
1347
acceptance_lengths.append(acceptance_len)
1348
1349
mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths)
1350
- expected_len = expected_len * 0.9
1350
+ # Fairly tight tolerance of 95% against the paper's figures,
1351
+ # watching for regressions. Can be relaxed if test is flaky but be sure to
1352
+ # check for genuine issues such as #40727.
1353
+ expected_len = expected_len * 0.95
1354
print(
1355
f"DFlash acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}"
1356
f" (expected at least {expected_len:.2f})"
vllm/transformers_utils/configs/speculators/algos.py
+2
-1
@@ -99,7 +99,8 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
99
aux_layer_ids = config_dict["aux_hidden_state_layer_ids"]
100
pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids
101
102
+ # DFlash configs use different indexing for the target layers, see #40727
103
pre_trained_config["dflash_config"] = {
104
"mask_token_id": config_dict["mask_token_id"],
104
- "target_layer_ids": aux_layer_ids,
105
+ "target_layer_ids": [i - 1 for i in aux_layer_ids],
106
}
vllm/v1/worker/gpu_model_runner.py
+4
-1
@@ -5213,7 +5213,10 @@ class GPUModelRunner(
5213
eagle_config = getattr(hf_config, "eagle_config", None)
5214
5215
if dflash_config and isinstance(dflash_config, dict):
5216
- layer_ids = dflash_config.get("target_layer_ids")
5216
+ # Add 1 to convert DFlash's aux layer id semantics
5217
+ layer_ids = [
5218
+ i + 1 for i in (dflash_config.get("target_layer_ids") or [])
5219
+ ]
5220
5221
if eagle_config and isinstance(eagle_config, dict):
5222
layer_ids = eagle_config.get("eagle_aux_hidden_state_layer_ids")