Skip to content

fix(agentic): propagate speculative decoding to sweep jobs #2183

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cquil11 merged 1 commit into from
Jul 13, 2026
Merged

fix(agentic): propagate speculative decoding to sweep jobs #2183

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/e2e-tests.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,7 @@ jobs:
isl: '0'
osl: '0'
max-model-len: '0'
spec-decoding: 'none'
spec-decoding: ${{ matrix.config.spec-decoding }}
disagg: 'false'
run-eval: false
scenario-type: agentic-coding
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/run-sweep.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -521,7 +521,7 @@ jobs:
isl: '0'
osl: '0'
max-model-len: '0'
spec-decoding: 'none'
spec-decoding: ${{ matrix.config.spec-decoding }}
disagg: ${{ 'false' }}
run-eval: false
scenario-type: agentic-coding
Expand Down
6 changes: 6 additions & 0 deletions utils/matrix_logic/generate_sweep_configs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -645,6 +645,7 @@
pcp_size = bmk.get(Fields.PCP_SIZE.value, 1)
ep = bmk.get(Fields.EP.value)
dp_attn = bmk.get(Fields.DP_ATTN.value)
spec_decoding = bmk.get(Fields.SPEC_DECODING.value, "none")

@claude claude Bot Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📊 Line Count Report

  • Total Lines: 1313
  • Base Lines: 1307
  • Change: 📈 +6 lines

kv_offloading = bmk[Fields.KV_OFFLOADING.value]
kv_offload_backend = bmk.get(Fields.KV_OFFLOAD_BACKEND.value)
total_cpu_dram_gb = (
Expand Down Expand Up @@ -733,16 +734,18 @@
Fields.PCP_SIZE.value: pcp_size,
Fields.EP.value: ep if ep is not None else 1,
Fields.DP_ATTN.value: dp_attn if dp_attn is not None else False,
Fields.SPEC_DECODING.value: spec_decoding,
Fields.CONC.value: conc,
Fields.KV_OFFLOADING.value: kv_offloading,
Fields.TOTAL_CPU_DRAM_GB.value: total_cpu_dram_gb,
Fields.DURATION.value: duration,
Fields.EXP_NAME.value: (
f"{model_code}_tp{tp}_conc{conc}_"
f"{agentic_kv_offload_suffix(kv_offloading, kv_offload_backend)}"
+ (f"_spec-{spec_decoding}" if spec_decoding != "none" else "")
),
Fields.SCENARIO_TYPE.value: "agentic-coding",
}

Check warning on line 748 in utils/matrix_logic/generate_sweep_configs.py

View check run for this annotation

Claude / Claude Code Review

Multi-node agentic exp-name omits spec-decoding suffix 

The multi-node agentic exp-name (around lines 700-712 and 1000-1012) wasn't updated to append `_spec-{spec_decoding}` the way the single-node agentic exp-name in this PR was, even though the multi-node agentic entry does carry `spec-decoding`. Two multi-node agentic configs differing only in spec-decoding will get identical exp-names — consider adding the same suffix logic there for consistency.
Comment on lines 742 to 748

@claude claude Bot Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 The multi-node agentic exp-name (around lines 700-712 and 1000-1012) wasn't updated to append _spec-{spec_decoding} the way the single-node agentic exp-name in this PR was, even though the multi-node agentic entry does carry spec-decoding. Two multi-node agentic configs differing only in spec-decoding will get identical exp-names — consider adding the same suffix logic there for consistency.

Extended reasoning...

This PR's stated goal is to "include non-default speculative-decoding modes in experiment names," and it does this correctly for the single-node agentic path in both generate_full_sweep and generate_test_config_sweep:

Fields.EXP_NAME.value: (
    f"{model_code}_tp{tp}_conc{conc}_"
    f"{agentic_kv_offload_suffix(kv_offloading, kv_offload_backend)}"
    + (f"_spec-{spec_decoding}" if spec_decoding != "none" else "")
),

