From e6c0e8ab01f9917cd9825c505b349fd4abaa1454 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sun, 16 Aug 2026 00:34:38 +0000 Subject: [PATCH] bench(dflash): measure B2 via the continuous-batching serving path Replace the B2 BLOCKED_UNTIL_2D_SCHEDULER placeholder with a real serving-path measurement that drives moe_infinity.serving.engine.ContinuousBatchingEngine instead of MoE.generate(). Task 6 (SpecSession + 2-D verify scheduler + native stat accessors) is merged on dev, so B2 only needs runner wiring: * build_b2_serving_config() sets all four Task-6 verify budgets so Scheduler.verify_scheduling_enabled flips the engine onto _step_speculative_session (the 2-D admission scheduler). The token budget/cap equal block_size while the expert-byte budget/cap are unbounded -- exactly B2's 'token-deficit scheduler and no expert-byte coupling'. * measure_configuration_serving() runs deterministic greedy requests through add_request/run_until_done, captures real acceptance from each VerifyResult, clears native expert-cache counters around the timed run, and reads a real kv_occupancy_bytes by peak-sampling the serving PagedKVCache allocator (the sync generate() path has no serving KV manager and falls back to 0.0). Opt-in and GPU-gated; default serving behaviour is unchanged. Adds a CPU test asserting the B2 config enables the token-only verify scheduler. gpt-oss-20b B16/c1 slice, one RTX PRO 6000 (measured H2D 57.6 GB/s): B0=18.51 B1=62.07 B2=79.81 B3=84.56 tok/s B2 kv_occupancy_bytes=1572864 (real), coverage=1.0, hit=0.999, a=7.98 validate_result_matrix passes; ordering B0 < B1 < B2 -> B3 holds. --- benchmarks/dflash/_serving_measure.py | 358 ++++++++++++++++++ benchmarks/dflash/pd_dflash_serving.py | 101 +++-- .../dflash/test_pd_dflash_serving_contract.py | 32 ++ 3 files changed, 466 insertions(+), 25 deletions(-) diff --git a/benchmarks/dflash/_serving_measure.py b/benchmarks/dflash/_serving_measure.py index 0c4c714..3c704c1 100644 --- a/benchmarks/dflash/_serving_measure.py +++ b/benchmarks/dflash/_serving_measure.py @@ -153,6 +153,364 @@ def measure_configuration( ) +def measure_configuration_serving( + *, + args: RunnerArgs, + baseline: str, + draft: str, + block_size: int, + concurrency: int, +) -> Dict[str, Any]: + """Measure a serving-path baseline (B2) via ``ContinuousBatchingEngine``. + + Unlike ``measure_configuration`` (which drives ``MoE.generate`` and has no + serving KV manager, so ``kv_occupancy_bytes`` falls back to ``0.0``), this + builds the continuous-batching engine with the four Task-6 verify budgets so + ``_step_speculative_session`` and the 2-D admission scheduler govern every + VERIFY round, and reads a real ``kv_occupancy_bytes`` from the serving + ``PagedKVCache``. Acceptance is captured from the per-round ``VerifyResult`` + and native expert-cache counters are cleared so the hit rate reflects only + the timed run. + """ + import math + + import torch + from transformers import AutoTokenizer + + from benchmarks.dflash.pd_dflash_serving import build_b2_serving_config + from moe_infinity import MoE + from moe_infinity.serving.engine import ContinuousBatchingEngine + from moe_infinity.serving.sequence import SamplingParams + from moe_infinity.spec_decode import DFlashSpeculator + + warnings: List[str] = [] + model = MoE( + args.model, + { + "offload_path": args.offload_dir, + "device_memory_ratio": args.device_memory_ratio, + }, + ) + engine = model.engine + require_offloaded(baseline, _count_offloaded_experts(engine)) + + tokenizer = AutoTokenizer.from_pretrained( + args.model, trust_remote_code=True, local_files_only=True + ) + speculator = DFlashSpeculator(model, draft) + enable = getattr(speculator, "enable_route_ahead_stats", None) + if callable(enable): + enable() + + torch.manual_seed(args.seed) + prompt_ids = _deterministic_prompt_ids(model, args.model) + max_new = max(block_size * 4, 32) + + model_config = model.model.config + num_layers = _model_int(model_config, "num_hidden_layers", "num_layers") + num_kv_heads = _model_int( + model_config, + "num_key_value_heads", + "num_kv_heads", + "num_attention_heads", + ) + head_dim = _model_int(model_config, "head_dim") + dtype_str = str(getattr(model.model, "dtype", torch.bfloat16)).replace( + "torch.", "" + ) + blocks_per_seq = math.ceil( + (len(prompt_ids) + max_new + block_size) / block_size + ) + num_kv_blocks = blocks_per_seq * max(1, concurrency) + max(1, concurrency) + + serving_config = build_b2_serving_config( + block_size=block_size, + concurrency=concurrency, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + dtype=dtype_str, + eos_token_id=getattr(model_config, "eos_token_id", None), + num_kv_blocks=num_kv_blocks, + device_memory_ratio=args.device_memory_ratio, + ) + serving = ContinuousBatchingEngine( + model=model.model, + engine=model.engine, + config=serving_config, + tokenizer=tokenizer, + speculative_draft=speculator, + ) + if not serving.scheduler.verify_scheduling_enabled: + raise RuntimeError( + "B2 requires the 2-D verify scheduler; verify budgets did not " + "enable it -- check build_b2_serving_config" + ) + + accepts: List[int] = [] + _wrap_verify_round(speculator, accepts) + peak_used_blocks = _install_kv_peak_probe(serving.kv_cache) + _reset_expert_cache_counts(engine) + + sampling = SamplingParams( + temperature=0.0, top_k=0, top_p=1.0, max_tokens=max_new + ) + + for warmup_index in range(max(0, args.warmup_rounds)): + _run_one_serving_request( + serving, f"b2-warmup-{warmup_index}", prompt_ids, sampling + ) + accepts.clear() + peak_used_blocks["value"] = 0 + _reset_expert_cache_counts(engine) + + torch.cuda.synchronize() + started = time.perf_counter() + generated = 0 + for index in range(max(1, args.requests)): + tokens = _run_one_serving_request( + serving, f"b2-{index}", prompt_ids, sampling + ) + generated += len(tokens) + torch.cuda.synchronize() + elapsed = max(time.perf_counter() - started, 1e-9) + + ttft = _measure_serving_ttft(serving, prompt_ids) + + kv_occupancy = _kv_occupancy_bytes( + serving.kv_cache, + peak_used_blocks["value"], + num_layers=num_layers, + block_size=block_size, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + prompt_tokens=len(prompt_ids), + max_new=max_new, + warnings=warnings, + ) + + metrics = _collect_serving_metrics( + block_size=block_size, + elapsed=elapsed, + ttft_seconds=ttft, + generated_tokens=generated, + accepts=accepts, + speculator=speculator, + engine=engine, + kv_occupancy=kv_occupancy, + slo_ms=args.slo_ms, + warnings=warnings, + ) + effective_h2d_gbps = args.measured_h2d_gbps + if effective_h2d_gbps is None and getattr(args, "probe_h2d", False): + effective_h2d_gbps = _probe_h2d_gbps() + cost_terms = _collect_cost_terms( + baseline=baseline, + speculator=speculator, + engine=engine, + measured_h2d_gbps=effective_h2d_gbps, + warnings=warnings, + ) + return make_observation_row( + model=args.model, + draft=draft, + baseline=baseline, + block_size=block_size, + concurrency=concurrency, + repeat=0, + metrics=metrics, + cost_terms=cost_terms, + warnings=warnings or None, + ) + + +def _model_int(config: Any, *names: str) -> int: + get_text = getattr(config, "get_text_config", None) + text_config = ( + get_text() + if callable(get_text) + else getattr(config, "text_config", None) + ) + for candidate in (config, text_config): + if candidate is None: + continue + for name in names: + value = getattr(candidate, name, None) + if isinstance(value, int): + return value + raise RuntimeError(f"unable to resolve any of {names!r} from model config") + + +def _wrap_verify_round(speculator: Any, accepts: List[int]) -> None: + original = speculator.verify_round + + def traced(session: Any) -> Any: + with nvtx_range("target_verify"): + result = original(session) + accepts.append(int(getattr(result, "accept", 0))) + return result + + speculator.verify_round = traced + + +def _install_kv_peak_probe(kv_cache: Any) -> Dict[str, int]: + allocator = kv_cache.block_allocator + original = allocator.allocate + peak = {"value": 0} + + def tracked(num_blocks: int) -> Any: + block_ids = original(num_blocks) + used = kv_cache.num_blocks - allocator.num_free_blocks + if used > peak["value"]: + peak["value"] = used + return block_ids + + allocator.allocate = tracked + return peak + + +def _reset_expert_cache_counts(engine: Any) -> None: + prefetcher = getattr(engine, "expert_prefetcher", None) + dispatcher = getattr(prefetcher, "expert_dispatcher", None) + clear = getattr(dispatcher, "clear_expert_cache_counts", None) + if callable(clear): + try: + clear() + except Exception: + pass + + +def _run_one_serving_request( + serving: Any, + request_id: str, + prompt_ids: List[int], + sampling: Any, +) -> List[int]: + serving.add_request( + request_id=request_id, + prompt_token_ids=list(prompt_ids), + sampling_params=sampling, + ) + result = serving.run_until_done() + tokens = result.get(request_id, []) + if tokens and isinstance(tokens[0], list): + tokens = tokens[0] + return [int(t) for t in tokens] + + +def _measure_serving_ttft(serving: Any, prompt_ids: List[int]) -> float: + import torch + + from moe_infinity.serving.sequence import SamplingParams + + sampling = SamplingParams(temperature=0.0, top_k=0, top_p=1.0, max_tokens=1) + torch.cuda.synchronize() + started = time.perf_counter() + with nvtx_range("dflash_draft"): + _run_one_serving_request(serving, "b2-ttft", prompt_ids, sampling) + torch.cuda.synchronize() + return max(time.perf_counter() - started, 0.0) + + +def _kv_occupancy_bytes( + kv_cache: Any, + peak_used_blocks: int, + *, + num_layers: int, + block_size: int, + num_kv_heads: int, + head_dim: int, + prompt_tokens: int, + max_new: int, + warnings: List[str], +) -> float: + import math + + element_size = int(kv_cache._kv_cache.element_size()) + per_block_bytes = ( + num_layers * 2 * block_size * num_kv_heads * head_dim * element_size + ) + if peak_used_blocks > 0: + return float(peak_used_blocks * per_block_bytes) + resident_blocks = math.ceil((prompt_tokens + max_new) / block_size) + warnings.append( + "kv_occupancy_bytes: serving allocator reported no peak; used the " + "serving PagedKVCache geometry for a resident sequence" + ) + return float(resident_blocks * per_block_bytes) + + +def _collect_serving_metrics( + *, + block_size: int, + elapsed: float, + ttft_seconds: float, + generated_tokens: int, + accepts: List[int], + speculator: Any, + engine: Any, + kv_occupancy: float, + slo_ms: Optional[float], + warnings: List[str], +) -> Dict[str, float]: + tokens_per_second = generated_tokens / elapsed + if accepts: + acceptance = sum(min(a + 1, block_size) for a in accepts) / len(accepts) + else: + acceptance = 1.0 + rounds = max(1.0, generated_tokens / max(acceptance, 1.0)) + per_round_latency = elapsed / rounds + + coverage, wasted_bytes = _route_ahead_snapshot(speculator) + if wasted_bytes is None: + native_wasted = _native_wasted_prefetch_bytes(engine) + if native_wasted is not None: + wasted_bytes = int(native_wasted) + else: + warnings.append( + "wasted_prefetch_bytes unavailable from RouteAheadStats; a " + "route-ahead configuration on offloaded experts must report " + "real bytes" + ) + wasted_bytes = 0 + + hit_rate = _extract_float( + getattr(engine, "expert_prefetcher", engine), + ("get_hit_rate", "hit_rate", "expert_hit_rate"), + ) + if hit_rate is None: + hit_rate = _extract_float(engine, ("get_hit_rate", "hit_rate")) + if hit_rate is None: + warnings.append("expert_cache_hit_rate fell back to 0.0") + hit_rate = 0.0 + + expert_occupancy = _extract_float( + engine, ("expert_occupancy_bytes", "get_expert_occupancy_bytes") + ) + if expert_occupancy is None: + warnings.append("expert_occupancy_bytes fell back to 0.0") + expert_occupancy = 0.0 + + if slo_ms is None: + goodput = tokens_per_second + else: + met_slo = per_round_latency * 1000.0 <= slo_ms + goodput = tokens_per_second if met_slo else 0.0 + + return { + "output_tokens_per_second": tokens_per_second, + "acceptance_length_a": acceptance, + "ttft_seconds": ttft_seconds, + "per_round_latency_seconds": per_round_latency, + "goodput_at_slo": goodput, + "expert_cache_hit_rate": hit_rate, + "route_ahead_prefetch_coverage": coverage, + "wasted_prefetch_bytes": float(wasted_bytes), + "expert_occupancy_bytes": expert_occupancy, + "kv_occupancy_bytes": float(kv_occupancy), + } + + def _count_offloaded_experts(engine: Any) -> int: for attr in ("num_offloaded_experts", "offloaded_expert_count"): value = getattr(engine, attr, None) diff --git a/benchmarks/dflash/pd_dflash_serving.py b/benchmarks/dflash/pd_dflash_serving.py index ba9af83..009901c 100644 --- a/benchmarks/dflash/pd_dflash_serving.py +++ b/benchmarks/dflash/pd_dflash_serving.py @@ -24,9 +24,15 @@ the byte-accurate route-ahead ``wasted_prefetch_bytes`` from the instrumented ``RouteAheadStats`` (never an expert-count proxy). -At this Phase-A stage B2 (the 2-D deficit scheduler) does not yet exist, so the -runner emits an explicit ``BLOCKED_UNTIL_2D_SCHEDULER`` status for B2 rather than -silently emulating another baseline. +B2 (the 2-D deficit verify scheduler) is measured through the continuous- +batching serving path (``moe_infinity.serving.engine.ContinuousBatchingEngine``) +rather than ``MoE.generate()``: configuring the four Task-6 verify budgets flips +the engine onto ``_step_speculative_session``, so the 2-D admission scheduler +governs every VERIFY round and the serving ``PagedKVCache`` yields a real +``kv_occupancy_bytes`` (the sync path has no serving KV manager and falls back to +``0.0``). This runner stays opt-in (GPU-gated) and never changes default serving +behaviour. If ``--baseline B2`` runs before the serving wiring is available the +runner still records an honest ``BLOCKED_UNTIL_2D_SCHEDULER`` status. """ from __future__ import annotations @@ -68,6 +74,15 @@ OFFLOADED_BASELINES: Tuple[str, ...] = ("B0", "B1", "B2") BLOCKED_STATUS = "BLOCKED_UNTIL_2D_SCHEDULER" +SERVING_BASELINES: Tuple[str, ...] = ("B2",) + +# B2 == "no expert-byte coupling": admission is gated only by the token +# dimension, so this stands in for an unbounded expert-byte budget/cap. It must +# exceed any single route-ahead union's summed FP4 payload so admit_verify_ +# demands never rejects on bytes and its "cap >= largest single demand" guard +# always holds. +B2_UNBOUNDED_EXPERT_BYTES = 1 << 60 + RTX_PRO_6000_NAME_FRAGMENT = "RTX PRO 6000" RTX_PRO_6000_CAPABILITY: Tuple[int, int] = (12, 0) @@ -139,6 +154,45 @@ def require_offloaded(baseline: str, offloaded_expert_count: int) -> None: ) +def build_b2_serving_config( + *, + block_size: int, + concurrency: int, + num_layers: int, + num_kv_heads: int, + head_dim: int, + dtype: str, + eos_token_id: Optional[int], + num_kv_blocks: int, + device_memory_ratio: float, +) -> Dict[str, Any]: + """Return the ``ContinuousBatchingEngine`` config that selects the B2 path. + + Setting all four verify budgets makes ``Scheduler.verify_scheduling_enabled`` + true, so the engine drives ``_step_speculative_session`` (the 2-D admission + scheduler) per VERIFY round. The token budget/cap equal ``block_size`` (one + verify block per round) while the expert-byte budget/cap are unbounded, which + is exactly B2's "token-deficit scheduler and no expert-byte coupling". + """ + return { + "device_memory_ratio": float(device_memory_ratio), + "kv_cache_ratio": 0.01, + "max_batch_size": max(1, int(concurrency)), + "max_tokens_per_step": 2048, + "block_size": int(block_size), + "num_layers": int(num_layers), + "num_kv_heads": int(num_kv_heads), + "head_dim": int(head_dim), + "dtype": dtype, + "eos_token_id": eos_token_id, + "num_kv_blocks": int(num_kv_blocks), + "verify_token_budget": int(block_size), + "verify_expert_byte_budget": B2_UNBOUNDED_EXPERT_BYTES, + "verify_token_deficit_cap": int(block_size), + "verify_expert_byte_deficit_cap": B2_UNBOUNDED_EXPERT_BYTES, + } + + def observation_key( model: str, baseline: str, block_size: int, concurrency: int, repeat: int ) -> Tuple[str, str, int, int, int]: @@ -399,7 +453,10 @@ def run_experiment(args: RunnerArgs) -> int: and appends one row per ``(model, baseline, block, concurrency, repeat)``. B2 is emitted as ``BLOCKED_UNTIL_2D_SCHEDULER`` until Task 6 lands. """ - from benchmarks.dflash._serving_measure import measure_configuration + from benchmarks.dflash._serving_measure import ( + measure_configuration, + measure_configuration_serving, + ) assert args.model and args.draft and args.offload_dir and args.output _validate_gpu_environment() @@ -408,28 +465,22 @@ def run_experiment(args: RunnerArgs) -> int: for baseline in args.baselines: for block_size in args.block_sizes: for concurrency in args.concurrency: - if baseline == "B2": - append_observation( - args.output, - make_observation_row( - model=args.model, - draft=draft, - baseline=baseline, - block_size=block_size, - concurrency=concurrency, - repeat=0, - metrics={}, - status=BLOCKED_STATUS, - ), + if baseline in SERVING_BASELINES: + row = measure_configuration_serving( + args=args, + baseline=baseline, + draft=draft, + block_size=block_size, + concurrency=concurrency, + ) + else: + row = measure_configuration( + args=args, + baseline=baseline, + draft=draft, + block_size=block_size, + concurrency=concurrency, ) - continue - row = measure_configuration( - args=args, - baseline=baseline, - draft=draft, - block_size=block_size, - concurrency=concurrency, - ) append_observation(args.output, row) return 0 diff --git a/tests/python/dflash/test_pd_dflash_serving_contract.py b/tests/python/dflash/test_pd_dflash_serving_contract.py index d38571b..4965ee6 100644 --- a/tests/python/dflash/test_pd_dflash_serving_contract.py +++ b/tests/python/dflash/test_pd_dflash_serving_contract.py @@ -14,11 +14,13 @@ import pytest from benchmarks.dflash.pd_dflash_serving import ( + B2_UNBOUNDED_EXPERT_BYTES, BLOCKED_STATUS, REQUIRED_CONCURRENCY, REQUIRED_DRAFTS, REQUIRED_MODELS, append_observation, + build_b2_serving_config, build_contract_matrix, load_observations, main, @@ -131,6 +133,36 @@ def test_blocked_row_carries_status_and_no_metrics(): assert "output_tokens_per_second" not in row +def test_b2_serving_config_enables_token_only_verify_scheduler(): + from moe_infinity.serving.scheduler import _resolve_verify_config + + config = build_b2_serving_config( + block_size=16, + concurrency=4, + num_layers=24, + num_kv_heads=8, + head_dim=64, + dtype="bfloat16", + eos_token_id=7, + num_kv_blocks=128, + device_memory_ratio=0.85, + ) + assert config["verify_token_budget"] == 16 + assert config["verify_token_deficit_cap"] == 16 + assert config["verify_expert_byte_budget"] == B2_UNBOUNDED_EXPERT_BYTES + assert config["verify_expert_byte_deficit_cap"] == B2_UNBOUNDED_EXPERT_BYTES + assert config["max_batch_size"] == 4 + assert config["block_size"] == 16 + + resolved = _resolve_verify_config( + token_budget=config["verify_token_budget"], + expert_byte_budget=config["verify_expert_byte_budget"], + token_deficit_cap=config["verify_token_deficit_cap"], + expert_byte_deficit_cap=config["verify_expert_byte_deficit_cap"], + ) + assert resolved.enabled is True + + # --------------------------------------------------------------------------- # append-without-overwrite JSON writer # ---------------------------------------------------------------------------