Skip to content

Commit 2993627

Browse files
cquil11cquil11
authored
enable KV offloading for multinode agentic sweeps (#2141)
1 parent b2c32a1 commit 2993627

3 files changed

Lines changed: 73 additions & 7 deletions

File tree

‎utils/matrix_logic/generate_sweep_configs.py‎

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -604,8 +604,8 @@ def generate_full_sweep(args, all_config_data, runner_data):
604604
prefill = bmk[Fields.PREFILL.value]
605605
decode = bmk[Fields.DECODE.value]
606606
spec_decoding = bmk.get(Fields.SPEC_DECODING.value, "none")
607-
kv_offloading = "none"
608-
kv_offload_backend = None
607+
kv_offloading = bmk.get(Fields.KV_OFFLOADING.value, "none")
608+
kv_offload_backend = bmk.get(Fields.KV_OFFLOAD_BACKEND.value)
609609
else:
610610
tp = bmk[Fields.TP.value]
611611
dcp_size = bmk.get(Fields.DCP_SIZE.value, 1)
@@ -648,6 +648,13 @@ def generate_full_sweep(args, all_config_data, runner_data):
648648
runners_for_entry = runner_nodes_to_use if runner_nodes_to_use else [runner]
649649

650650
if is_multinode:
651+
# Preserve historical exp-names for the default (no offload)
652+
# case; only append a suffix when KV offloading is active.
653+
offload_suffix = (
654+
f"_{agentic_kv_offload_suffix(kv_offloading, kv_offload_backend)}"
655+
if kv_offloading != "none"
656+
else ""
657+
)
651658
for runner_value in runners_for_entry:
652659
for conc_batch in chunk_multinode_agentic_concurrencies(conc_values):
653660
entry = {
@@ -661,16 +668,19 @@ def generate_full_sweep(args, all_config_data, runner_data):
661668
Fields.PREFILL.value: prefill,
662669
Fields.DECODE.value: decode,
663670
Fields.CONC.value: conc_batch,
664-
Fields.KV_OFFLOADING.value: "none",
671+
Fields.KV_OFFLOADING.value: kv_offloading,
665672
Fields.DURATION.value: duration,
666673
Fields.EXP_NAME.value: (
667674
f"{model_code}_p{prefill[Fields.NUM_WORKER.value]}x{prefill[Fields.TP.value]}"
668675
f"_d{decode[Fields.NUM_WORKER.value]}x{decode[Fields.TP.value]}"
669676
f"_conc{'x'.join(str(c) for c in conc_batch)}"
677+
f"{offload_suffix}"
670678
),
671679
Fields.DISAGG.value: disagg,
672680
Fields.SCENARIO_TYPE.value: "agentic-coding",
673681
}
682+
if kv_offload_backend is not None:
683+
entry[Fields.KV_OFFLOAD_BACKEND.value] = kv_offload_backend
674684
validate_agentic_matrix_entry(entry)
675685
matrix_values.append(entry)
676686
else:
@@ -885,8 +895,8 @@ def generate_test_config_sweep(args, all_config_data, runner_data=None):
885895
prefill = bmk[Fields.PREFILL.value]
886896
decode = bmk[Fields.DECODE.value]
887897
spec_decoding = bmk.get(Fields.SPEC_DECODING.value, "none")
888-
kv_offloading = "none"
889-
kv_offload_backend = None
898+
kv_offloading = bmk.get(Fields.KV_OFFLOADING.value, "none")
899+
kv_offload_backend = bmk.get(Fields.KV_OFFLOAD_BACKEND.value)
890900
else:
891901
tp = bmk[Fields.TP.value]
892902
dcp_size = bmk.get(Fields.DCP_SIZE.value, 1)
@@ -923,6 +933,13 @@ def generate_test_config_sweep(args, all_config_data, runner_data=None):
923933
continue
924934

925935
if is_multinode:
936+
# Preserve historical exp-names for the default (no offload)
937+
# case; only append a suffix when KV offloading is active.
938+
offload_suffix = (
939+
f"_{agentic_kv_offload_suffix(kv_offloading, kv_offload_backend)}"
940+
if kv_offloading != "none"
941+
else ""
942+
)
926943
for runner_value in runners_for_entry:
927944
for conc_batch in chunk_multinode_agentic_concurrencies(conc_values):
928945
entry = {
@@ -936,16 +953,19 @@ def generate_test_config_sweep(args, all_config_data, runner_data=None):
936953
Fields.PREFILL.value: prefill,
937954
Fields.DECODE.value: decode,
938955
Fields.CONC.value: conc_batch,
939-
Fields.KV_OFFLOADING.value: "none",
956+
Fields.KV_OFFLOADING.value: kv_offloading,
940957
Fields.DURATION.value: duration,
941958
Fields.EXP_NAME.value: (
942959
f"{model_code}_p{prefill[Fields.NUM_WORKER.value]}x{prefill[Fields.TP.value]}"
943960
f"_d{decode[Fields.NUM_WORKER.value]}x{decode[Fields.TP.value]}"
944961
f"_conc{'x'.join(str(c) for c in conc_batch)}"
962+
f"{offload_suffix}"
945963
),
946964
Fields.DISAGG.value: disagg,
947965
Fields.SCENARIO_TYPE.value: "agentic-coding",
948966
}
967+
if kv_offload_backend is not None:
968+
entry[Fields.KV_OFFLOAD_BACKEND.value] = kv_offload_backend
949969
matrix_values.append(validate_agentic_matrix_entry(entry))
950970
else:
951971
for conc in conc_values:

‎utils/matrix_logic/test_generate_sweep_configs.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2099,6 +2099,45 @@ def test_multinode_agentic_groups_concurrencies_per_search_entry(self):
20992099
assert result[1]["conc"] == [256]
21002100
assert result[1]["exp-name"] == "dsv4_p2x8_d1x8_conc256"
21012101

2102+
def test_multinode_agentic_preserves_kv_offload_fields(self):
2103+
config = {
2104+
"dsv4-agentic-hicache": {
2105+
"image": "sglang-rocm",
2106+
"model": "deepseek-ai/DeepSeek-V4-Pro",
2107+
"model-prefix": "dsv4",
2108+
"precision": "fp4",
2109+
"framework": "sglang-disagg",
2110+
"runner": "cluster:mi355x-amds",
2111+
"multinode": True,
2112+
"disagg": True,
2113+
"scenarios": {
2114+
"agentic-coding": [{
2115+
"search-space": [{
2116+
"conc-list": [16],
2117+
"kv-offloading": "dram",
2118+
"kv-offload-backend": "hicache",
2119+
"prefill": {"num-worker": 1, "tp": 8, "ep": 1, "dp-attn": False},
2120+
"decode": {"num-worker": 1, "tp": 8, "ep": 1, "dp-attn": False},
2121+
}],
2122+
}],
2123+
},
2124+
},
2125+
}
2126+
args = argparse.Namespace(
2127+
config_keys=["dsv4-agentic-hicache"],
2128+
seq_lens=None,
2129+
conc=None,
2130+
scenario_type=["agentic-coding"],
2131+
runner_node_filter=None,
2132+
)
2133+
2134+
result = generate_test_config_sweep(args, config)
2135+
2136+
assert len(result) == 1
2137+
assert result[0]["kv-offloading"] == "dram"
2138+
assert result[0]["kv-offload-backend"] == "hicache"
2139+
assert result[0]["exp-name"] == "dsv4_p1x8_d1x8_conc16_kvdram-hicache"
2140+
21022141

21032142
# =============================================================================
21042143
# Test apply_node_type_defaults

‎utils/matrix_logic/validation.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -240,7 +240,10 @@ class MultiNodeAgenticMatrixEntry(BaseModel):
240240
prefill: WorkerConfig
241241
decode: WorkerConfig
242242
conc: list[int]
243-
kv_offloading: Literal["none"] = Field(alias=Fields.KV_OFFLOADING.value)
243+
kv_offloading: Literal["none", "dram"] = Field(alias=Fields.KV_OFFLOADING.value)
244+
kv_offload_backend: Optional[str] = Field(
245+
default=None, alias=Fields.KV_OFFLOAD_BACKEND.value
246+
)
244247
duration: int = Field(alias=Fields.DURATION.value)
245248
exp_name: str = Field(alias=Fields.EXP_NAME.value)
246249
disagg: bool
@@ -250,6 +253,10 @@ class MultiNodeAgenticMatrixEntry(BaseModel):
250253
def validate_worker_hardware_pair(self):
251254
return _validate_worker_hardware_pair(self)
252255

256+
@model_validator(mode='after')
257+
def validate_kv_offload_fields(self):
258+
return _validate_kv_offload_fields(self)
259+
253260

254261
AgenticMatrixEntry = Union[SingleNodeAgenticMatrixEntry, MultiNodeAgenticMatrixEntry]
255262

0 commit comments

Comments
 (0)