However, the analogous multi-node agentic exp-name construction (generate_sweep_configs.py lines ~700-712 and ~1000-1012 in the two functions) was not given the same treatment:

Fields.EXP_NAME.value: (
    f"{model_code}_p{prefill[Fields.NUM_WORKER.value]}x{prefill[Fields.TP.value]}"
    f"_d{decode[Fields.NUM_WORKER.value]}x{decode[Fields.TP.value]}"
    f"_conc{'x'.join(str(c) for c in conc_batch)}"
    f"{offload_suffix}"
),

This block only ever appends offload_suffix (for KV offloading) — it never references spec_decoding, despite the fact that the multi-node agentic entry a few lines above already resolves spec_decoding = bmk.get(Fields.SPEC_DECODING.value, "none") and stores it via Fields.SPEC_DECODING.value: spec_decoding in the same entry dict.

Concrete proof of the gap: take two multi-node agentic search-space entries with identical prefill/decode topology and the same conc batch, differing only in spec-decoding: mtp vs spec-decoding: none. Both will produce entry["spec-decoding"] values of "mtp" and "none" respectively (correct, and correctly forwarded to run-sweep.yml/e2e-tests.yml via matrix.config.spec-decoding), but both will produce the exact same exp-name string, e.g. kimik2.6_p1x8_d1x8_conc16_kvnone. This is inconsistent with the single-node agentic path, where the same scenario would produce ..._kvnone_spec-mtp vs ..._kvnone.

Why existing tests don't catch it: the new tests added in this PR (test_agentic_spec_decoding_is_propagated, test_single_node_agentic_spec_decoding_is_propagated) only exercise the single-node agentic path's sample_single_node_agentic_config fixture; there is no equivalent multi-node agentic fixture/test asserting exp-name behavior for spec-decoding.

Impact: this is cosmetic, not functional. The correct spec-decoding value is still stored on the entry and propagated through to the workflow inputs/launcher regardless of exp-name — so no benchmark runs with the wrong speculative-decoding setting. It only makes it harder to visually distinguish multi-node agentic result rows that differ solely by spec-decoding mode when reading exp-name in dashboards/logs. It's also worth noting that exp-name is already not a strict uniqueness key elsewhere (e.g. fixed-seq-len paths use the coarser {model_code}_{seq_len_str} across many differing configs), so this doesn't introduce a new class of problem, just extends an existing gap to a scenario this PR specifically claims to have addressed.

Suggested fix: mirror the single-node suffix logic in both multi-node agentic exp-name blocks, e.g. append + (f"_spec-{spec_decoding}" if spec_decoding != "none" else "") after {offload_suffix} in both generate_full_sweep and generate_test_config_sweep.

if kv_offload_backend is not None:
entry[Fields.KV_OFFLOAD_BACKEND.value] = kv_offload_backend
entry.update(component_metadata(bmk, val))
Expand Down Expand Up @@ -948,6 +951,7 @@
pcp_size = bmk.get(Fields.PCP_SIZE.value, 1)
ep = bmk.get(Fields.EP.value)
dp_attn = bmk.get(Fields.DP_ATTN.value)
spec_decoding = bmk.get(Fields.SPEC_DECODING.value, "none")
kv_offloading = bmk[Fields.KV_OFFLOADING.value]
kv_offload_backend = bmk.get(Fields.KV_OFFLOAD_BACKEND.value)
total_cpu_dram_gb = (
Expand Down Expand Up @@ -1029,13 +1033,15 @@
Fields.PCP_SIZE.value: pcp_size,
Fields.EP.value: ep if ep is not None else 1,
Fields.DP_ATTN.value: dp_attn if dp_attn is not None else False,
Fields.SPEC_DECODING.value: spec_decoding,
Fields.CONC.value: conc,
Fields.KV_OFFLOADING.value: kv_offloading,
Fields.TOTAL_CPU_DRAM_GB.value: total_cpu_dram_gb,
Fields.DURATION.value: duration,
Fields.EXP_NAME.value: (
f"{model_code}_tp{tp}_conc{conc}_"
f"{agentic_kv_offload_suffix(kv_offloading, kv_offload_backend)}"
+ (f"_spec-{spec_decoding}" if spec_decoding != "none" else "")
),
Fields.SCENARIO_TYPE.value: "agentic-coding",
}
Expand Down
72 changes: 72 additions & 0 deletions utils/matrix_logic/test_generate_sweep_configs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,39 @@ def sample_runner_config():
}


