FP8 gemma-4-12B-it + DSpark speculative decoding (fast + long-context + streaming/thinking recipe)

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,152 @@
+---
+base_model: google/gemma-4-12B-it
+license: gemma
+library_name: transformers
+pipeline_tag: text-generation
+language:
+- en
+tags:
+- gemma
+- gemma4
+- fp8
+- torchao
+- quantization
+- speculative-decoding
+- dspark
+- long-context
+- blackwell
+---
+
+# gemma-4-12B-it — FP8 + DSpark speculative decoding
+
+A **load-and-go FP8** quantization of [`google/gemma-4-12B-it`](https://huggingface.co/google/gemma-4-12B-it),
+plus a reproducible recipe for running it with **DeepSeek [DSpark](https://github.com/deepseek-ai/DeepSpec)
+speculative decoding** (draft head: [`deepseek-ai/dspark_gemma4_12b_block7`](https://huggingface.co/deepseek-ai/dspark_gemma4_12b_block7))
+— on a single 32 GB Blackwell GPU (validated on an RTX 5090).
+
+Two recipes are documented here, both starting from the same FP8 target:
+
+| recipe | speed | context | notes |
+|---|---|---|---|
+| **Fast** (torch.compile max-autotune) | ~150 tok/s on code (~2× a plain bf16 12B) | ~32 k | short/medium chat + code |
+| **Long-context** (windowed KV cache) | ~40–55 tok/s | **128 k = 26.6 GB, 256 k = 28.7 GB, in-VRAM** | full 256 k on 32 GB |
+
+The weights in this repo are the FP8 target. The DSpark draft, DeepSpec loop, and the two recipes are
+described below (code in [`recipe/`](./recipe)).
+
+## The FP8 quantization (this repo's weights)
+
+torchao dynamic-activation / dynamic-weight float8, **per-row**, forced onto torch's native
+`_scaled_mm` kernel. On Blackwell (sm_120) `_scaled_mm` fp8 matmul is ~2.5× a bf16 matmul; the default
+`KernelPreference.AUTO` instead tries a cutlass kernel that doesn't load on sm_120/py3.12 and silently
+falls back to a slow dequant path — **so `KernelPreference.TORCH` is essential**.
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Load-and-go: the fp8 quantization_config is baked into config.json — no config needed at load.
+model = AutoModelForCausalLM.from_pretrained(
+    "skibare87/gemma-4-12B-it-FP8-DSpark",
+    dtype=torch.bfloat16, device_map="cuda", attn_implementation="sdpa",
+).eval()
+tok = AutoTokenizer.from_pretrained("skibare87/gemma-4-12B-it-FP8-DSpark")
+```
+
+To re-quantize `google/gemma-4-12B-it` yourself instead of using these weights:
+
+```python
+from transformers import AutoModelForCausalLM, TorchAoConfig
+from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, PerRow
+from torchao.quantization.quantize_.common.kernel_preference import KernelPreference
+
+cfg = Float8DynamicActivationFloat8WeightConfig(
+    granularity=PerRow(), kernel_preference=KernelPreference.TORCH,  # native _scaled_mm, not AUTO
+)
+model = AutoModelForCausalLM.from_pretrained(
+    "google/gemma-4-12B-it", quantization_config=TorchAoConfig(cfg),
+    dtype=torch.bfloat16, device_map="cuda", attn_implementation="sdpa",
+).eval()
+# model.save_pretrained("gemma-4-12B-it-FP8")  # <- produces the checkpoint in this repo
+```
+
+FP8 target ≈ 13 GB (vs ~24 GB bf16). Draft head ≈ 7 GB.
+
+## Recipe 1 — Fast (torch.compile max-autotune), ~150 tok/s on code
+
+DSpark verifies the draft's proposals against the FP8 target. Speculation shines where the draft is
+predictable: on **code**, accept-length ≈ 5 and ~150 tok/s (136–171 measured, ~2× a plain 12B); on
+open-ended prose accept-length ≈ 2.5 (~90 tok/s). Speed comes from `torch.compile(mode=
+"max-autotune-no-cudagraphs", dynamic=True)` on the target, which fuses the fp8 activation-quant +
+`_scaled_mm` into proper triton kernels.
+
+- Compile is slow (~30 min cold). Persist it: set `TORCHINDUCTOR_CACHE_DIR` off `/tmp`, and use torch
+  2.11's **mega-cache** (`torch.compiler.save_cache_artifacts()` / `load_cache_artifacts()`) so restarts
+  are a cache hit (~4 min) instead of a recompile. `torch._dynamo.config.caching_precompile` does **not**
+  work with torchao fp8 (it can't serialize `Float8Tensor` guards).
+- (AOTInductor gives ~21 s startup / ~173 tok/s but its static cache caps context at ~32 k on 32 GB —
+  fine for short context, superseded by Recipe 2 for long context.)
+
+## Recipe 2 — Long-context (128 k–256 k in-VRAM on 32 GB)
+
+gemma-4-12B has 40 sliding-attention layers (window 1024) + 8 full-attention layers. A plain
+`DynamicCache` stores the sliding layers **full-length**, so 256 k KV would be ~90 GB. Windowing the
+sliding layers makes 256 k KV **~5 GB**. The pieces (patches in [`recipe/`](./recipe), applied to a
+[DeepSpec](https://github.com/deepseek-ai/DeepSpec) checkout):
+
+1. **`windowed_cache.py` — `SpecSlidingLayer`**: a crop-safe sliding cache (stores `window + pad` so a
+   speculative reject never eats into the real window; `get_mask_sizes` reports the true stored length so
+   gemma's sliding mask stays aligned). **Validated logit-exact vs the full forward past 1024 tokens,
+   through crop cycles.**
+2. **Chunked prefill** (`base_evaluator.patch`): prefill long prompts in chunks and keep only the draft's
+   target hidden-state layers, with a rolling window so `target_hidden_states` never materializes
+   full-length. This was the real memory lever (128 k: 38.7 → 26.6 GB).
+3. **Windowed draft context** (`evaluator.patch`): the draft only proposes, so its context is windowed
+   (`DSPARK_DRAFT_CTX_WINDOW`, default 16384); a cumulative-offset trick keeps absolute positions correct.
+4. **Efficient SDPA backend**: the 8 full-attn layers have `head_dim=512`; **flash-attn caps at 256**, so
+   force `torch.nn.attention.sdpa_kernel([EFFICIENT_ATTENTION, MATH])` — the math backend uses 32 GB for
+   one such attention, efficient uses 4.4 GB.
+
+Measured on the RTX 5090: **128 k = 26.6 GB / 56 s, 256 k = 28.7 GB / 181 s**, both fully in-VRAM.
+
+### Heads-up: a transformers bug you'll hit at long context
+
+`DynamicCache(config=...)` (and `get_head_shapes`) do `layer_types[:-num_kv_shared_layers]`; gemma-4-12B
+has `num_kv_shared_layers = 0`, so `[:-0]` is an **empty list** and the sliding-window cache layers are
+silently never created (everything becomes full-storage). The workaround (build the cache layers
+manually) is in `windowed_cache.py`; details + a minimal repro in
+[`recipe/transformers-num-kv-shared-layers-bug.md`](./recipe/transformers-num-kv-shared-layers-bug.md).
+
+## Serving it
+
+[`recipe/server.py`](./recipe/server.py) is a self-contained OpenAI-compatible `/v1/chat/completions`
+shim wrapping DeepSpec's `Gemma4DSparkEvaluator`, with env knobs for both recipes
+(`DSPARK_COMPILE=1` → fast path; `DSPARK_DRAFT_CTX_WINDOW` / `DSPARK_PREFILL_CHUNK` → long context).
+
+It supports **real token-by-token streaming** (a `stream_callback` added to DSpark's generate loop
+pushes accepted tokens as speculation commits them — see the patches) and **thinking**: gemma4 reasons
+in a `<|channel>thought … <channel|>` channel, which the shim exposes as OpenAI-style `reasoning_content`
+(streamed separately from the answer `content`; toggle with `DSPARK_THINKING=0`). Works through a
+LiteLLM gateway into Open WebUI.
+
+## Gotchas
+
+- **Use the `-it` (instruct) variant.** The base pattern-completes and never stops; `-it` ships its chat
+  template and `eos_token_id: [1, 106, 50]`. gemma-4 uses a `<|channel|>`/`<|think|>` (harmony-style)
+  format, not `<start_of_turn>`.
+- Validated stack: **torch 2.11 + torchao 0.17 + transformers 5.x, CUDA 12.8, RTX 5090 (sm_120), WSL2.**
+- `attn_implementation="sdpa"` (not flash — head_dim 512).
+
+## Attribution & license
+
+- Base model: **[google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it)** — © Google,
+  distributed under the **[Gemma Terms of Use](https://ai.google.dev/gemma/terms)**. This is a derivative
+  (FP8 quantization); the Gemma Terms and use restrictions apply. "Gemma" is a trademark of Google.
+- Draft head + speculative-decoding method: **DeepSeek DSpark** —
+  [`deepseek-ai/dspark_gemma4_12b_block7`](https://huggingface.co/deepseek-ai/dspark_gemma4_12b_block7),
+  [DeepSpec](https://github.com/deepseek-ai/DeepSpec).
+- Quantization: [torchao](https://github.com/pytorch/ao). Serving: [transformers](https://github.com/huggingface/transformers).
+
+Recipe assembled while getting DSpark + gemma-4-12B running at long context on a single 5090; shared so
+others don't have to rediscover the Blackwell fp8 kernel choice, the mega-cache persistence, or the
+sliding-cache/`num_kv_shared_layers` interactions.

chat_template.jinja

@@ -0,0 +1,363 @@
+{%- macro format_parameters(properties, required, filter_keys=false) -%}
+    {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
+    {%- set ns = namespace(found_first=false) -%}
+    {%- for key, value in properties | dictsort -%}
+        {%- set add_comma = false -%}
+        {%- if not filter_keys or key not in standard_keys -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {{ key }}:{
+            {%- if value['description'] -%}
+                description:<|"|>{{ value['description'] }}<|"|>
+                {%- set add_comma = true -%}
+            {%- endif -%}
+            {%- if value['type'] | upper == 'STRING' -%}
+                {%- if value['enum'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    enum:{{ format_argument(value['enum']) }}
+                {%- endif -%}
+            {%- elif value['type'] | upper == 'ARRAY' -%}
+                {%- if value['items'] is mapping and value['items'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    items:{
+                    {%- set ns_items = namespace(found_first=false) -%}
+                    {%- for item_key, item_value in value['items'] | dictsort -%}
+                        {%- if item_value is not none -%}
+                            {%- if ns_items.found_first %},{% endif -%}
+                            {%- set ns_items.found_first = true -%}
+                            {%- if item_key == 'properties' -%}
+                                properties:{
+                                {%- if item_value is mapping -%}
+                                    {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
+                                {%- endif -%}
+                                }
+                            {%- elif item_key == 'required' -%}
+                                required:[
+                                {%- for req_item in item_value -%}
+                                    <|"|>{{- req_item -}}<|"|>
+                                    {%- if not loop.last %},{% endif -%}
+                                {%- endfor -%}
+                                ]
+                            {%- elif item_key == 'type' -%}
+                                {%- if item_value is string -%}
+                                    type:{{ format_argument(item_value | upper) }}
+                                {%- else -%}
+                                    type:{{ format_argument(item_value | map('upper') | list) }}
+                                {%- endif -%}
+                            {%- else -%}
+                                {{ item_key }}:{{ format_argument(item_value) }}
+                            {%- endif -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                    }
+                {%- endif -%}
+            {%- endif -%}
+            {%- if value['nullable'] %}
+                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                nullable:true
+            {%- endif -%}
+            {%- if value['type'] | upper == 'OBJECT' -%}
+                {%- if value['properties'] is defined and value['properties'] is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+                    }
+                {%- elif value is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
+                    }
+                {%- endif -%}
+                {%- if value['required'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    required:[
+                    {%- for item in value['required'] | default([]) -%}
+                        <|"|>{{- item -}}<|"|>
+                        {%- if not loop.last %},{% endif -%}
+                    {%- endfor -%}
+                    ]
+                {%- endif -%}
+            {%- endif -%}
+            {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+            type:<|"|>{{ value['type'] | upper }}<|"|>}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endmacro -%}
+{%- macro format_function_declaration(tool_data) -%}
+    declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
+    {%- set params = tool_data['function']['parameters'] -%}
+    {%- if params -%}
+        ,parameters:{
+        {%- if params['properties'] -%}
+            properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
+        {%- endif -%}
+        {%- if params['required'] -%}
+            required:[
+            {%- for item in params['required'] -%}
+                <|"|>{{- item -}}<|"|>
+                {{- ',' if not loop.last -}}
+            {%- endfor -%}
+            ],
+        {%- endif -%}
+        {%- if params['type'] -%}
+            type:<|"|>{{- params['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if 'response' in tool_data['function'] -%}
+        {%- set response_declaration = tool_data['function']['response'] -%}
+        ,response:{
+        {%- if response_declaration['description'] -%}
+            description:<|"|>{{- response_declaration['description'] -}}<|"|>,
+        {%- endif -%}
+        {%- if response_declaration['type'] | upper == 'OBJECT' -%}
+            type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    }
+{%- endmacro -%}
+{%- macro format_argument(argument, escape_keys=True) -%}
+    {%- if argument is string -%}
+        {{- '<|"|>' + argument + '<|"|>' -}}
+    {%- elif argument is boolean -%}
+        {{- 'true' if argument else 'false' -}}
+    {%- elif argument is mapping -%}
+        {{- '{' -}}
+        {%- set ns = namespace(found_first=false) -%}
+        {%- for key, value in argument | dictsort -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {%- if escape_keys -%}
+                {{- '<|"|>' + key + '<|"|>' -}}
+            {%- else -%}
+                {{- key -}}
+            {%- endif -%}
+            :{{- format_argument(value, escape_keys=escape_keys) -}}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- elif argument is sequence -%}
+        {{- '[' -}}
+        {%- for item in argument -%}
+            {{- format_argument(item, escape_keys=escape_keys) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- ']' -}}
+    {%- else -%}
+        {{- argument -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro strip_thinking(text) -%}
+    {%- set ns = namespace(result='') -%}
+    {%- for part in text.split('<channel|>') -%}
+        {%- if '<|channel>' in part -%}
+            {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
+        {%- else -%}
+            {%- set ns.result = ns.result + part -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- ns.result | trim -}}
+{%- endmacro -%}
+
+{%- macro format_tool_response_block(tool_name, response) -%}
+    {{- '<|tool_response>' -}}
+    {%- if response is mapping -%}
+        {{- 'response:' + tool_name + '{' -}}
+        {%- for key, value in response | dictsort -%}
+            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- else -%}
+        {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
+    {%- endif -%}
+    {{- '<tool_response|>' -}}
+{%- endmacro -%}
+
+{%- set ns = namespace(prev_message_type=None) -%}
+{%- set loop_messages = messages -%}
+{{- bos_token -}}
+{#- Handle System/Tool Definitions Block -#}
+{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}
+    {{- '<|turn>system\n' -}}
+    {#- Inject Thinking token at the very top of the FIRST system turn -#}
+    {%- if enable_thinking is defined and enable_thinking -%}
+        {{- '<|think|>\n' -}}
+        {%- set ns.prev_message_type = 'think' -%}
+    {%- endif -%}
+    {%- if messages[0]['role'] in ['system', 'developer'] -%}
+        {%- if messages[0]['content'] is string -%}
+            {{- messages[0]['content'] | trim -}}
+        {%- elif messages[0]['content'] is sequence -%}
+            {%- for item in messages[0]['content'] -%}
+                {{- item['text'] | trim + ' '-}}
+            {%- endfor -%}
+        {%- endif -%}
+        {%- set loop_messages = messages[1:] -%}
+    {%- endif -%}
+    {%- if tools -%}
+        {%- for tool in tools %}
+            {{- '<|tool>' -}}
+            {{- format_function_declaration(tool) | trim -}}
+            {{- '<tool|>' -}}
+        {%- endfor %}
+        {%- set ns.prev_message_type = 'tool' -%}
+    {%- endif -%}
+    {{- '<turn|>\n' -}}
+{%- endif %}
+
+{#- Pre-scan: find last user message index for reasoning guard -#}
+{%- set ns_turn = namespace(last_user_idx=-1) -%}
+{%- for i in range(loop_messages | length) -%}
+    {%- if loop_messages[i]['role'] == 'user' -%}
+        {%- set ns_turn.last_user_idx = i -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- Loop through messages -#}
+{%- for message in loop_messages -%}
+    {%- if message['role'] != 'tool' -%}
+    {%- set ns.prev_message_type = None -%}
+    {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
+    {#- Detect continuation: suppress duplicate <|turn>model when previous non-tool message was also assistant -#}
+    {%- set prev_nt = namespace(role=None, found=false) -%}
+    {%- if loop.index0 > 0 -%}
+        {%- for j in range(loop.index0 - 1, -1, -1) -%}
+            {%- if not prev_nt.found -%}
+                {%- if loop_messages[j]['role'] != 'tool' -%}
+                    {%- set prev_nt.role = loop_messages[j]['role'] -%}
+                    {%- set prev_nt.found = true -%}
+                {%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+    {%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%}
+    {%- if not continue_same_model_turn -%}
+        {{- '<|turn>' + role + '\n' }}
+    {%- endif -%}
+
+    {#- Render reasoning/reasoning_content as thinking channel -#}
+    {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
+    {%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}
+        {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
+    {%- endif -%}
+
+            {%- if message['tool_calls'] -%}
+                {%- for tool_call in message['tool_calls'] -%}
+                    {%- set function = tool_call['function'] -%}
+                    {{- '<|tool_call>call:' + function['name'] + '{' -}}
+                    {%- if function['arguments'] is mapping -%}
+                        {%- set ns_args = namespace(found_first=false) -%}
+                        {%- for key, value in function['arguments'] | dictsort -%}
+                            {%- if ns_args.found_first %},{% endif -%}
+                            {%- set ns_args.found_first = true -%}
+                            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+                        {%- endfor -%}
+                    {%- elif function['arguments'] is string -%}
+                        {{- function['arguments'] -}}
+                    {%- endif -%}
+                    {{- '}<tool_call|>' -}}
+                {%- endfor -%}
+                {%- set ns.prev_message_type = 'tool_call' -%}
+            {%- endif -%}
+
+            {%- set ns_tr_out = namespace(flag=false) -%}
+            {%- if message.get('tool_responses') -%}
+                {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
+                {%- for tool_response in message['tool_responses'] -%}
+                    {{- format_tool_response_block(tool_response['name'] | default('unknown'), tool_response['response']) -}}
+                    {%- set ns_tr_out.flag = true -%}
+                    {%- set ns.prev_message_type = 'tool_response' -%}
+                {%- endfor -%}
+            {%- elif message.get('tool_calls') -%}
+                {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
+                {%- set ns_tool_scan = namespace(stopped=false) -%}
+                {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
+                    {%- if ns_tool_scan.stopped -%}
+                    {%- elif loop_messages[k]['role'] != 'tool' -%}
+                        {%- set ns_tool_scan.stopped = true -%}
+                    {%- else -%}
+                        {%- set follow = loop_messages[k] -%}
+                        {#- Resolve tool_call_id to function name -#}
+                        {%- set ns_tname = namespace(name=follow.get('name') | default('unknown')) -%}
+                        {%- for tc in message['tool_calls'] -%}
+                            {%- if tc.get('id') == follow.get('tool_call_id') -%}
+                                {%- set ns_tname.name = tc['function']['name'] -%}
+                            {%- endif -%}
+                        {%- endfor -%}
+                        {#- Handle content as string or content-parts array -#}
+                        {%- set tool_body = follow.get('content') -%}
+                        {%- if tool_body is string -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- elif tool_body is sequence and tool_body is not string -%}
+                            {%- set ns_txt = namespace(s='') -%}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'text' -%}
+                                    {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
+                                {%- endif -%}
+                            {%- endfor -%}
+                            {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'image' -%}
+                                    {{- '<|image|>' -}}
+                                {%- elif part.get('type') == 'audio' -%}
+                                    {{- '<|audio|>' -}}
+                                {%- elif part.get('type') == 'video' -%}
+                                    {{- '<|video|>' -}}
+                                {%- endif -%}
+                            {%- endfor -%}
+                        {%- else -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- endif -%}
+                        {%- set ns_tr_out.flag = true -%}
+                        {%- set ns.prev_message_type = 'tool_response' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+
+            {%- set captured_content -%}
+            {%- if message['content'] is string -%}
+                {%- if role == 'model' -%}
+                    {{- strip_thinking(message['content']) -}}
+                {%- else -%}
+                    {{- message['content'] | trim -}}
+                {%- endif -%}
+            {%- elif message['content'] is sequence -%}
+                {%- for item in message['content'] -%}
+                    {%- if item['type'] == 'text' -%}
+                        {%- if role == 'model' -%}
+                            {{- strip_thinking(item['text']) -}}
+                        {%- else -%}
+                            {{- item['text'] | trim -}}
+                        {%- endif -%}
+                    {%- elif item['type'] == 'image' -%}
+                        {{- '<|image|>' -}}
+                        {%- set ns.prev_message_type = 'image' -%}
+                    {%- elif item['type'] == 'audio' -%}
+                        {{- '<|audio|>' -}}
+                        {%- set ns.prev_message_type = 'audio' -%}
+                    {%- elif item['type'] == 'video' -%}
+                        {{- '<|video|>' -}}
+                        {%- set ns.prev_message_type = 'video' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+            {%- endset -%}
+
+            {{- captured_content -}}
+            {%- set has_content = captured_content | trim | length > 0 -%}
+
+        {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
+            {{- '<|tool_response>' -}}
+        {%- elif not (ns_tr_out.flag and not has_content) -%}
+            {{- '<turn|>\n' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
+        {{- '<|turn>model\n' -}}
+        {%- if not enable_thinking | default(false) -%}
+            {{- '<|channel>thought\n<channel|>' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}

config.json

@@ -0,0 +1,227 @@
+{
+  "architectures": [
+    "Gemma4UnifiedForConditionalGeneration"
+  ],
+  "audio_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "audio_embed_dim": 640,
+    "chunk_size_feed_forward": 0,
+    "dtype": "bfloat16",
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "model_type": "gemma4_unified_audio",
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "problem_type": null,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06
+  },
+  "audio_token_id": 258881,
+  "boa_token_id": 256000,
+  "boi_token_id": 255999,
+  "dtype": "bfloat16",
+  "eoa_token_index": 258883,
+  "eoi_token_id": 258882,
+  "eos_token_id": [
+    1,
+    106
+  ],
+  "image_token_id": 258880,
+  "initializer_range": 0.02,
+  "model_type": "gemma4_unified",
+  "quantization_config": {
+    "include_input_output_embeddings": false,
+    "modules_to_not_convert": null,
+    "quant_method": "torchao",
+    "quant_type": {
+      "default": {
+        "_data": {
+          "activation_dtype": {
+            "_data": "float8_e4m3fn",
+            "_type": "torch.dtype"
+          },
+          "activation_value_lb": null,
+          "activation_value_ub": null,
+          "granularity": [
+            {
+              "_data": {
+                "dim": -1
+              },
+              "_type": "PerRow",
+              "_version": 1
+            },
+            {
+              "_data": {
+                "dim": -1
+              },
+              "_type": "PerRow",
+              "_version": 1
+            }
+          ],
+          "kernel_preference": {
+            "_data": "TORCH",
+            "_type": "KernelPreference"
+          },
+          "mm_config": {
+            "_data": {
+              "emulate": false,
+              "pad_inner_dim": false,
+              "use_fast_accum": true
+            },
+            "_type": "Float8MMConfig",
+            "_version": 1
+          },
+          "packing_format": {
+            "_data": "PLAIN",
+            "_type": "Float8PackingFormat"
+          },
+          "set_inductor_config": true,
+          "weight_dtype": {
+            "_data": "float8_e4m3fn",
+            "_type": "torch.dtype"
+          }
+        },
+        "_type": "Float8DynamicActivationFloat8WeightConfig",
+        "_version": 2
+      }
+    },
+    "untie_embedding_weights": false
+  },
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attention_k_eq_v": true,
+    "bos_token_id": 2,
+    "dtype": "bfloat16",
+    "enable_moe_block": false,
+    "eos_token_id": 1,
+    "final_logit_softcapping": 30.0,
+    "global_head_dim": 512,
+    "head_dim": 256,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 3840,
+    "hidden_size_per_layer_input": 0,
+    "initializer_range": 0.02,
+    "intermediate_size": 15360,
+    "layer_types": [
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 262144,
+    "model_type": "gemma4_unified_text",
+    "moe_intermediate_size": null,
+    "num_attention_heads": 16,
+    "num_experts": null,
+    "num_global_key_value_heads": 1,
+    "num_hidden_layers": 48,
+    "num_key_value_heads": 8,
+    "num_kv_shared_layers": 0,
+    "pad_token_id": 0,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "full_attention": {
+        "partial_rotary_factor": 0.25,
+        "rope_theta": 1000000.0,
+        "rope_type": "proportional"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 1024,
+    "tie_word_embeddings": true,
+    "top_k_experts": null,
+    "use_bidirectional_attention": "vision",
+    "use_cache": true,
+    "use_double_wide_mlp": false,
+    "vocab_size": 262144,
+    "vocab_size_per_layer_input": 262144
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.10.2",
+  "video_token_id": 258884,
+  "vision_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "chunk_size_feed_forward": 0,
+    "dtype": "bfloat16",
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "mm_embed_dim": 3840,
+    "mm_posemb_size": 1120,
+    "model_type": "gemma4_unified_vision",
+    "num_soft_tokens": 280,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_proj_dims": 3840,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "problem_type": null,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06
+  }
+}

generation_config.json

@@ -0,0 +1,18 @@
+{
+  "bos_token_id": 2,
+  "do_sample": true,
+  "eos_token_id": [
+    1,
+    106,
+    50
+  ],
+  "pad_token_id": 0,
+  "suppress_tokens": [
+    258883,
+    258882
+  ],
+  "temperature": 1.0,
+  "top_k": 64,
+  "top_p": 0.95,
+  "transformers_version": "5.10.2"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c42e8ea59a29c065d56925ad817967af29816e3871f3c919bf84a1faaec5f23d
+size 12985253984

recipe/base_evaluator.patch

@@ -0,0 +1,86 @@
+--- /tmp/base_eval.bak	2026-07-03 13:59:36.997249952 -0400
++++ deepspec/eval/base_evaluator.py	2026-07-03 18:43:15.290004401 -0400
+@@ -317,6 +317,8 @@
+     propose: Callable[..., DraftProposal],
+     update: Callable[[Any, VerificationResult], None],
+     post_verify: Callable[[DraftProposal, VerificationResult], None] | None = None,
++    prefill_keep_hidden_layers: list[int] | None = None,
++    stream_callback=None,
+ ) -> SimpleNamespace:
+     """Speculative-decoding loop.
+ 
+@@ -343,14 +345,48 @@
+     from deepspec.eval.windowed_cache import build_target_cache
+     past_key_values_target = build_target_cache(target_model, pad=max(64, int(max_proposal_tokens)+8))
+ 
+-    output = target_model(
+-        input_ids=input_ids,
+-        position_ids=position_ids[:, :num_input_tokens],
+-        past_key_values=past_key_values_target,
+-        use_cache=True,
+-        output_hidden_states=True,
+-        logits_to_keep=1,
+-    )
++    _chunk = int(os.environ.get("DSPARK_PREFILL_CHUNK", "4096"))
++    if num_input_tokens > _chunk and prefill_keep_hidden_layers is not None:
++        # chunked prefill: bound activation memory (never materialize all layers x all positions);
++        # keep only the draft target layers hidden states, concatenated across chunks.
++        _keep = sorted(set(int(l) for l in prefill_keep_hidden_layers))
++        _acc = {li: [] for li in _keep}
++        _last_logits = None
++        for _i in range(0, num_input_tokens, _chunk):
++            _j = min(_i + _chunk, num_input_tokens)
++            _islast = _j == num_input_tokens
++            _o = target_model(
++                input_ids=input_ids[:, _i:_j],
++                position_ids=position_ids[:, _i:_j],
++                past_key_values=past_key_values_target,
++                use_cache=True,
++                output_hidden_states=True,
++                logits_to_keep=1 if _islast else 0,
++            )
++            for li in _keep:
++                _acc[li].append(_o.hidden_states[li])
++            if _islast:
++                _last_logits = _o.logits
++            del _o
++            _hw = int(os.environ.get("DSPARK_DRAFT_CTX_WINDOW", "0"))
++            if _hw:
++                for li in _keep:
++                    _tot = sum(t.shape[1] for t in _acc[li])
++                    while len(_acc[li]) > 1 and _tot - _acc[li][0].shape[1] >= _hw:
++                        _tot -= _acc[li].pop(0).shape[1]
++        _hs = [None] * (max(_keep) + 1)
++        for li in _keep:
++            _hs[li] = torch.cat(_acc[li], dim=1)
++        output = SimpleNamespace(logits=_last_logits, hidden_states=tuple(_hs))
++    else:
++        output = target_model(
++            input_ids=input_ids,
++            position_ids=position_ids[:, :num_input_tokens],
++            past_key_values=past_key_values_target,
++            use_cache=True,
++            output_hidden_states=True,
++            logits_to_keep=1,
++        )
+ 
+     output_ids[:, :num_input_tokens] = input_ids
+     output_ids[:, num_input_tokens : num_input_tokens + 1] = sample_from_probs(
+@@ -384,6 +420,8 @@
+     )
+ 
+     while start < max_length:
++        if stream_callback is not None:
++            stream_callback(output_ids[:, num_input_tokens : start + 1])
+         proposal = propose(
+             context=context,
+             output_ids=output_ids,
+@@ -429,6 +467,8 @@
+         if has_stop_token(new_token_ids, stop_token_ids):
+             break
+ 
++    if stream_callback is not None:
++        stream_callback(output_ids[:, num_input_tokens : start + 1])
+     output_ids = output_ids[:, : min(start + 1, max_length)]
+     output_ids = trim_output_ids(output_ids, num_input_tokens, stop_token_ids)
+     return SimpleNamespace(

recipe/evaluator.patch

@@ -0,0 +1,59 @@
+diff --git a/deepspec/eval/dspark/evaluator.py b/deepspec/eval/dspark/evaluator.py
+index eba2b34..bc4abda 100644
+--- a/deepspec/eval/dspark/evaluator.py
++++ b/deepspec/eval/dspark/evaluator.py
+@@ -88,12 +88,32 @@ class Qwen3DSparkEvaluator(BaseEvaluator):
+         initial_output,
+         **kwargs,
+     ) -> SimpleNamespace:
++        import os
++        _w = int(os.environ.get("DSPARK_DRAFT_CTX_WINDOW", "0"))
++        target_hidden_states = extract_context_feature(
++            initial_output.hidden_states,
++            self.draft_model.target_layer_ids,
++        )
++        if _w:
++            from deepspec.eval.windowed_cache import make_draft_cache
++            past_key_values_draft = make_draft_cache(
++                self.draft_model, _w, pad=int(self.max_proposal_tokens) + 64
++            )
++            # window the init context tensor (its 5-layer hidden states are the big init spike) and
++            # offset the draft cache cumulative so absolute positions stay correct downstream.
++            _num_input = int(kwargs.get("num_input_tokens", target_hidden_states.shape[1]))
++            if target_hidden_states.shape[1] > _w:
++                target_hidden_states = target_hidden_states[:, -_w:, :].contiguous()
++            # absolute offset = where the retained window starts (prefill may have pre-trimmed it)
++            _offset = _num_input - target_hidden_states.shape[1]
++            if _offset > 0:
++                for _layer in past_key_values_draft.layers:
++                    _layer.cumulative_length = _offset
++        else:
++            past_key_values_draft = DynamicCache()
+         return SimpleNamespace(
+-            past_key_values_draft=DynamicCache(),
+-            target_hidden_states=extract_context_feature(
+-                initial_output.hidden_states,
+-                self.draft_model.target_layer_ids,
+-            ),
++            past_key_values_draft=past_key_values_draft,
++            target_hidden_states=target_hidden_states,
+         )
+ 
+     def _propose(
+@@ -164,6 +184,7 @@ class Qwen3DSparkEvaluator(BaseEvaluator):
+         *,
+         input_ids: torch.Tensor,
+         stop_token_ids: list[int] | None,
++        stream_callback=None,
+     ) -> SimpleNamespace:
+         return generate_decoding_sample(
+             target_model=self.target_model,
+@@ -176,6 +197,8 @@ class Qwen3DSparkEvaluator(BaseEvaluator):
+             propose=self._propose,
+             update=self._update,
+             post_verify=self._post_verify,
++            prefill_keep_hidden_layers=[0 if l == -1 else l + 1 for l in self.draft_model.target_layer_ids],
++            stream_callback=stream_callback,
+         )
+ 
+     def evaluate(self) -> None:

recipe/server.py

@@ -0,0 +1,288 @@
+"""OpenAI-compatible shim for DSpark + Gemma4-12B on hiro (RTX 5090).
+
+FP8 target (torchao float8 dynamic activation/weight) + BF16 draft head. Wraps DeepSeek DeepSpec's
+DSpark speculative generate loop (Gemma4DSparkEvaluator.generate_one_sample) behind
+/v1/chat/completions. Runs from ~/DeepSpec (imports `deepspec`) in its .venv.
+
+- FP8 target ~13 GB + BF16 draft ~7 GB, plus a windowed KV cache -> 128k fits in 26.6 GB and 256k in
+  28.7 GB, both fully in-VRAM on the 32 GB card (no host spilling).
+- Long context needs DeepSpec's patched base_evaluator (chunked prefill + rolling hidden-state window)
+  and windowed_cache.py (SpecSlidingLayer): sliding-attn layers window to 1024, so 256k KV is ~5 GB
+  instead of ~90 GB. Validated: windowed cache is logit-exact vs the full forward past 1024 tokens.
+- bsz=1 in the DSpark loop -> requests are serialized under a lock (fine for a single-tenant gemma).
+- Subclasses the evaluator to load the FP8 target; the long-context patches live in the DeepSpec repo.
+"""
+import os
+import time
+import threading
+import queue
+import uuid
+
+os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+os.environ.setdefault("MASTER_PORT", "29500")
+os.environ.setdefault("RANK", "0")
+os.environ.setdefault("WORLD_SIZE", "1")
+# persist inductor's compiled/autotuned kernels off /tmp (which is wiped on reboot) so restarts
+# reuse them instead of re-running the ~30-min max-autotune compile. fx_graph_cache is on by default.
+os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", os.path.expanduser("~/DeepSpec/.inductor_cache"))
+# long-context (128k-256k) defaults: window the draft's context + chunk the prefill so the whole run
+# fits in the 32 GB card (128k=26.6GB, 256k=28.7GB, both in-VRAM). Implemented in DeepSpec's patched
+# base_evaluator (chunked prefill + rolling hidden-state window) and windowed_cache (SpecSlidingLayer).
+os.environ.setdefault("DSPARK_DRAFT_CTX_WINDOW", "16384")
+os.environ.setdefault("DSPARK_PREFILL_CHUNK", "2048")
+
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+# NOTE: torch._dynamo.config.caching_precompile (frontend guard cache) is INCOMPATIBLE with the
+# torchao fp8 target — serializing guards calls empty_like on a Float8Tensor, which torchao doesn't
+# implement (NotImplementedError). It crashes the cache save. So we can't cache the dynamo frontend;
+# the ~200s warmup tracing stays. (AOTInductor would avoid dynamo entirely — separate effort.)
+from types import SimpleNamespace
+from typing import List, Optional
+import json as _json
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
+from torchao.quantization import (
+    Float8DynamicActivationFloat8WeightConfig, Float8WeightOnlyConfig, PerRow)
+from torchao.quantization.quantize_.common.kernel_preference import KernelPreference
+from deepspec.eval.base_evaluator import assert_no_final_target_layer, resolve_stop_token_ids
+from deepspec.eval.dspark import Gemma4DSparkEvaluator
+
+TARGET = os.environ.get("DSPARK_TARGET", "google/gemma-4-12B-it")
+DRAFT = os.environ.get("DSPARK_DRAFT", "deepseek-ai/dspark_gemma4_12b_block7")
+MODEL_NAME = os.environ.get("DSPARK_MODEL_NAME", "gemma4-dspark")
+# gemma4 reasons in a `<|channel>thought\n ... <channel|>` channel; enable_thinking leaves it open so
+# the model reasons, then emits the answer after the <channel|> token (id 101). We expose the reasoning
+# as `reasoning_content` and the answer as `content`. Toggle with DSPARK_THINKING=0.
+THINKING = os.environ.get("DSPARK_THINKING", "1") == "1"
+_CH_CLOSE = 101  # <channel|> — separates the thought channel from the answer
+
+
+class FP8Gemma4DSparkEvaluator(Gemma4DSparkEvaluator):
+    """Same as the stock evaluator but loads the TARGET in torchao FP8 (draft stays BF16)."""
+    def build_models(self):
+        fp8 = os.environ.get("DSPARK_TARGET_DTYPE", "fp8").lower() == "fp8"
+        if fp8:
+            recipe = os.environ.get("DSPARK_FP8_RECIPE", "dynamic").lower()
+            if recipe == "weightonly":
+                # fp8 WEIGHTS (13 GB, memory-bandwidth win), bf16 activations -> no per-forward
+                # activation-quant overhead (which slowed dynamic fp8 below bf16 in DSpark's small forwards)
+                fp8_cfg = Float8WeightOnlyConfig()
+            else:
+                # dynamic activation fp8, native torch._scaled_mm path (AUTO tries a broken cutlass kernel here)
+                fp8_cfg = Float8DynamicActivationFloat8WeightConfig(
+                    granularity=PerRow(), kernel_preference=KernelPreference.TORCH)
+            target = AutoModelForCausalLM.from_pretrained(
+                self.args.target_name_or_path,
+                quantization_config=TorchAoConfig(fp8_cfg),
+                dtype=torch.bfloat16,
+                device_map={"": self.device},
+                attn_implementation=self.EVAL_ATTN_IMPLEMENTATION,
+            ).eval()
+        else:
+            target = AutoModelForCausalLM.from_pretrained(
+                self.args.target_name_or_path,
+                dtype=torch.bfloat16,
+                attn_implementation=self.EVAL_ATTN_IMPLEMENTATION,
+            ).to(self.device).eval()
+        if os.environ.get("DSPARK_COMPILE", "0") == "1":
+            # let inductor fuse the fp8 activation-quant + _scaled_mm into one kernel (no cudagraphs:
+            # the DSpark loop has dynamic control flow); dynamic shapes for the varying verify/prefill M
+            target = torch.compile(target, mode="max-autotune-no-cudagraphs", dynamic=True)
+        draft = self.draft_model_cls.from_pretrained(
+            self.args.draft_name_or_path,
+            dtype=torch.bfloat16,
+            attn_implementation=self.EVAL_ATTN_IMPLEMENTATION,
+        ).to(self.device).eval()
+        assert_no_final_target_layer(target, draft.target_layer_ids)
+        tokenizer = AutoTokenizer.from_pretrained(self.args.target_name_or_path)
+        return target, draft, tokenizer
+
+
+_args = SimpleNamespace(
+    target_name_or_path=TARGET, draft_name_or_path=DRAFT,
+    max_new_tokens=512, temperature=1.0, confidence_threshold=0.0,
+    tensorboard_dir=None, step=None, seed=980406, tasks=[],
+)
+print(f"[dspark] loading target(FP8)={TARGET} + draft(BF16)={DRAFT} ...", flush=True)
+_T0 = time.time()
+EV = FP8Gemma4DSparkEvaluator(0, _args)
+print(f"[dspark] TIMING model+quant load: {time.time()-_T0:.1f}s", flush=True)
+EV.confidence_head_recorder = None  # metrics recorder is only started inside evaluate(); we bypass it
+TOK = EV.tokenizer
+STOP = resolve_stop_token_ids(EV.target_model, TOK)  # gemma real eos set (e.g. [1,106,50]), not a guess
+_LOCK = threading.Lock()
+
+if os.environ.get("DSPARK_COMPILE", "0") == "1":
+    # torch 2.11 mega-cache: the portable "saved compiled version". Load it before compiling so the
+    # warmup is a near-instant cache hit (full graph codegen + autotune), not a ~30-min recompile.
+    _MEGA = os.path.expanduser("~/DeepSpec/.dspark_megacache.bin")
+    if os.path.exists(_MEGA):
+        try:
+            _t = time.time()
+            with open(_MEGA, "rb") as _f:
+                torch.compiler.load_cache_artifacts(_f.read())
+            print(f"[dspark] TIMING mega-cache load: {time.time()-_t:.1f}s", flush=True)
+        except Exception as _e:
+            print("[dspark] mega-cache load failed:", _e, flush=True)
+    print("[dspark] warming up torch.compile ...", flush=True)
+    for _i, _wp in enumerate(["Hello there.", "def fib(n):", "Explain gravity briefly."]):
+        _enc = TOK.apply_chat_template([{"role": "user", "content": _wp}],
+                                       add_generation_prompt=True, return_tensors="pt")
+        _ids = (_enc["input_ids"] if hasattr(_enc, "keys") else _enc).to(EV.device)
+        EV.args.max_new_tokens = 48
+        _t = time.time()
+        with torch.no_grad():
+            EV.generate_one_sample(input_ids=_ids, stop_token_ids=STOP)
+        print(f"[dspark] TIMING warmup[{_i}]: {time.time()-_t:.1f}s", flush=True)
+    # persist the full compiled graph so future starts skip the compile entirely
+    try:
+        _art = torch.compiler.save_cache_artifacts()
+        _blob = _art[0] if isinstance(_art, tuple) else _art
+        if _blob:
+            with open(_MEGA, "wb") as _f:
+                _f.write(_blob)
+            print(f"[dspark] saved compile mega-cache ({len(_blob)//1024} KB)", flush=True)
+    except Exception as _e:
+        print("[dspark] mega-cache save failed:", _e, flush=True)
+    print("[dspark] warmup done", flush=True)
+
+print(f"[dspark] ready. VRAM={torch.cuda.memory_allocated()/1e9:.1f}GB stop={STOP}", flush=True)
+
+app = FastAPI()
+
+
+class Msg(BaseModel):
+    role: str
+    content: str
+
+
+class ChatReq(BaseModel):
+    model: Optional[str] = None
+    messages: List[Msg]
+    max_tokens: Optional[int] = 512
+    temperature: Optional[float] = 0.7
+    stream: Optional[bool] = False
+
+
+@app.get("/health")
+def health():
+    return {"status": "ok", "model": MODEL_NAME, "target": TARGET, "draft": DRAFT}
+
+
+@app.get("/v1/models")
+def models():
+    return {"object": "list", "data": [{"id": MODEL_NAME, "object": "model", "owned_by": "deepspec-dspark"}]}
+
+
+def _split_think(ids):
+    """Split generated token ids into (reasoning, content). With thinking on, the model emits the
+    thought channel then the <channel|> token (101) then the answer; before 101 appears it's all
+    reasoning. The leading `thought` label (regular token after the special <|channel>) is stripped."""
+    if not THINKING:
+        return "", TOK.decode(ids, skip_special_tokens=True)
+    if _CH_CLOSE in ids:
+        i = ids.index(_CH_CLOSE)
+        r = TOK.decode(ids[:i], skip_special_tokens=True)
+        c = TOK.decode(ids[i + 1:], skip_special_tokens=True)
+    else:
+        r, c = TOK.decode(ids, skip_special_tokens=True), ""
+    if r.startswith("thought"):
+        r = r[len("thought"):]
+    return r.strip("\n"), c
+
+
+@app.post("/v1/chat/completions")
+def chat(req: ChatReq):
+    msgs = [{"role": m.role, "content": m.content} for m in req.messages]
+    ids = TOK.apply_chat_template(
+        msgs, add_generation_prompt=True, enable_thinking=THINKING,
+        return_tensors="pt", return_dict=True,
+    )["input_ids"].to(EV.device)
+    cid = "chatcmpl-" + uuid.uuid4().hex[:12]
+    created = int(time.time())
+    model = req.model or MODEL_NAME
+    max_new = int(req.max_tokens or 512)
+    if THINKING:
+        # gemma4 reasons deeply; keep the client's max_tokens as the ANSWER budget and add a SEPARATE,
+        # GENEROUS reasoning ceiling on top, so nothing gets truncated mid-thought. This is a ceiling, not
+        # a forced length — simple prompts stop early on their own, so being generous costs them nothing;
+        # only genuinely hard prompts spend it. (We have the context for it — don't cap thinking small.)
+        max_new += int(os.environ.get("DSPARK_THINK_BUDGET", "16384"))
+    temp = max(float(req.temperature if req.temperature is not None else 1.0), 0.05)
+
+    if req.stream:
+        # Real streaming: run the (blocking) DSpark loop in a thread; its stream_callback pushes the
+        # cumulative generated ids onto a queue as speculative decoding accepts them; the SSE generator
+        # decodes, splits thought/answer, and emits reasoning_content + content deltas.
+        q: "queue.Queue" = queue.Queue()
+
+        def _run():
+            try:
+                with _LOCK:
+                    EV.args.max_new_tokens = max_new
+                    EV.args.temperature = temp
+                    with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]), torch.no_grad():
+                        EV.generate_one_sample(
+                            input_ids=ids, stop_token_ids=STOP,
+                            stream_callback=lambda t: q.put(t[0].tolist()),
+                        )
+            except Exception as e:  # surface generation errors to the stream instead of hanging
+                q.put(("__error__", str(e)))
+            finally:
+                q.put(None)
+
+        threading.Thread(target=_run, daemon=True).start()
+
+        def _chunk(delta, finish=None):
+            return "data: " + _json.dumps({
+                "id": cid, "object": "chat.completion.chunk", "created": created, "model": model,
+                "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
+            }) + "\n\n"
+
+        def _sse():
+            yield _chunk({"role": "assistant"})
+            last_r = last_c = ""
+            while True:
+                item = q.get()
+                if item is None:
+                    break
+                if isinstance(item, tuple) and item[0] == "__error__":
+                    yield _chunk({"content": f"\n[error: {item[1]}]"})
+                    break
+                r, c = _split_think(item)
+                if len(r) > len(last_r):
+                    yield _chunk({"reasoning_content": r[len(last_r):]})
+                    last_r = r
+                if len(c) > len(last_c):
+                    yield _chunk({"content": c[len(last_c):]})
+                    last_c = c
+            yield _chunk({}, finish="stop")
+            yield "data: [DONE]\n\n"
+
+        return StreamingResponse(_sse(), media_type="text/event-stream")
+
+    with _LOCK:
+        EV.args.max_new_tokens = max_new
+        EV.args.temperature = temp
+        t = time.time()
+        with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]), torch.no_grad():
+            res = EV.generate_one_sample(input_ids=ids, stop_token_ids=STOP)
+        dt = time.time() - t
+    gen = res.output_ids[0, res.num_input_tokens:].tolist()
+    reasoning, content = _split_think(gen)
+    al = res.acceptance_lengths
+    msg = {"role": "assistant", "content": content}
+    if reasoning:
+        msg["reasoning_content"] = reasoning
+    return {
+        "id": cid, "object": "chat.completion", "created": created, "model": model,
+        "choices": [{"index": 0, "message": msg, "finish_reason": "stop"}],
+        "usage": {"prompt_tokens": res.num_input_tokens, "completion_tokens": res.num_output_tokens,
+                  "total_tokens": res.num_input_tokens + res.num_output_tokens},
+        "dspark": {"mean_accept_len": round(sum(al) / len(al), 2) if al else 0,
+                   "verify_passes": len(al), "gen_seconds": round(dt, 2),
+                   "tokens_per_sec": round(res.num_output_tokens / dt, 1) if dt > 0 else 0},
+    }

recipe/transformers-num-kv-shared-layers-bug.md

@@ -0,0 +1,77 @@
+# Bug report draft — `num_kv_shared_layers == 0` silently disables hybrid cache layer typing
+
+**Repo:** huggingface/transformers · **Type:** bug (silent correctness/perf regression)
+
+## Title
+
+`DynamicCache(config=...)` (and `get_head_shapes`) drop all typed layers when `num_kv_shared_layers == 0` — sliding-window layers stop windowing
+
+## Summary
+
+When a decoder config exposes `num_kv_shared_layers` **set to 0** (e.g. `google/gemma-4-12B-it`),
+`DynamicCache.__init__` computes `layer_types[: -num_kv_shared_layers]`, which for `n == 0` is
+`layer_types[:-0]` → **an empty list**. The per-layer type loop then runs zero times, so no
+`DynamicSlidingWindowLayer`s are created; every layer lazily falls back to a plain full-storage
+`DynamicLayer`. The **sliding-attention layers silently stop windowing**, so KV memory grows with the
+full sequence length instead of being capped at the window.
+
+For gemma-4-12B (40 sliding layers, window 1024, at 256k context) this is the difference between
+**~5 GB and ~90 GB** of KV cache — i.e. it silently makes long-context generation OOM on hardware
+that should handle it.
+
+The identical slicing pattern appears in `transformers/integrations/executorch.py::get_head_shapes`,
+where it yields empty `num_heads`/`head_dim` lists and raises
+`ValueError: num_head was provided as a list of length 0, but the Cache currently has N layers`.
+
+## Offending code
+
+`src/transformers/cache_utils.py`, `DynamicCache.__init__`:
+
+```python
+# Some models have shared layers thus no cache is needed for them (e.g. Gemma3n)
+if hasattr(decoder_config, "num_kv_shared_layers"):
+    layer_types = layer_types[: -decoder_config.num_kv_shared_layers]   # n == 0 -> [] (empty!)
+```
+
+`src/transformers/integrations/executorch.py`, `get_head_shapes`:
+
+```python
+head_dim = [... for layer in config.layer_types[: -config.num_kv_shared_layers]]   # same bug
+num_heads = [... for layer in config.layer_types[: -config.num_kv_shared_layers]]
+```
+
+## Minimal repro
+
+```python
+from transformers import AutoConfig
+from transformers.cache_utils import DynamicCache, DynamicSlidingWindowLayer
+
+cfg = AutoConfig.from_pretrained("google/gemma-4-12B-it").get_text_config()
+print("num_kv_shared_layers:", cfg.num_kv_shared_layers)          # -> 0
+print("layer_types:", len(cfg.layer_types), cfg.layer_types[:2])  # -> 48 ['sliding_attention', ...]
+
+cache = DynamicCache(config=cfg)
+print("typed layers built from config:", len(cache.layers))       # -> 0  (BUG; expected 48)
+# after a forward, every layer lazily becomes a plain DynamicLayer, so sliding layers never window:
+# none are DynamicSlidingWindowLayer, and KV for sliding layers grows unbounded with seq len.
+```
+
+Expected: 48 typed layers, with `sliding_attention` entries → `DynamicSlidingWindowLayer` (windowed).
+Actual: 0 typed layers built from config → all lazy `DynamicLayer` (full storage, no windowing).
+
+## Suggested fix
+
+Guard the slice on a positive value (a config can legitimately carry `num_kv_shared_layers = 0`):
+
+```python
+if getattr(decoder_config, "num_kv_shared_layers", 0):
+    layer_types = layer_types[: -decoder_config.num_kv_shared_layers]
+```
+
+Apply the same guard in `get_head_shapes` (and anywhere else the `[:-num_kv_shared_layers]` idiom is
+used). A quick `grep -rn "num_kv_shared_layers\]" src/transformers` finds the sites.
+
+## Environment
+
+Present on `main` (verified) and transformers 5.x. Triggered by any hybrid-attention model whose config
+sets `num_kv_shared_layers = 0` while defining `layer_types` (gemma-4-12B is one released example).

recipe/windowed_cache.py

@@ -0,0 +1,62 @@
+"""Windowed KV cache for hybrid-attention models (gemma4) under DSpark speculative decoding.
+
+Sliding-attention layers store only `sliding_window + pad` tokens (so 256k context fits in ~5 GB
+instead of ~90 GB) while full-attention layers store everything. The `pad` keeps a small speculative
+`crop` (rejecting proposal tokens) from eating into the real window. `get_mask_sizes` reports the
+true stored length so the model's sliding mask (which still windows to config.sliding_window) aligns.
+Validated: 100% argmax match vs the full forward past 1024 tokens, including crop cycles.
+"""
+import transformers.cache_utils as cu
+from transformers.cache_utils import DynamicCache, DynamicLayer
+
+
+class SpecSlidingLayer(cu.DynamicSlidingWindowLayer):
+    def __init__(self, real_window, pad):
+        super().__init__(sliding_window=real_window + pad)
+
+    def get_mask_sizes(self, query_length):
+        stored = self.keys.shape[-2] if (self.is_initialized and self.keys is not None) else 0
+        return stored + query_length, max(self.cumulative_length - stored, 0)
+
+    def crop(self, max_length):
+        if max_length < 0:
+            max_length = self.cumulative_length + max_length
+        remove = self.cumulative_length - max_length
+        if remove <= 0:
+            return
+        n = self.keys.shape[-2]
+        self.keys = self.keys[:, :, : n - remove, :]
+        self.values = self.values[:, :, : n - remove, :]
+        self.cumulative_length = max_length
+
+
+def build_target_cache(model, pad=64):
+    try:
+        tc = model.config.get_text_config()
+        layer_types = getattr(tc, "layer_types", None)
+        sw = getattr(tc, "sliding_window", None)
+    except Exception:
+        return DynamicCache()
+    if not layer_types or not sw:
+        return DynamicCache()
+    c = DynamicCache()
+    c.layers = [SpecSlidingLayer(sw, pad) if lt in ("sliding_attention", "chunked_attention")
+                else DynamicLayer() for lt in layer_types]
+    return c
+
+
+def make_draft_cache(draft_model, window, pad=64):
+    """Sliding cache for the DSpark draft's accumulated context keys. The draft is full-attention and
+    otherwise grows its context cache to the full sequence length; windowing it to the last `window`
+    tokens bounds memory. SpecSlidingLayer keeps cumulative_length == absolute position, so the draft's
+    position bookkeeping (position_ids[get_seq_length():...]) stays correct. Correctness-safe: the draft
+    only proposes; the target verifies every token.
+    """
+    from transformers.cache_utils import DynamicCache
+    try:
+        n = draft_model.config.get_text_config().num_hidden_layers
+    except Exception:
+        n = draft_model.config.num_hidden_layers
+    c = DynamicCache()
+    c.layers = [SpecSlidingLayer(window, pad) for _ in range(n)]
+    return c

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f
+size 32169626

tokenizer_config.json

@@ -0,0 +1,96 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>"
+  },
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "processor_class": "Gemma4UnifiedProcessor",
+  "response_schema": {
+    "properties": {
+      "content": {
+        "type": "string"
+      },
+      "role": {
+        "const": "assistant"
+      },
+      "thinking": {
+        "type": "string"
+      },
+      "tool_calls": {
+        "items": {
+          "properties": {
+            "function": {
+              "properties": {
+                "arguments": {
+                  "additionalProperties": {},
+                  "type": "object",
+                  "x-parser": "gemma4-tool-call"
+                },
+                "name": {
+                  "type": "string"
+                }
+              },
+              "type": "object",
+              "x-regex": "call\\:(?P<name>\\w+)(?P<arguments>\\{.*\\})"
+            },
+            "type": {
+              "const": "function"
+            }
+          },
+          "type": "object"
+        },
+        "type": "array",
+        "x-regex-iterator": "<\\|tool_call>(.*?)<tool_call\\|>"
+      }
+    },
+    "type": "object",
+    "x-regex": "(\\<\\|channel\\>thought\\n(?P<thinking>.*?)\\<channel\\|\\>)?(?P<tool_calls>\\<\\|tool_call\\>.*\\<tool_call\\|\\>)?(?P<content>(?:(?!\\<turn\\|\\>)(?!\\<\\|tool_response\\>).)+)?(?:\\<turn\\|\\>|\\<\\|tool_response\\>)?"
+  },
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>"
+}