Skip to content

[quantization] Add generation step to StaticGemma4Runtime #849

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
dvsav wants to merge 1 commit into
base: main
Choose a base branch
from
+217 −9
Draft

[quantization] Add generation step to StaticGemma4Runtime #849

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
226 changes: 217 additions & 9 deletions tico/quantization/recipes/debug/static_gemma4_runtime.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -607,14 +607,22 @@ def prefill(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:

if isinstance(out, tuple) and len(out) == 3:
hidden_states, new_k, new_v = out
# Write full-sequence K/V into the pre-allocated fixed cache.
# For prefill, new_k/new_v cover the entire max_seq, so no
# positional slicing is needed — copy_ preserves the buffer.
self.layer_caches[layer_idx].past_k.copy_(new_k)
self.layer_caches[layer_idx].past_v.copy_(new_v)
# Store full-length K/V for shared-KV consumer layers
# Write only valid K/V (excluding padding) into the cache.
# new_k/new_v from prefill have shape (B, kv_heads, max_seq, head_dim)
# but positions >= valid_length contain garbage from padding tokens.
valid_len = int(batch["valid_length"].item())
self.layer_caches[layer_idx].past_k[:, :, :valid_len, :].copy_(
new_k[:, :, :valid_len, :]
)
self.layer_caches[layer_idx].past_v[:, :, :valid_len, :].copy_(
new_v[:, :, :valid_len, :]
)
# Store valid-length K/V for shared-KV consumer layers
if bool(getattr(attn, "store_full_length_kv", False)):
shared_kv_states[layer_type] = (new_k, new_v)
shared_kv_states[layer_type] = (
new_k[:, :, :valid_len, :],
new_v[:, :, :valid_len, :],
)
else:
hidden_states = out

Expand Down Expand Up @@ -700,6 +708,13 @@ def decode_one(self, input_ids: torch.Tensor) -> torch.Tensor:
dtype=hidden_states.dtype,
)

# DEBUG: Print mask info
print(
f"[DEBUG decode_one] past_len={self.past_len}, mask shape={attention_masks['full_attention'].shape}"
)
valid_positions = (attention_masks["full_attention"] == 0.0).sum().item()
print(f"[DEBUG decode_one] Valid positions in mask: {valid_positions}")

# Compute PLE for the single decode token if enabled
per_layer_inputs = None
if self.text_model.hidden_size_per_layer_input:
Expand All @@ -713,6 +728,24 @@ def decode_one(self, input_ids: torch.Tensor) -> torch.Tensor:
# keyed by layer_type, updated after each store layer writes its delta.
shared_kv_states: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}

# Pre-populate shared_kv_states from layer_caches for shared consumer layers.
# This ensures consumer layers have access to KV cache from previous tokens
# before store layers update it in this decode step.
for layer_idx, layer_type in enumerate(self.text_config.layer_types):
if layer_type not in shared_kv_states:
attn = self.decode_layers[layer_idx].wrapped.self_attn.wrapped
is_shared = bool(getattr(attn, "is_kv_shared_layer", False))

if is_shared:
cache = self.layer_caches[layer_idx]
shared_kv_states[layer_type] = (
cache.past_k[:, :, : self.past_len, :],
cache.past_v[:, :, : self.past_len, :],
)
print(
f"[DEBUG decode_one] Pre-populated shared_kv_states[{layer_type}] from layer {layer_idx}, shape = {cache.past_k[:, :, : self.past_len, :].shape}"
)

