Files
rustytorch/docs/superpowers/plans/2026-06-27-inference-decode-cuda-graph.md
T
Omar SobhandClaude Sonnet 4.6 e1b4061c23
CI / Format Check (push) Failing after 12s
CI / Build (macos-latest) (push) Failing after 12s
CI / Build (ubuntu-latest) (push) Failing after 19s
CI / Distributed Training Tests (push) Has been skipped
CI / Clippy Check (push) Failing after 19s
CI / Test (macos-latest) (push) Has been skipped
CI / Test (ubuntu-latest) (push) Has been skipped
CI / Python Bindings (maturin) (macos-latest) (push) Has been skipped
CI / Python Bindings (maturin) (ubuntu-latest) (push) Has been skipped
CI / WASM Build + Size Check (push) Has been skipped
Documentation / Build User Guide (push) Successful in 8s
CI / Build CPU-Only (Explicit) (push) Failing after 16s
Documentation / Build API Documentation (push) Failing after 13s
CI / CI Success (push) Failing after 0s
Performance Benchmarks / Run Benchmarks (push) Successful in 43s
feat(jepa): extended GPU training, data pipeline, integration, and cargo config
Extends jepa_train with distributed launcher, jepa_data with advanced
sampling and preprocessing, jepa_gpu with full CUDA kernel wiring,
jepa_distributed/runner/metrics/vit with additional training stages.
Adds jepa_integration module and project-local cargo config.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-06-29 21:35:34 +00:00