@pytest.fixture
def sample_single_node_agentic_config():
"""Single-node agentic config with explicit and default spec decoding."""
return {
"kimik2.6-fp4-b300-trt-agentic": {
"image": "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc5",
"model": "moonshotai/Kimi-K2.5",
"model-prefix": "kimik2.6",
"precision": "fp4",
"framework": "trt",
"runner": "cluster:b300-nv",
"multinode": False,
"scenarios": {
"agentic-coding": [{
"search-space": [
{
"tp": 8,
"spec-decoding": "mtp",
"kv-offloading": "none",
"conc-list": [16],
},
{
"tp": 8,
"kv-offloading": "none",
"conc-list": [32],
},
],
}],
},
},
}


@pytest.fixture
def full_sweep_args_single_node():
"""Args for full-sweep single-node command."""
Expand Down Expand Up @@ -689,6 +722,22 @@ def test_matrix_entry_structure(self, sample_single_node_config, sample_runner_c
for row in explicit_result
} == {(2, 2, 2)}

def test_agentic_spec_decoding_is_propagated(
self,
sample_single_node_agentic_config,
sample_runner_config,
full_sweep_args_single_node,
):
result = generate_full_sweep(
full_sweep_args_single_node,
sample_single_node_agentic_config,
sample_runner_config,
)

assert [entry["spec-decoding"] for entry in result] == ["mtp", "none"]
assert result[0]["exp-name"].endswith("_kvnone_spec-mtp")
assert result[1]["exp-name"].endswith("_kvnone")

def test_filter_by_model_prefix(self, sample_single_node_config, sample_runner_config, full_sweep_args_single_node):
"""Filter by model prefix should work."""
full_sweep_args_single_node.model_prefix = ["dsr1"]
Expand Down Expand Up @@ -1914,6 +1963,29 @@ def test_single_node_parallelism_fields_are_generated(
for row in explicit_result
] == [(2, 2, 2)]

def test_single_node_agentic_spec_decoding_is_propagated(
self,
sample_single_node_agentic_config,
sample_runner_config,
):
args = argparse.Namespace(
config_keys=["kimik2.6-fp4-b300-trt-agentic"],
seq_lens=None,
conc=None,
scenario_type=["agentic-coding"],
runner_node_filter=None,
)

result = generate_test_config_sweep(
args,
sample_single_node_agentic_config,
sample_runner_config,
)

assert [entry["spec-decoding"] for entry in result] == ["mtp", "none"]
assert result[0]["exp-name"].endswith("_kvnone_spec-mtp")
assert result[1]["exp-name"].endswith("_kvnone")

def test_multinode_parallelism_fields_are_generated(
self,
sample_multinode_config,
Expand Down
3 changes: 3 additions & 0 deletions utils/matrix_logic/validation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -276,6 +276,9 @@ class SingleNodeAgenticMatrixEntry(BaseModel):
pcp_size: int = Field(alias=Fields.PCP_SIZE.value, gt=0, strict=True)
ep: int
dp_attn: bool = Field(alias=Fields.DP_ATTN.value)
spec_decoding: Literal["mtp", "draft_model", "none"] = Field(
default="none", alias=Fields.SPEC_DECODING.value
)
conc: int
kv_offloading: Literal["none", "dram"] = Field(
alias=Fields.KV_OFFLOADING.value
Expand Down