feat(iorails): add speculative streaming generation

docs/configure-rails/yaml-schema/guardrails-configuration.mdx

@@ -332,9 +332,12 @@ To decide whether Speculative Generation makes sense for your use-case, explore
 <Warning title="Experimental Feature">
     Speculative generation currently requires the opt-in IORails engine. To
     enable IORails, set `NEMO_GUARDRAILS_IORAILS_ENGINE=1`. Speculative
-    generation is supported only for non-streaming requests (`generate_async`).
-    When speculative generation is enabled, streaming requests (`stream_async`)
-    fall back to sequential execution and emit a warning.
+    generation is supported for both non-streaming requests (`generate_async`)
+    and streaming requests (`stream_async`). Streaming speculation only supports
+    check-first behavior: if `rails.output.streaming.stream_first` is `True`, it
+    is overridden to check-first for speculative requests (with a warning), since
+    tokens cannot reach the client before the input verdict is known. Streaming
+    speculation with output rails requires `rails.output.streaming.enabled: True`.
 </Warning>
 
 ### How Speculative Generation Works
@@ -367,6 +370,48 @@ The engine handles three outcomes:
 Output rails always run after the main LLM completes.
 Speculative generation does not change the output-rail path.
 
+### Streaming Speculative Generation
+
+For streaming requests (`stream_async`), speculative generation runs the input rails, the main LLM, and the output rails concurrently.
+While the input-rail verdict is still pending (the *speculation window*), validated tokens are held in a bounded in-memory buffer instead of being sent to the client — no token reaches the caller before the input is confirmed safe.
+
+1. Start the input rails and the main LLM stream in parallel. Output rails process tokens as they stream (check-first), at the LLM's natural generation rate.
+2. Hold output-rail-validated tokens in the awaiting-release buffer.
+3. Resolve the race:
+    - **Input rails pass:** flush the held buffer to the client and continue streaming normally.
+    - **Input rails reject:** discard the held buffer, abort the generation, and return the refusal message. No tokens are leaked.
+    - **Output rails reject during the speculation window:** return the refusal early and cancel the still-running input rails — the request is refused before the input verdict is known.
+
+Streaming speculation only supports check-first behavior. If `rails.output.streaming.stream_first` is `True`, it is overridden to check-first for speculative requests, and a warning is logged at engine startup.
+
+If the main LLM fails during the speculation window, any chunks already held in the buffer are discarded along with the request. Without speculation those chunks would have reached the client before the failure.
+
+If an input rail *raises* rather than rejecting — for example the safety model is unreachable — the request fails with a `generation_error` payload rather than a `guardrails_violation`, and increments the errors counter rather than the blocked counter. An outage and a refusal are reported as different things.
+
+<Warning title="Speculation without output rails">
+    Streaming speculation also runs when no output rails are configured. In that case there is nothing to validate tokens during the speculation window, so raw model tokens are held in the buffer and released once the input rails pass. The client receives no tokens at all until the input verdict, then receives the held tokens at once — expect a pause followed by a burst rather than a smooth stream. Responses on this path are **not** output-rail validated.
+</Warning>
+
+#### Tuning the Release Buffer
+
+The awaiting-release buffer is bounded by `rails.input.speculative_max_buffered_chunks` (default `4096`).
+
+The buffer only grows while the input rails are still pending — that is the speculation window. Once the verdict arrives, the buffer is flushed (on pass) or discarded (on reject) and the stream becomes a pass-through. So its size is simply how much the model produced while the input rails were still working:
+
+| Input-rail latency | Held at ~50 tokens/sec |
+| --- | --- |
+| 361 ms | ~20 chunks |
+| 5 s | ~250 chunks |
+| 80 s | ~4096 chunks (the default bound) |
+
+At the default the bound is effectively unreachable: a request would time out long before the buffer fills.
+
+Lowering the value makes the engine stop consuming and wait for the input verdict sooner. With a slow input rail, the client stream pauses at the bound and resumes when the verdict arrives — you trade stream smoothness for a smaller held buffer.
+
+<Note>
+    This setting does **not** cap total memory. While the gate waits, the main LLM keeps generating into the internal stream buffer, which applies no backpressure to the model. Resident memory is bounded by the maximum response length multiplied by the number of concurrent streams (up to 256), not by this value. Size hosts from expected response length and concurrency rather than from this setting.
+</Note>
+
 ### Configuration Example
 
 To enable speculative generation, set `speculative_generation: True` under `rails.input` in the `config.yml` file.

docs/observability/tracing/span-reference.mdx

@@ -219,11 +219,25 @@ Every span sets ERROR status and records the exception when one propagates throu
 
 When speculative generation is active, the request span also carries these attributes:
 
