diff --git a/.github/workflows/e2e-tests.yml b/.github/workflows/e2e-tests.yml index 803ea98ba9..51c68e1d85 100644 --- a/.github/workflows/e2e-tests.yml +++ b/.github/workflows/e2e-tests.yml @@ -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 diff --git a/.github/workflows/run-sweep.yml b/.github/workflows/run-sweep.yml index 66e429284c..b983982af8 100644 --- a/.github/workflows/run-sweep.yml +++ b/.github/workflows/run-sweep.yml @@ -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 diff --git a/utils/matrix_logic/generate_sweep_configs.py b/utils/matrix_logic/generate_sweep_configs.py index c5141c209b..babc328a81 100644 --- a/utils/matrix_logic/generate_sweep_configs.py +++ b/utils/matrix_logic/generate_sweep_configs.py @@ -645,6 +645,7 @@ def generate_full_sweep(args, all_config_data, runner_data): 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 = ( @@ -733,6 +734,7 @@ def generate_full_sweep(args, all_config_data, runner_data): 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, @@ -740,6 +742,7 @@ def generate_full_sweep(args, all_config_data, runner_data): 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", } @@ -948,6 +951,7 @@ def generate_test_config_sweep(args, all_config_data, runner_data=None): 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 = ( @@ -1029,6 +1033,7 @@ def generate_test_config_sweep(args, all_config_data, runner_data=None): 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, @@ -1036,6 +1041,7 @@ def generate_test_config_sweep(args, all_config_data, runner_data=None): 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", } diff --git a/utils/matrix_logic/test_generate_sweep_configs.py b/utils/matrix_logic/test_generate_sweep_configs.py index 8072fc3570..7f2583f57e 100644 --- a/utils/matrix_logic/test_generate_sweep_configs.py +++ b/utils/matrix_logic/test_generate_sweep_configs.py @@ -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.""" @@ -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"] @@ -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, diff --git a/utils/matrix_logic/validation.py b/utils/matrix_logic/validation.py index 805859e14e..6b7c671289 100644 --- a/utils/matrix_logic/validation.py +++ b/utils/matrix_logic/validation.py @@ -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