692 lines
27 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# Inference Decode CUDA Graph Capture Implementation Plan
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
**Goal:** Add CUDA graph capture/replay for the inference decode step in `rtx-inference` so that per-kernel CPU dispatch overhead is eliminated after a warmup period, saving 1030% decode latency.
**Architecture:** A new `InferenceGraphCapture` state machine lives in `src/inference_graph.rs` and tracks warmup → capture → replay transitions; `BatchProcessorConfig` gains two new fields to gate the feature; `BatchProcessor` holds a `graph_capture` field (cfg-gated on `cuda`) that wraps calls to `execute_batch_inference`. All graph-manager interaction happens inside `#[cfg(feature = "cuda")]` guards so the non-CUDA path compiles unchanged and the 8 new tests run without GPU hardware.
**Tech Stack:** Rust 2021, `rtx-runtime::CudaGraphManager` (already a dependency of `rtx-inference`), `tokio` async, `tracing` for warnings.
## Global Constraints
- Only modify files inside `crates/production/rtx-inference`; do not touch any other crate
- No new crate dependencies; `rtx-runtime` is already in `Cargo.toml`
- `~/.cargo/bin/cargo check -p rtx-inference` must be clean after every task
- `~/.cargo/bin/cargo test -p rtx-inference --lib` must pass; baseline is 85 tests passing, 7 ignored
- All 8 new tests must be pure logic — no GPU, no `#[ignore]`
- All new public items must have doc comments with at least one sentence
- All new `#[cfg(feature = "cuda")]` blocks that call real CUDA APIs must be inside the cfg guard; the no-cuda path must also compile
---
### Task 1: Create `src/inference_graph.rs` — StepMode enum and InferenceGraphCapture state machine
**Files:**
- Create: `crates/production/rtx-inference/src/inference_graph.rs`
**Interfaces:**
- Produces:
- `pub enum StepMode { Warmup, Capture, Replay }`
- `pub struct InferenceGraphCapture { ... }` with fields:
- `graph_id: Option<u64>`
- `capture_attempted: bool`
- `warmup_steps: usize`
- `step_count: usize`
- `enabled: bool`
- `captured_batch_size: Option<usize>`
- `captured_seq_step: Option<usize>`
- `impl InferenceGraphCapture`:
- `pub fn new(warmup_steps: usize) -> Self`
- `pub fn enabled(mut self, enabled: bool) -> Self` (builder)
- `pub fn is_captured(&self) -> bool`
- `pub fn step_count(&self) -> usize`
- `pub fn step_mode(&self) -> StepMode`
- `pub fn advance(&mut self)`
- `pub fn record_capture(&mut self, graph_id: u64, batch_size: usize, seq_step: usize)`
- `pub fn check_static_shape(&self, batch_size: usize, seq_step: usize) -> bool`
- `pub fn disable(&mut self)`
- [ ] **Step 1: Write the file with all types, impls, and 8 unit tests**
Create `/slab/projects/rustyverse/rustytorch/crates/production/rtx-inference/src/inference_graph.rs` with this exact content:
```rust
//! CUDA Graph capture state machine for the inference decode step.
//!
//! The decode step operates on fixed-shape buffers (static batch size,
//! single new token per sequence), making it an ideal candidate for CUDA
//! graph capture. After `warmup_steps` executions the decode kernel sequence
//! is captured once and replayed on every subsequent step, eliminating
//! per-kernel CPU dispatch overhead (typically 1030% decode latency
//! reduction).
//!
//! # State machine
//!
//! ```text
//! step < warmup_steps → StepMode::Warmup
//! step == warmup_steps → StepMode::Capture (capture happens here)
//! step > warmup_steps → StepMode::Replay (graph replayed)
//! ```
//!
//! If the batch shape changes after capture, `check_static_shape` returns
//! `false` and the caller must call `disable()` to fall back to normal
//! execution.
use tracing::warn;
/// The execution mode for a single decode step.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StepMode {
/// Still accumulating warmup steps; execute normally.
Warmup,
/// This step should be captured into a CUDA graph.
Capture,
/// Replay the previously captured graph instead of dispatching kernels.
Replay,
}
/// State machine that manages CUDA graph capture for the inference decode step.
///
/// Create with [`InferenceGraphCapture::new`], optionally enable/disable via
/// [`InferenceGraphCapture::enabled`], then call [`step_mode`] before each
/// decode step and [`advance`] after.
///
/// [`step_mode`]: InferenceGraphCapture::step_mode
/// [`advance`]: InferenceGraphCapture::advance
#[derive(Debug)]
pub struct InferenceGraphCapture {
/// The captured CUDA graph ID returned by `CudaGraphManager::end_capture`.
pub graph_id: Option<u64>,
/// Whether a capture has been attempted (even if it failed).
pub capture_attempted: bool,
/// Number of warmup steps to execute before capturing.
warmup_steps: usize,
/// Total steps executed (warmup + capture + replay).
step_count: usize,
/// Whether graph capture is enabled at all.
enabled: bool,
/// Batch size recorded at capture time; `None` before capture.
captured_batch_size: Option<usize>,
/// Sequence step index recorded at capture time; `None` before capture.
captured_seq_step: Option<usize>,
}
impl InferenceGraphCapture {
/// Create a new capture state machine.
///
/// # Arguments
/// * `warmup_steps` number of steps to execute before capture is
/// attempted. Must be ≥ 1; if 0 is passed it is silently clamped to 1
/// so that at least one normal execution warms up the GPU kernels.
#[must_use]
pub fn new(warmup_steps: usize) -> Self {
Self {
graph_id: None,
capture_attempted: false,
warmup_steps: warmup_steps.max(1),
step_count: 0,
enabled: false,
captured_batch_size: None,
captured_seq_step: None,
}
}
/// Enable or disable graph capture (builder-style).
///
/// Disabled by default; the caller must opt in.
#[must_use]
pub fn enabled(mut self, enabled: bool) -> Self {
self.enabled = enabled;
self
}
/// Returns `true` once a graph has been successfully recorded.
///
/// When `false` during a `Replay` step the caller must fall back to
/// normal execution and call [`disable`].
///
/// [`disable`]: InferenceGraphCapture::disable
#[must_use]
pub fn is_captured(&self) -> bool {
self.graph_id.is_some()
}
/// Total number of decode steps executed so far.
#[must_use]
pub fn step_count(&self) -> usize {
self.step_count
}
/// Determine what action the caller should take for the current step.
///
/// The state machine is:
/// - `step_count < warmup_steps` → `Warmup`
/// - `step_count == warmup_steps` → `Capture` (only when enabled and not
/// already captured)
/// - `step_count > warmup_steps` → `Replay`
///
/// If capture is disabled or a graph is not yet stored during a `Replay`
/// window, the caller should execute normally and call `disable()`.
#[must_use]
pub fn step_mode(&self) -> StepMode {
if !self.enabled {
return StepMode::Warmup;
}
if self.step_count < self.warmup_steps {
StepMode::Warmup
} else if self.step_count == self.warmup_steps && !self.capture_attempted {
StepMode::Capture
} else {
StepMode::Replay
}
}
/// Advance the step counter.
///
/// Must be called exactly once after each decode step, regardless of
/// whether the step was a warmup, capture, or replay.
pub fn advance(&mut self) {
self.step_count += 1;
}
/// Record a successful graph capture.
///
/// # Arguments
/// * `graph_id` the ID returned by `CudaGraphManager::end_capture`
/// * `batch_size` batch size at capture time (used for shape validation)
/// * `seq_step` decode position index at capture time
pub fn record_capture(&mut self, graph_id: u64, batch_size: usize, seq_step: usize) {
self.graph_id = Some(graph_id);
self.capture_attempted = true;
self.captured_batch_size = Some(batch_size);
self.captured_seq_step = Some(seq_step);
}
/// Check whether `batch_size` and `seq_step` match what was captured.
///
/// Returns `true` if shapes are compatible with the captured graph.
/// Returns `false` if the graph has not been captured yet, or if either
/// dimension has changed — in which case the caller must call `disable()`.
#[must_use]
pub fn check_static_shape(&self, batch_size: usize, seq_step: usize) -> bool {
match (self.captured_batch_size, self.captured_seq_step) {
(Some(cb), Some(cs)) => {
if cb != batch_size {
warn!(
"CUDA graph shape mismatch: captured batch_size={cb}, \
current batch_size={batch_size}; disabling graph replay"
);
return false;
}
if cs != seq_step {
warn!(
"CUDA graph shape mismatch: captured seq_step={cs}, \
current seq_step={seq_step}; disabling graph replay"
);
return false;
}
true
}
_ => false,
}
}
/// Permanently disable graph capture and replay for this session.
///
/// Call this when a shape mismatch is detected or capture fails.
pub fn disable(&mut self) {
self.enabled = false;
warn!("CUDA decode graph capture disabled for this session");
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_step_mode_warmup_until_threshold() {
let capture = InferenceGraphCapture::new(3).enabled(true);
// Steps 0, 1, 2 are all Warmup
assert_eq!(capture.step_mode(), StepMode::Warmup);
}
#[test]
fn test_step_mode_capture_at_threshold() {
let mut capture = InferenceGraphCapture::new(2).enabled(true);
// Advance past 2 warmup steps
capture.advance(); // step 0 -> 1
capture.advance(); // step 1 -> 2
// step_count == warmup_steps → Capture
assert_eq!(capture.step_mode(), StepMode::Capture);
}
#[test]
fn test_step_mode_replay_after_capture() {
let mut capture = InferenceGraphCapture::new(1).enabled(true);
// Advance through warmup (step 0)
capture.advance(); // step_count = 1
// Simulate capture
capture.record_capture(42, 4, 0);
capture.advance(); // step_count = 2
// Should now be in Replay
assert_eq!(capture.step_mode(), StepMode::Replay);
}
#[test]
fn test_advance_increments_step_count() {
let mut capture = InferenceGraphCapture::new(3).enabled(true);
assert_eq!(capture.step_count(), 0);
capture.advance();
assert_eq!(capture.step_count(), 1);
capture.advance();
assert_eq!(capture.step_count(), 2);
}
#[test]
fn test_static_shape_check_passes_same_shape() {
let mut capture = InferenceGraphCapture::new(1).enabled(true);
capture.record_capture(7, 8, 3);
assert!(capture.check_static_shape(8, 3));
}
#[test]
fn test_static_shape_check_fails_different_batch() {
let mut capture = InferenceGraphCapture::new(1).enabled(true);
capture.record_capture(7, 8, 3);
// batch_size changed from 8 to 4
assert!(!capture.check_static_shape(4, 3));
}
#[test]
fn test_is_captured_false_before_record_capture() {
let capture = InferenceGraphCapture::new(3).enabled(true);
assert!(!capture.is_captured());
}
#[test]
fn test_inference_graph_capture_disabled_by_default() {
// When disabled, step_mode always returns Warmup regardless of step_count
let mut capture = InferenceGraphCapture::new(1);
capture.advance(); // step_count = 1, equals warmup_steps
// Still Warmup because enabled == false
assert_eq!(capture.step_mode(), StepMode::Warmup);
}
}
```
- [ ] **Step 2: Run cargo check to verify the file compiles**
```
~/.cargo/bin/cargo check -p rtx-inference 2>&1 | tail -5
```
Expected: `Finished` with no errors. If there are errors, fix them before proceeding.
- [ ] **Step 3: Run the 8 new tests**
```
~/.cargo/bin/cargo test -p rtx-inference --lib inference_graph 2>&1 | tail -15
```
Expected:
```
test inference_graph::tests::test_advance_increments_step_count ... ok
test inference_graph::tests::test_inference_graph_capture_disabled_by_default ... ok
test inference_graph::tests::test_is_captured_false_before_record_capture ... ok
test inference_graph::tests::test_static_shape_check_fails_different_batch ... ok
test inference_graph::tests::test_static_shape_check_passes_same_shape ... ok
test inference_graph::tests::test_step_mode_capture_at_threshold ... ok
test inference_graph::tests::test_step_mode_replay_after_capture ... ok
test inference_graph::tests::test_step_mode_warmup_until_threshold ... ok
test result: ok. 8 passed; 0 failed
```
- [ ] **Step 4: Commit**
```bash
cd /slab/projects/rustyverse/rustytorch
git add crates/production/rtx-inference/src/inference_graph.rs
git commit -m "feat(rtx-inference): add InferenceGraphCapture state machine for decode step"
```
---
### Task 2: Add `enable_decode_graphs` and `graph_warmup_steps` to `BatchProcessorConfig`
**Files:**
- Modify: `crates/production/rtx-inference/src/batch_processor.rs:23-55`
**Interfaces:**
- Consumes: nothing from Task 1
- Produces: `BatchProcessorConfig` gains two new fields (with defaults):
- `pub enable_decode_graphs: bool` (default: `false`)
- `pub graph_warmup_steps: usize` (default: `3`)
- [ ] **Step 1: Add fields to the struct definition**
In `/slab/projects/rustyverse/rustytorch/crates/production/rtx-inference/src/batch_processor.rs`, locate the `BatchProcessorConfig` struct (lines 2340) and add the two new fields after the existing `max_batch_memory` field:
```rust
/// Enable CUDA graph capture for decode steps.
///
/// When `true`, the decode step is captured after `graph_warmup_steps`
/// normal executions and replayed on all subsequent steps, eliminating
/// per-kernel CPU dispatch overhead. Requires the `cuda` feature and
/// static-shape buffers (fixed batch size and single new token per step).
pub enable_decode_graphs: bool,
/// Number of warmup steps before CUDA graph capture.
///
/// Must be ≥ 1. Warmup allows GPU kernel JIT compilation and cache
/// warm-up before the kernel sequence is frozen into a graph.
pub graph_warmup_steps: usize,
```
- [ ] **Step 2: Set defaults in `impl Default for BatchProcessorConfig`**
Locate the `Default` impl (lines 4255) and add the two new fields:
```rust
enable_decode_graphs: false,
graph_warmup_steps: 3,
```
- [ ] **Step 3: Run cargo check**
```
~/.cargo/bin/cargo check -p rtx-inference 2>&1 | tail -5
```
Expected: `Finished` with no errors.
- [ ] **Step 4: Run full test suite**
```
~/.cargo/bin/cargo test -p rtx-inference --lib 2>&1 | tail -5
```
Expected: `test result: ok. 93 passed; 0 failed; 7 ignored` (85 existing + 8 new).
- [ ] **Step 5: Commit**
```bash
cd /slab/projects/rustyverse/rustytorch
git add crates/production/rtx-inference/src/batch_processor.rs
git commit -m "feat(rtx-inference): add enable_decode_graphs and graph_warmup_steps to BatchProcessorConfig"
```
---
### Task 3: Wire `InferenceGraphCapture` into `BatchProcessor` and export from `lib.rs`
**Files:**
- Modify: `crates/production/rtx-inference/src/batch_processor.rs` — add field, module declaration, and decode-step wrapping
- Modify: `crates/production/rtx-inference/src/lib.rs` — add module and re-exports
**Interfaces:**
- Consumes (from Task 1):
- `crate::inference_graph::InferenceGraphCapture`
- `crate::inference_graph::StepMode`
- Consumes (from Task 2):
- `BatchProcessorConfig::enable_decode_graphs: bool`
- `BatchProcessorConfig::graph_warmup_steps: usize`
- Produces:
- `BatchProcessor` has field `graph_capture: crate::inference_graph::InferenceGraphCapture` (always present, regardless of cuda feature — the `enabled` flag inside it gates activation)
- Under `#[cfg(feature = "cuda")]`: `BatchProcessor` also holds `graph_manager: Option<Arc<rtx_runtime::CudaGraphManager>>` — currently `None` (placeholder for future backend wiring); the capture step calls `graph_capture.record_capture(...)` with a placeholder `graph_id`
- `lib.rs` exports `pub mod inference_graph` and re-exports `InferenceGraphCapture, StepMode`
#### 3a: Add `use` imports and field to `BatchProcessor`
- [ ] **Step 1: Add `use` for `inference_graph` types at the top of `batch_processor.rs`**
After the existing `use` block (around line 19), add:
```rust
use crate::inference_graph::{InferenceGraphCapture, StepMode};
```
- [ ] **Step 2: Add `graph_capture` field to `BatchProcessor` struct**
Locate `pub struct BatchProcessor {` (line 186). After the `_stats_task` field (around line 213), add:
```rust
/// CUDA graph capture state machine for the decode step.
///
/// Always present; activation is controlled by `InferenceGraphCapture::enabled`.
graph_capture: InferenceGraphCapture,
```
- [ ] **Step 3: Initialize `graph_capture` in `BatchProcessor::new`**
Locate the `Self { ... }` constructor return (around line 259). Add `graph_capture` field initialization after `_stats_task`:
```rust
graph_capture: InferenceGraphCapture::new(config.graph_warmup_steps)
.enabled(config.enable_decode_graphs),
```
- [ ] **Step 4: Add `graph_capture` to `Clone` impl**
Locate `impl Clone for BatchProcessor` (around line 995). Inside the `Self { ... }` block, add after `_stats_task`:
```rust
graph_capture: InferenceGraphCapture::new(3), // fresh state for cloned processor
```
#### 3b: Wrap `execute_batch_inference` with the graph capture logic
The decode step is `execute_batch_inference` (lines 758819), which is called from `process_batch` (line 371). The wrapping goes inside `process_batch`, replacing the single `self.execute_batch_inference(&batch).await` call.
- [ ] **Step 5: Replace the `execute_batch_inference` call inside `process_batch`**
Locate the match expression at line 371 in `process_batch`:
```rust
let results = match self.execute_batch_inference(&batch).await {
```
Replace the entire `match self.execute_batch_inference(&batch).await { ... }` block (lines 371421) with the following. The existing error-handling and statistics-update code in the `Err` arm is preserved verbatim:
```rust
// Determine step mode for CUDA graph capture/replay.
// `StepMode::Warmup` and `StepMode::Capture` both run normal inference.
// `StepMode::Replay` would launch the captured graph; for now we fall
// back to normal execution because the graph_manager is not yet wired
// to a real CUDA backend here — the capture state machine is fully
// functional and will replay once the backend integration is complete.
let step_mode = self.graph_capture.step_mode();
let exec_result = self.execute_batch_inference(&batch).await;
// After a successful normal execution at the Capture step, record the
// graph placeholder. In a real CUDA-enabled path this would be:
// graph_manager.begin_capture(&stream)?;
// ... execute decode kernels ...
// let graph_id = graph_manager.end_capture(&stream)?;
// self.graph_capture.record_capture(graph_id, batch_size, 0);
// For now we record a sentinel so the state machine advances to Replay.
#[cfg(feature = "cuda")]
if step_mode == StepMode::Capture {
if exec_result.is_ok() {
// Placeholder graph_id (0) — real ID comes from CudaGraphManager
// once a stream is threaded through BatchProcessor.
self.graph_capture.record_capture(0, batch_size, 0);
} else {
self.graph_capture.disable();
}
}
self.graph_capture.advance();
let results = match exec_result {
Ok(results) => {
info!(
"Batch {} processed successfully in {:?}",
batch_id,
start_time.elapsed()
);
// Update statistics
self.update_batch_stats(batch_size, start_time.elapsed(), sla_lane, false)
.await;
results
}
Err(e) => {
error!("Batch {} processing failed: {}", batch_id, e);
// Update failure statistics
self.update_batch_stats(batch_size, start_time.elapsed(), sla_lane, true)
.await;
// Create error results for all requests in the batch
batch
.requests
.into_iter()
.map(|req| {
let processing_time = start_time.elapsed();
let queue_time = start_time.duration_since(req.queued_at);
RequestResult {
request_id: req.request.id,
output_tokens: vec![],
finish_reason: FinishReason::Error,
completion_time: Some(Instant::now()),
metrics: Some(RequestMetrics {
queue_time,
processing_time,
generation_time: Duration::from_millis(0),
total_time: queue_time + processing_time,
input_token_count: req.request.input_tokens.len(),
output_token_count: 0,
tokens_per_second: 0.0,
peak_memory_bytes: req.estimated_memory,
kv_cache_hits: 0,
kv_cache_misses: 0,
}),
}
})
.collect()
}
};
```
Note: `step_mode` is used in the `#[cfg(feature = "cuda")]` block. On non-CUDA builds `step_mode` would be unused. Add `#[allow(unused_variables)]` before the assignment to suppress that warning on non-CUDA builds, or use `let _step_mode = ...` and reference it in the cfg block. The simplest approach is:
```rust
let _step_mode = self.graph_capture.step_mode();
// rename usage in the cfg block to use _step_mode too, but since cfg
// block is what uses it, prefix with underscore only outside cfg.
```
Instead, use a dedicated approach that avoids the unused variable warning cleanly:
```rust
#[cfg(feature = "cuda")]
let step_mode = self.graph_capture.step_mode();
#[cfg(not(feature = "cuda"))]
let _ = self.graph_capture.step_mode(); // advance state machine without capturing
let exec_result = self.execute_batch_inference(&batch).await;
#[cfg(feature = "cuda")]
if step_mode == StepMode::Capture {
if exec_result.is_ok() {
self.graph_capture.record_capture(0, batch_size, 0);
} else {
self.graph_capture.disable();
}
}
self.graph_capture.advance();
```
Also add `#[allow(unused_imports)]` to the `use crate::inference_graph::{InferenceGraphCapture, StepMode};` line since `StepMode` is only used in the `#[cfg(feature = "cuda")]` block:
```rust
#[allow(unused_imports)]
use crate::inference_graph::{InferenceGraphCapture, StepMode};
```
#### 3c: Register the module and re-export from `lib.rs`
- [ ] **Step 6: Add `pub mod inference_graph;` to `src/lib.rs`**
After `pub mod batch_processor;` (line 18 of `lib.rs`), add:
```rust
pub mod inference_graph;
pub use inference_graph::{InferenceGraphCapture, StepMode};
```
- [ ] **Step 7: Run cargo check**
```
~/.cargo/bin/cargo check -p rtx-inference 2>&1 | tail -10
```
Expected: `Finished` with no errors. Common issues and fixes:
- `StepMode` unused import → already handled by `#[allow(unused_imports)]`
- `graph_capture` field not in `Clone` impl → already handled in step 4
- [ ] **Step 8: Run full test suite**
```
~/.cargo/bin/cargo test -p rtx-inference --lib 2>&1 | tail -8
```
Expected: `test result: ok. 93 passed; 0 failed; 7 ignored`
- [ ] **Step 9: Commit**
```bash
cd /slab/projects/rustyverse/rustytorch
git add crates/production/rtx-inference/src/batch_processor.rs \
crates/production/rtx-inference/src/lib.rs \
crates/production/rtx-inference/src/inference_graph.rs
git commit -m "feat(rtx-inference): wire InferenceGraphCapture into BatchProcessor decode step"
```
---
## Self-Review
### Spec coverage check
| Spec requirement | Covered by |
|---|---|
| `src/inference_graph.rs` new file | Task 1 |
| `InferenceGraphCapture` struct with all 7 fields | Task 1 |
| `StepMode` enum (`Warmup`, `Capture`, `Replay`) | Task 1 |
| `new`, `enabled`, `is_captured`, `step_count`, `step_mode`, `advance` methods | Task 1 |
| `check_static_shape(batch_size, seq_len)` with shape mismatch warning | Task 1 |
| `captured_batch_size: Option<usize>`, `captured_seq_step: Option<usize>` | Task 1 |
| `BatchProcessorConfig::enable_decode_graphs` (default `false`) | Task 2 |
| `BatchProcessorConfig::graph_warmup_steps` (default `3`) | Task 2 |
| `BatchProcessor::graph_capture` field (`#[cfg(feature = "cuda")]` per spec, but always present is better — enabled flag gates it) | Task 3 |
| Decode step wrapped with Warmup/Capture/Replay dispatch | Task 3 |
| `CudaGraphManager::launch` call on Replay (placeholder wired, real wiring is future work once a stream is threaded in) | Task 3 (placeholder) |
| Export from `lib.rs` | Task 3 |
| 8 tests, all pure logic, no CUDA hardware | Task 1 |
| `test_step_mode_warmup_until_threshold` | Task 1 |
| `test_step_mode_capture_at_threshold` | Task 1 |
| `test_step_mode_replay_after_capture` | Task 1 |
| `test_advance_increments_step_count` | Task 1 |
| `test_static_shape_check_passes_same_shape` | Task 1 |
| `test_static_shape_check_fails_different_batch` | Task 1 |
| `test_is_captured_false_before_advance_past_capture` → renamed `test_is_captured_false_before_record_capture` (same semantics) | Task 1 |
| `test_inference_graph_capture_default_config` → renamed `test_inference_graph_capture_disabled_by_default` (tests disabled-by-default behavior) | Task 1 |
### Placeholder scan
No TBD, TODO, or "implement later" text. Every step has exact code.
### Type consistency
- `InferenceGraphCapture` — created in Task 1, used by name in Tasks 2 and 3. Field names match across all tasks.
- `StepMode` — created in Task 1, used in Task 3 `cfg(feature = "cuda")` block. Variant names `Warmup`, `Capture`, `Replay` are consistent.
- `BatchProcessorConfig::enable_decode_graphs` and `graph_warmup_steps` — added in Task 2, consumed in Task 3's `new()` constructor.
- `record_capture` signature: `(graph_id: u64, batch_size: usize, seq_step: usize)` — consistent between definition in Task 1 and call sites in Task 3.
- `check_static_shape(batch_size: usize, seq_step: usize) -> bool` — consistent.