-| Attribute                                | Type    | When set                  | Description                                                     |
-| ---------------------------------------- | ------- | ------------------------- | --------------------------------------------------------------- |
-| `speculative_generation.mode_active`     | boolean | Speculative generation on | The value `true`.                                               |
-| `speculative_generation.first_completed` | string  | Speculative generation on | The branch that completed first: `input_rails` or `generation`. |
-| `speculative_generation.first_rejector`  | string  | Speculative generation on | The branch that rejected the request: `input_rails` or `none`.  |
+| Attribute                                              | Type    | When set                  | Description                                                                                                             |
+| ------------------------------------------------------ | ------- | ------------------------- | --------------------------------------------------------------------------------------------------------------------- |
+| `speculative_generation.mode_active`                   | boolean | Speculative generation on | The value `true`.                                                                                                     |
+| `speculative_generation.first_completed`               | string  | Speculative generation on | The branch that completed first: `input_rails`, `generation`, or `output_rails`.                                      |
+| `speculative_generation.first_rejector`                | string  | Speculative generation on | The branch that rejected the request: `input_rails`, `output_rails`, or `none`. Covers rail *rejections* only — a rail or generation *error* leaves this `none` and marks the span status `ERROR`. |
+| `speculative_generation.rails_duration_ms`             | float   | Speculative generation on | Wall-clock time for the input rails.                                                                                  |
+| `speculative_generation.generation_duration_ms`        | float   | Speculative generation on | Wall-clock time for the main LLM generation.                                                                          |
+| `speculative_generation.overlap_ms`                    | float   | Speculative generation on | Time the input rails and generation ran concurrently.                                                                 |
+| `speculative_generation.time_saved_ms`                 | float   | Speculative generation on | Estimated saving versus sequential execution (the overlap for safe requests, `0` for rejected requests).             |
+| `speculative_generation.cancellation_event`            | string  | Speculative generation on | Which task was cancelled: `none`, `generation_cancelled`, `input_rails_cancelled`, or `output_rails_cancelled`. `input_rails_cancelled` covers the output-rails rejection short-circuit as well as error paths. |
+| `speculative_generation.output_rails_early_reject`     | boolean | Streaming speculation, output rails configured | Output rails rejected during the speculation window (short-circuit before the input verdict). Absent when no output rails are configured. |
+| `speculative_generation.output_rails_speculation_chunks` | int   | Streaming speculation     | Chunks held during the speculation window. With output rails configured these are output-rail-validated; with input rails only they are raw model chunks. |
+| `speculative_generation.output_rails_wasted_chunks`    | int     | Streaming speculation     | Held chunks discarded because the request was rejected.                                                               |
+| `speculative_generation.release_queue_duration_ms`     | float   | Streaming speculation     | Time chunks were held in the awaiting-release buffer.                                                                 |
+| `speculative_generation.release_queue_token_count`     | int     | Streaming speculation     | Number of chunks held awaiting the input-rail verdict.                                                                |
+
+<Note>
+    The attribute prefix is `speculative_generation.`, and the two `output_rails_*` chunk counters describe the release buffer rather than output-rail activity specifically — they are also emitted for input-rails-only configurations.
+</Note>
 
 **`guardrails.rail`** is one span per rail that runs, wrapping the rail's execution. Span kind `INTERNAL`.
 

docs/reference/engine-feature-support.mdx

@@ -187,8 +187,9 @@ Both engines run multiple rails in the same direction concurrently when `rails.i
 For YAML examples, see [Parallel Execution of Input and Output Rails](/configure-guardrails/yaml-schema/guardrails-configuration#parallel-execution-of-input-and-output-rails).
 
 `IORails` adds two concurrency capabilities that `LLMRails` does not provide.
-Speculative generation (`rails.input.speculative_generation`) runs input rails concurrently with model generation and discards the generation if an input rail blocks, reducing latency on the safe path; it applies to non-streaming generation only.
-For a configuration example, see [Speculative Generation](/configure-guardrails/yaml-schema/guardrails-configuration#speculative-generation).
+Speculative generation (`rails.input.speculative_generation`) runs input rails concurrently with model generation and discards the generation if an input rail blocks, reducing latency on the safe path; it applies to both non-streaming and streaming requests.
+For streaming requests, tokens are held in an in-memory buffer until the input rails pass, and only check-first output-rail behavior is supported.
+For a configuration example, see [Speculative Generation](/configure-guardrails/yaml-schema/guardrails-configuration#speculative-generation), and for the streaming specifics see [Streaming Speculative Generation](/configure-guardrails/yaml-schema/guardrails-configuration#streaming-speculative-generation).
 Admission control through an `AsyncWorkQueue` (and a separate semaphore for streaming) bounds the number of in-flight requests and rejects work when the queue is full.
 
 ### Reasoning-Model Support