for layer_idx, layer in enumerate(self.decode_layers):
layer_type = self.text_config.layer_types[layer_idx]
per_layer_input = (
Expand All @@ -728,6 +761,11 @@ def decode_one(self, input_ids: torch.Tensor) -> torch.Tensor:

cache = self.layer_caches[layer_idx]

# DEBUG: Print cache state before layer call
print(
f"[DEBUG decode_one] Layer {layer_idx}: past_key_value shape = {cache.past_k[:, :, : self.past_len, :].shape if not is_shared else 'N/A (shared)'}"
)

out = layer(
hidden_states=hidden_states,
attention_mask=attention_masks[layer_type],
Expand All @@ -746,9 +784,22 @@ def decode_one(self, input_ids: torch.Tensor) -> torch.Tensor:

if isinstance(out, tuple) and len(out) == 3:
hidden_states, new_k, new_v = out
# DEBUG: Print new_k/v info
print(
f"[DEBUG decode_one] Layer {layer_idx}: new_k shape={new_k.shape}, norm={new_k.norm().item():.6f}"
)

# Write single-token K/V delta into the fixed cache
cache.past_k[:, :, self.past_len : self.past_len + 1, :].copy_(new_k)
cache.past_v[:, :, self.past_len : self.past_len + 1, :].copy_(new_v)

# DEBUG: Verify write
written_k = cache.past_k[:, :, self.past_len, :]
matches = torch.allclose(written_k, new_k)
print(
f"[DEBUG decode_one] Layer {layer_idx}: written_k norm={written_k.norm().item():.6f}, matches={matches}"
)

# Update shared_kv_states for consumer layers.
# The valid K/V (prefill + all decode deltas so far) occupies
# positions 0..past_len inclusive in the cache. Slice to that
Expand All @@ -758,13 +809,65 @@ def decode_one(self, input_ids: torch.Tensor) -> torch.Tensor:
cache.past_k[:, :, : self.past_len + 1, :],
cache.past_v[:, :, : self.past_len + 1, :],
)
print(
f"[DEBUG decode_one] Layer {layer_idx}: Updated shared_kv_states, full K/V shape = {cache.past_k[:, :, : self.past_len + 1, :].shape}"
)
else:
hidden_states = out
print(
f"[DEBUG decode_one] Layer {layer_idx}: No KV returned, hidden_states shape = {hidden_states.shape}"
)

print(f"[DEBUG decode_one] Before past_len increment: {self.past_len}")
self.past_len += 1
print(f"[DEBUG decode_one] After past_len increment: {self.past_len}")
logits = self.lm_head(hidden_states)[:, -1, :]
return self._apply_final_logit_softcapping(logits)

@torch.no_grad()
def generate_greedy(
self,
batch: dict[str, torch.Tensor],
max_new_tokens: int,
eos_token_id: Optional[int] = None,
) -> torch.LongTensor:
"""Generate tokens greedily from a preprocessed static batch.

Args:
batch: The batch dict returned by ``build_static_inputs``.
max_new_tokens: Maximum number of new tokens to generate.
eos_token_id: Optional EOS token ID to stop generation early.
If not provided, uses ``text_config.eos_token_id``.

Returns:
Generated token IDs of shape ``(1, prompt_len + num_generated)``.
"""
# Run prefill
logits = self.prefill(batch)

# Get prompt length
prompt_len = int(batch["valid_length"].item())

# Extract valid input_ids (without padding)
input_ids = batch["llm_input_ids"][0, :prompt_len].clone() # (prompt_len,)

if eos_token_id is None:
eos_token_id = int(getattr(self.text_config, "eos_token_id", -1))

# Generate loop
for _ in range(max_new_tokens):
next_token = torch.argmax(logits, dim=-1) # (1,)
input_ids = torch.cat([input_ids, next_token], dim=0)

# Check EOS
if eos_token_id >= 0 and next_token.item() == eos_token_id:
break

# Decode step: next_token is (1,), need (1, 1) for decode_one
logits = self.decode_one(next_token.unsqueeze(0))

return input_ids.unsqueeze(0) # (1, generated_len)

def _apply_final_logit_softcapping(self, logits: torch.Tensor) -> torch.Tensor:
"""Apply Gemma4 final logit softcapping.

Expand Down Expand Up @@ -1581,6 +1684,98 @@ def verify_step_decode(
return rt_decode_logits_list[-1]


@torch.no_grad()
def verify_step_generation(
runtime: StaticGemma4Runtime,
batch: dict[str, torch.Tensor],
prompt: str,
image,
max_new_tokens: int = 16,
) -> torch.LongTensor:
"""Side-by-side validation of ``generate_greedy`` against HF reference.

Runs greedy generation from the runtime and HF reference, comparing
generated token sequences. The runtime uses quantized weights while
HF uses FP, so exact match is not expected — we report token-level
accuracy and first mismatch position.

Args:
runtime: The ``StaticGemma4Runtime`` instance.
batch: The batch dict returned by ``build_static_inputs``.
prompt: Original prompt string.
image: Original image.
max_new_tokens: Maximum number of new tokens to generate.

Returns:
Runtime-generated token IDs of shape ``(1, prompt_len + num_generated)``.
"""
# Runtime generation
runtime.reset_cache()
rt_generated = runtime.generate_greedy(batch, max_new_tokens, eos_token_id=None)

# HF reference generation
raw_inputs = batch.get("_raw_inputs", None)
if raw_inputs is None:
raw_inputs = runtime.processor(
text=prompt, images=image, return_tensors="pt", padding=False
)
raw_input_ids = raw_inputs["input_ids"].to(runtime.device)
pixel_values = raw_inputs["pixel_values"].to(runtime.device).to(runtime.model.dtype)
image_position_ids = raw_inputs.get("image_position_ids", None)
if image_position_ids is not None:
image_position_ids = image_position_ids.to(runtime.device)

# HF generate with explicit max_new_tokens
hf_generated = runtime.model.generate(
input_ids=raw_input_ids,
pixel_values=pixel_values,
image_position_ids=image_position_ids,
max_new_tokens=max_new_tokens,
do_sample=False,
pad_token_id=runtime.processor.tokenizer.pad_token_id,
eos_token_id=getattr(runtime.text_config, "eos_token_id", None),
)

# Compare sequences
rt_len = rt_generated.shape[1]
hf_len = hf_generated.shape[1]
min_len = min(rt_len, hf_len)

# Count matching tokens
matches = (rt_generated[:, :min_len] == hf_generated[:, :min_len]).sum().item()
accuracy = matches / min_len if min_len > 0 else 0.0

# Find first mismatch
first_mismatch = None
for i in range(min_len):
if rt_generated[0, i] != hf_generated[0, i]:
first_mismatch = i
break

print(f"[verify_step_generation] Runtime generated {rt_len} tokens")
print(f"[verify_step_generation] HF generated {hf_len} tokens")
print(
f"[verify_step_generation] Token accuracy: {accuracy * 100:.2f}% ({matches}/{min_len})"
)
if first_mismatch is not None:
print(f"[verify_step_generation] First mismatch at position {first_mismatch}")
else:
print("[verify_step_generation] All tokens match!")

# Print runtime generated text
rt_text = runtime.processor.decode(
rt_generated[0].tolist(), skip_special_tokens=True
)
hf_text = runtime.processor.decode(
hf_generated[0].tolist(), skip_special_tokens=True
)
print(f"[verify_step_generation] Runtime text: {rt_text}")
print(f"[verify_step_generation] HF text: {hf_text}")

print("[verify_step_generation] All checks passed.")
return rt_generated


def run_static_gemma4_runtime(cfg: StaticGemma4RuntimeConfig) -> None:
"""Run the Gemma4 E2B static runtime smoke flow.

Expand Down Expand Up @@ -1658,6 +1853,19 @@ def run_static_gemma4_runtime(cfg: StaticGemma4RuntimeConfig) -> None:
num_decode_steps=cfg.verify_steps,
)

# --- Step 8+: generation (not yet implemented) ---
print("[run_static_gemma4_runtime] SKIP: generation (not implemented in this PR)")
# --- Step 8: generation validation ---
print("[run_static_gemma4_runtime] Step 8: verify generation")
rt_generated = verify_step_generation(
runtime,
batch,
cfg.prompt,
image,
max_new_tokens=cfg.gen_steps,
)
# Detokenize and print generated text
rt_generated_text = processor.decode(
rt_generated[0].tolist(), skip_special_tokens=True
)
print(f"[run_static_gemma4_runtime] Generated text: {rt_generated_text}")

print("[run_static_gemma4_runtime] Done.")
Loading