Skip to content

debug spec decode #21

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

Open
wxsIcey wants to merge 1 commit into
base: main
Choose a base branch
from
Open

debug spec decode #21

Show file tree
Hide file tree
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
8 changes: 8 additions & 0 deletions vllm/model_executor/models/qwen2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -429,6 +429,14 @@ def forward(
start=self.start_layer,
):
hidden_states, residual = layer(positions, hidden_states, residual)
if positions.numel() == 6 and positions[-1].item() == 5:
model_name = getattr(self.config, "_name_or_path", "unknown")
print(
f"[layer-debug] model={model_name} layer={idx} pos=5 "
f"hidden={hidden_states[5, :8].detach().cpu().tolist()} "
f"residual={residual[5, :8].detach().cpu().tolist()}",
flush=True,
)
self._maybe_add_hidden_state(
aux_hidden_states, idx + 1, hidden_states, residual
)
Expand Down
59 changes: 59 additions & 0 deletions vllm/v1/attention/backends/flash_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1396,6 +1396,56 @@ def forward(
block_table = block_table[:, :num_pages]
num_splits = 1

layer_name = getattr(layer, "layer_name", "")
debug_flash_attn = (
"draft_model" in layer_name
and ".layers.0." in layer_name
and not getattr(self, "_debug_flash_attn_printed", False)
and cu_seqlens_q.numel() == 2
)
if debug_flash_attn:
kv_len = int(seqused_k[0].detach().cpu().item())
block_size = key_cache.shape[1]
num_blocks = min(
block_table.shape[1],
(kv_len + block_size - 1) // block_size,
)
block_ids = block_table[0, :num_blocks].to(
device=key_cache.device, dtype=torch.long
)
positions = torch.arange(kv_len, device=key_cache.device)
logical_blocks = block_ids[positions // block_size]
block_offsets = positions % block_size
logical_key = key_cache[logical_blocks, block_offsets]
logical_value = value_cache[logical_blocks, block_offsets]
print(
f"[fa-debug] layer={layer_name} "
f"fa_version={self.vllm_flash_attn_version} "
f"q_shape={tuple(query[:num_actual_tokens].shape)} "
f"q_dtype={query.dtype} "
f"cu_seqlens_q={cu_seqlens_q.detach().cpu().tolist()} "
f"seqused_k={seqused_k.detach().cpu().tolist()} "
f"max_seqlen_q={max_seqlen_q} max_seqlen_k={max_seqlen_k} "
f"block_size={block_size} "
f"block_table={block_ids.detach().cpu().tolist()} "
f"scale={self.scale} causal={causal} "
f"window={sliding_window_size} softcap={self.logits_soft_cap} "
f"num_splits={num_splits}",
flush=True,
)
print(
f"[fa-debug] op_input "
f"Q_shape={tuple(query[:num_actual_tokens].shape)} "
f"Q={query[:num_actual_tokens].detach().cpu().tolist()} "
f"logical_K_shape={tuple(logical_key.shape)} "
f"dtype={logical_key.dtype} "
f"logical_K={logical_key.detach().cpu().tolist()} "
f"logical_V_shape={tuple(logical_value.shape)} "
f"dtype={logical_value.dtype} "
f"logical_V={logical_value.detach().cpu().tolist()}",
flush=True,
)

_FA4_DENSE_ATTENTION_KERNEL(
q=query[:num_actual_tokens],
k=key_cache,
Expand All @@ -1422,6 +1472,15 @@ def forward(
mask_mod=rswa_mask_mod_fn or mm_mask_mod,
aux_tensors=rswa_aux or mm_aux,
)
if debug_flash_attn:
print(
f"[fa-debug] op_output "
f"shape={tuple(output[:num_actual_tokens].shape)} "
f"dtype={output.dtype} "
f"{output[:num_actual_tokens].detach().cpu().tolist()}",
flush=True,
)
self._debug_flash_attn_printed = True
return output

def do_kv_cache_update(
Expand Down
13 changes: 13 additions & 0 deletions vllm/v1/worker/gpu/spec_decode/speculator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -448,7 +448,10 @@ def sample_draft(
draft_logits: torch.Tensor | None,
spec_step_idx: int = 0,
num_steps: int = 1,
debug: bool = False,
) -> torch.Tensor:
if debug:
print("draft_hidden:", hidden_states[0].detach().cpu().tolist(), flush=True)
if draft_logits is None and self.use_local_argmax_reduction:
return self.get_draft_top_tokens(hidden_states, spec_step_idx)

Expand All @@ -473,6 +476,16 @@ def sample_draft(
)
else:
sampled = logits.argmax(dim=-1)
if debug:
top5_logits, top5_ids = logits[0].topk(min(5, logits.shape[-1]))
print(
f"[draft-debug] step={int(draft_step.detach().cpu().item())} "
f"pos={sample_src_positions.detach().cpu().tolist()} "
f"top5_ids={top5_ids.detach().cpu().tolist()} "
f"top5_logits={top5_logits.detach().cpu().tolist()} "
f"proposed={sampled.detach().cpu().tolist()}",
flush=True,
)
self._maybe_predict_acceptance(logits, idx_mapping, draft_step)
return sampled

Expand Down
51 changes: 41 additions & 10 deletions vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,7 @@ def _prefill(
num_reqs: int,
skip_attn: bool,
num_tokens_across_dp: torch.Tensor | None,
debug: bool = False,
) -> None:
if skip_attn:
src_tokens = input_batch.num_tokens
Expand Down Expand Up @@ -231,6 +232,17 @@ def _prefill(
prefill_slot_maps_by_layer = build_slot_mappings_by_layer(
prefill_slot_mappings, kv_cache_config
)
if debug:
input_ids = (
self.expanded_input_ids[:total_expanded].detach().cpu().tolist()
)
positions = (
self.expanded_positions[:total_expanded].detach().cpu().tolist()
)
print(
f"[draft-debug] step=0 input_ids={input_ids} positions={positions}",
flush=True,
)
hidden_states = self._run_model(
self.expanded_input_ids[:total_expanded],
self.expanded_positions[:total_expanded],
Expand All @@ -250,6 +262,7 @@ def _prefill(
self.seeds,
self.current_draft_step,
self.draft_logits,
debug=debug,
)

def _multi_step_decode(
Expand All @@ -259,6 +272,7 @@ def _multi_step_decode(
num_reqs: int,
skip_attn: bool,
num_tokens_across_dp: torch.Tensor | None,
debug: bool = False,
) -> None:
if skip_attn:
last_positions = input_batch.positions[
Expand Down Expand Up @@ -329,6 +343,18 @@ def _multi_step_decode(
# Include the correction token inserted during prefill.
step=step + 1,
)
if debug:
input_ids = (
self.input_buffers.input_ids[:num_reqs].detach().cpu().tolist()
)
positions = (
self.input_buffers.positions[:num_reqs].detach().cpu().tolist()
)
print(
f"[draft-debug] step={step} "
f"input_ids={input_ids} positions={positions}",
flush=True,
)
hidden_states = self._run_model(
self.input_buffers.input_ids[:num_reqs],
self.input_buffers.positions[:num_reqs],
Expand All @@ -346,6 +372,7 @@ def _multi_step_decode(
self.seeds,
self.current_draft_step,
self.draft_logits,
debug=debug,
)

@torch.inference_mode()
Expand All @@ -372,6 +399,7 @@ def propose(
assert self.model is not None

num_reqs = input_batch.num_reqs
debug_propose = num_reqs == 1 and not dummy_run and not is_profile
skip_attn = dummy_run and skip_attn_for_dummy_run
num_tokens_across_dp = (
dp_sync_state.num_tokens_across_dp if dp_sync_state is not None else None
Expand All @@ -393,20 +421,23 @@ def propose(
num_reqs,
skip_attn,
num_tokens_across_dp,
debug=debug_propose,
)

if self.num_speculative_steps == 1:
return self.draft_tokens[:num_reqs, :1]

self._multi_step_decode(
input_batch,
num_rejected,
num_reqs,
skip_attn,
num_tokens_across_dp,
)
proposed_tokens = self.draft_tokens[:num_reqs, :1]
else:
self._multi_step_decode(
input_batch,
num_rejected,
num_reqs,
skip_attn,
num_tokens_across_dp,
debug=debug_propose,
)
proposed_tokens = self.draft_tokens[:num_reqs]

return self.draft_tokens[:num_reqs]
return proposed_tokens


@triton.jit
Expand Down