DFlashProposer: platform-aware peak memory measurement (CUDA / MPS / CPU) — Step 3a

[FluffyAIcode]

Why this PR (Step 3a of post-merge plan)

PR #93's DFlashProposer.propose_block records peak activation bytes via:

peak = 0
if torch.cuda.is_available():
    torch.cuda.reset_peak_memory_stats()
tokens = self.drafter.draft_block(...)
if torch.cuda.is_available():
    peak = int(torch.cuda.max_memory_allocated())
return BlockProposal(..., peak_activation_bytes=peak)

This silently returns 0 on Mac MPS / CPU. The Mac MLX speculative-decoding eval (next PR, Step 3b) needs honest peak memory numbers on Apple Silicon for the BlockProposal accounting to be meaningful.

Fix

Three module-level helpers in dflash_drafter.py that dispatch by torch device type:

Helper	CUDA	MPS	CPU
_detect_device(model)	'cuda'	'mps'	'cpu'
_reset_peak_memory(device)	torch.cuda.reset_peak_memory_stats()	no-op	no-op
_peak_memory_bytes(device)	torch.cuda.max_memory_allocated()	torch.mps.driver_allocated_memory() (try/except → 0 on RuntimeError)	0 (signal: unmeasured, NOT a fake peak)

DFlashProposer.propose_block rewired to call these. CUDA path semantics unchanged (same helpers, same calls, same output values). MPS/CPU paths now produce honest values.

Documented caveats (inline)

• MPS has no peak counter. We use post-forward driver_allocated_memory as a tight upper bound on activations released after the forward — close enough for spec-decode-loop memory accounting in single-process scenarios. Stricter delta measurement requires the caller to snapshot before/after.
• CPU returns 0 deliberately (signal: unmeasured) rather than lying with a fake measurement. CPU peak measurement (psutil RSS or tracemalloc) is a different problem outside the scope of activation-byte accounting in BlockProposal.

Tests (TestPlatformAwarePeakMemory, 8 new tests)

test_detect_device_cpu                                   ✓
test_detect_device_raises_on_empty_model                 ✓ (defensive)
test_peak_memory_bytes_cpu_returns_zero                  ✓
test_peak_memory_bytes_unknown_device_returns_zero       ✓
test_reset_peak_memory_cpu_is_noop                       ✓
test_propose_block_records_zero_peak_on_cpu              ✓ (end-to-end)
test_peak_memory_bytes_mps_calls_driver_allocated_memory ✓ (synthetic torch.mps swap)
test_peak_memory_bytes_mps_handles_runtime_failure       ✓ (raise → 0)

Verified by stash-revert: 7 of 8 tests fail when fix is reverted (the 8th — empty-model raise — passes by luck because raise-on-empty was the original behaviour). Un-stash → 28/28.

tests/inference_engine/v04/: 315 passed (307 pre-existing + 8 new)

Why split Step 3 into 3a + 3b

• 3a (this PR) lands a small, fully-testable improvement to PR PR-K3 (Stage 1): native DFlash drafter module + proposer + Linux tests #93's already-merged code. Useful regardless of when 3b lands. Linux CI exercises the platform dispatch logic without requiring Apple Silicon.

• 3b (Mac MLX speculative decoding eval + reviewer aid) needs MLX bridge code: mx.array → torch.Tensor for hiddens, embed_fn / lm_head_fn callbacks that span the MLX verifier and PyTorch drafter runtimes. Writing that without ability to verify against a working mlx_lm verifier load = the same fake/fallback pattern the user just got us out of with PR PR-K3 (Stage 1): native DFlash drafter module + proposer + Linux tests #93. Better to wait for Step 4 evidence (user re-runs PR Retire PR #95/#96/#97/#98 wrapper code; defer DFlashDrafter to PR #93 #99's smoke on Mac mini, the new diagnostic captures the actual mlx_lm Gemma 4 MoE failure traceback) before authoring 3b.

Stack

main (post #93 + #99 + #94 merge)
└── THIS (#100, Step 3a)

待做:
  Step 3b — Mac MLX spec decode eval (off main, after Step 4 evidence)
  Step 4  — mlx_lm Gemma 4 MoE fix (waiting on user-side Mac mini diagnostic re-run)
  Step 5  — K2.A backport PR (dtype/transpose/HF Cache contract bugs from PR #98)
  Step 6  — alignment training corpus expansion (vast time)

Net effect

PR #93's BlockProposal.peak_activation_bytes now reflects honest measurement on whichever device the drafter actually runs on, instead of always being 0 on non-CUDA hardware. Sets up Step 3b (Mac MLX eval) to produce meaningful memory accounting in results/research/k3_dflash_specdecode_mac_*.json once Step 4 unblocks the verifier load.