[refactor] Refactor the test_chunked_prefill/decode.
This commit is contained in:
Zijie Tian
@@ -1,203 +1,111 @@
|
||||
"""
|
||||
Hook-based correctness test for chunked prefill attention.
|
||||
Correctness test for chunked prefill attention.
|
||||
|
||||
Uses PyTorch register_forward_hook() to capture real inference I/O,
|
||||
then compares against reference computation to locate bugs.
|
||||
|
||||
This test targets the integration layer (context setup, cpu_block_table management)
|
||||
which is where the needle test fails despite isolated attention tests passing.
|
||||
Captures Q and output during inference, then computes reference using
|
||||
CPU KV cache with standard flash attention.
|
||||
"""
|
||||
|
||||
import os
|
||||
os.environ["NANOVLLM_LOG_LEVEL"] = "DEBUG"
|
||||
os.environ["NANOVLLM_LOG_LEVEL"] = "WARNING"
|
||||
|
||||
import torch
|
||||
from random import randint, seed
|
||||
from typing import Dict, List
|
||||
from nanovllm import LLM, SamplingParams
|
||||
from nanovllm.utils.context import get_context
|
||||
from nanovllm.kvcache.chunked_attention import flash_attn_with_lse, merge_attention_outputs
|
||||
from flash_attn.flash_attn_interface import flash_attn_varlen_func
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Configuration
|
||||
# ============================================================
|
||||
|
||||
# Config
|
||||
MODEL_PATH = os.path.expanduser("~/models/Qwen3-4B-Instruct-2507/")
|
||||
MAX_MODEL_LEN = 32 * 1024
|
||||
MAX_MODEL_LEN = 128 * 1024
|
||||
NUM_GPU_BLOCKS = 2
|
||||
INPUT_LEN = 16 * 1024 # 4K tokens = 4 chunks with 1K block size
|
||||
INPUT_LEN = 16 * 1024
|
||||
BLOCK_SIZE = 1024
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Global capture storage
|
||||
# ============================================================
|
||||
|
||||
captures = []
|
||||
# State - capture Q and output for each (layer, chunk)
|
||||
captures: List[Dict] = []
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Hook Functions
|
||||
# ============================================================
|
||||
|
||||
def make_hook(layer_id):
|
||||
"""Create a forward hook for a specific layer."""
|
||||
def hook(module, inputs, output):
|
||||
q, k, v = inputs
|
||||
def make_ones_injection_hook():
|
||||
"""Inject Q=K=V=1.0 for deterministic testing."""
|
||||
def hook(module, inputs):
|
||||
ctx = get_context()
|
||||
if not ctx.is_prefill:
|
||||
return inputs
|
||||
|
||||
# Only capture prefill phase
|
||||
q, k, v = inputs[0], inputs[1], inputs[2]
|
||||
q_ones = torch.ones_like(q)
|
||||
k_ones = torch.ones_like(k)
|
||||
v_ones = torch.ones_like(v)
|
||||
return (q_ones, k_ones, v_ones) + inputs[3:]
|
||||
return hook
|
||||
|
||||
|
||||
def make_capture_hook(layer_id: int):
|
||||
"""Capture Q and output during prefill."""
|
||||
def hook(module, inputs, output):
|
||||
ctx = get_context()
|
||||
if not ctx.is_prefill:
|
||||
return
|
||||
|
||||
q, k, v = inputs
|
||||
chunk_idx = ctx.current_chunk_idx if hasattr(ctx, 'current_chunk_idx') else 0
|
||||
|
||||
capture_entry = {
|
||||
captures.append({
|
||||
'layer_id': layer_id,
|
||||
'chunk_idx': chunk_idx,
|
||||
'q': q.clone().cpu(),
|
||||
'k': k.clone().cpu(),
|
||||
'v': v.clone().cpu(),
|
||||
'output': output.clone().cpu(),
|
||||
'is_chunked_prefill': ctx.is_chunked_prefill,
|
||||
}
|
||||
|
||||
# For debugging: also capture CPU cache state for layer 0
|
||||
if layer_id == 0 and chunk_idx >= 2:
|
||||
kvcache_manager = ctx.kvcache_manager if hasattr(ctx, 'kvcache_manager') else None
|
||||
if kvcache_manager is not None and hasattr(kvcache_manager, 'offload_engine'):
|
||||
oe = kvcache_manager.offload_engine
|
||||
# Get what should have been loaded from CPU
|
||||
cpu_k0 = oe.k_cache_cpu[0, 0].clone().cpu() # Layer 0, CPU block 0
|
||||
cpu_k1 = oe.k_cache_cpu[0, 1].clone().cpu() # Layer 0, CPU block 1
|
||||
capture_entry['cpu_k0'] = cpu_k0
|
||||
capture_entry['cpu_k1'] = cpu_k1
|
||||
|
||||
captures.append(capture_entry)
|
||||
})
|
||||
return hook
|
||||
|
||||
|
||||
def register_hooks(llm):
|
||||
"""Register forward hooks on all Attention modules."""
|
||||
hooks = []
|
||||
model = llm.model_runner.model
|
||||
|
||||
for layer_idx, decoder_layer in enumerate(model.model.layers):
|
||||
attn_module = decoder_layer.self_attn.attn
|
||||
hook = attn_module.register_forward_hook(make_hook(layer_idx))
|
||||
hooks.append(hook)
|
||||
|
||||
return hooks
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Reference Computation
|
||||
# ============================================================
|
||||
|
||||
def compute_reference(layer_id, chunk_idx, scale, debug=False):
|
||||
def compute_reference(layer_id: int, chunk_idx: int, scale: float,
|
||||
k_cache_cpu: torch.Tensor, v_cache_cpu: torch.Tensor,
|
||||
block_size: int) -> torch.Tensor:
|
||||
"""
|
||||
Compute reference attention output for a specific layer and chunk.
|
||||
Compute reference output using CPU KV cache and standard flash attention.
|
||||
|
||||
Uses the captured k, v from all chunks up to and including chunk_idx.
|
||||
Concatenates all Q, K, V from chunks 0..chunk_idx and runs causal attention,
|
||||
then extracts output for the current chunk.
|
||||
"""
|
||||
# Filter captures for this layer
|
||||
# Get all captures for this layer up to chunk_idx
|
||||
layer_captures = [c for c in captures
|
||||
if c['layer_id'] == layer_id and c['chunk_idx'] <= chunk_idx]
|
||||
|
||||
if not layer_captures:
|
||||
return None
|
||||
|
||||
# Get current chunk's q
|
||||
current_capture = [c for c in layer_captures if c['chunk_idx'] == chunk_idx][0]
|
||||
q = current_capture['q'].cuda().unsqueeze(0) # [1, seqlen, nheads, headdim]
|
||||
|
||||
# Collect all k, v up to current chunk
|
||||
kv_list = []
|
||||
for c in sorted(layer_captures, key=lambda x: x['chunk_idx']):
|
||||
k = c['k'].cuda().unsqueeze(0) # [1, seqlen, nheads, headdim]
|
||||
v = c['v'].cuda().unsqueeze(0)
|
||||
kv_list.append((k, v, c['chunk_idx']))
|
||||
|
||||
if debug:
|
||||
print(f" Reference for L{layer_id} C{chunk_idx}:")
|
||||
print(f" q shape: {q.shape}, mean={q.mean().item():.4f}")
|
||||
print(f" kv_list: {len(kv_list)} chunks")
|
||||
for i, (k, v, cidx) in enumerate(kv_list):
|
||||
print(f" chunk {cidx}: k.mean={k.mean().item():.4f}, v.mean={v.mean().item():.4f}")
|
||||
|
||||
o_acc, lse_acc = None, None
|
||||
|
||||
# Previous chunks: non-causal attention
|
||||
for i in range(len(kv_list) - 1):
|
||||
k, v, _ = kv_list[i]
|
||||
o, lse = flash_attn_with_lse(q, k, v, softmax_scale=scale, causal=False)
|
||||
if o_acc is None:
|
||||
o_acc, lse_acc = o, lse
|
||||
else:
|
||||
o_acc, lse_acc = merge_attention_outputs(o_acc, lse_acc, o, lse)
|
||||
|
||||
# Current chunk: causal attention
|
||||
k_cur, v_cur, _ = kv_list[-1]
|
||||
o_cur, lse_cur = flash_attn_with_lse(q, k_cur, v_cur, softmax_scale=scale, causal=True)
|
||||
|
||||
if o_acc is None:
|
||||
return o_cur.squeeze(0).cpu()
|
||||
|
||||
final_o, _ = merge_attention_outputs(o_acc, lse_acc, o_cur, lse_cur)
|
||||
return final_o.squeeze(0).cpu()
|
||||
|
||||
|
||||
def compute_standard_reference(layer_id, chunk_idx, scale, debug=False):
|
||||
"""
|
||||
Compute reference using standard flash attention (single pass with all K, V).
|
||||
|
||||
This simulates what standard (non-chunked) prefill would produce.
|
||||
Concatenates all Q, K, V from chunks 0 to chunk_idx and runs a single
|
||||
causal attention pass, then extracts the output for the current chunk.
|
||||
"""
|
||||
# Filter captures for this layer
|
||||
layer_captures = [c for c in captures
|
||||
if c['layer_id'] == layer_id and c['chunk_idx'] <= chunk_idx]
|
||||
|
||||
if not layer_captures:
|
||||
return None
|
||||
|
||||
# Sort by chunk index
|
||||
layer_captures = sorted(layer_captures, key=lambda x: x['chunk_idx'])
|
||||
|
||||
# Concatenate all Q, K, V
|
||||
if not layer_captures:
|
||||
return None
|
||||
|
||||
# Collect Q from captures, K/V from CPU cache
|
||||
all_q = []
|
||||
all_k = []
|
||||
all_v = []
|
||||
chunk_lengths = []
|
||||
|
||||
for c in layer_captures:
|
||||
cidx = c['chunk_idx']
|
||||
q = c['q'].cuda() # [seqlen, nheads, headdim]
|
||||
k = c['k'].cuda()
|
||||
v = c['v'].cuda()
|
||||
all_q.append(q)
|
||||
all_k.append(k)
|
||||
all_v.append(v)
|
||||
chunk_lengths.append(q.shape[0])
|
||||
|
||||
# Concatenate along sequence dimension
|
||||
full_q = torch.cat(all_q, dim=0) # [total_seqlen, nheads, headdim]
|
||||
# Get K, V from CPU cache (already offloaded during prefill)
|
||||
# CPU cache shape: [num_layers, num_blocks, block_size, kv_heads, head_dim]
|
||||
k = k_cache_cpu[layer_id, cidx, :q.shape[0]].cuda()
|
||||
v = v_cache_cpu[layer_id, cidx, :q.shape[0]].cuda()
|
||||
all_k.append(k)
|
||||
all_v.append(v)
|
||||
|
||||
# Concatenate
|
||||
full_q = torch.cat(all_q, dim=0)
|
||||
full_k = torch.cat(all_k, dim=0)
|
||||
full_v = torch.cat(all_v, dim=0)
|
||||
|
||||
total_len = full_q.shape[0]
|
||||
|
||||
if debug:
|
||||
print(f" Standard Reference for L{layer_id} C{chunk_idx}:")
|
||||
print(f" full_q shape: {full_q.shape}, mean={full_q.mean().item():.4f}")
|
||||
print(f" full_k shape: {full_k.shape}, mean={full_k.mean().item():.4f}")
|
||||
print(f" chunk_lengths: {chunk_lengths}")
|
||||
|
||||
# Run standard causal flash attention
|
||||
# flash_attn_varlen_func expects: q, k, v with shape [total_seqlen, nheads, headdim]
|
||||
cu_seqlens = torch.tensor([0, total_len], dtype=torch.int32, device='cuda')
|
||||
|
||||
full_o = flash_attn_varlen_func(
|
||||
full_q, full_k, full_v,
|
||||
cu_seqlens_q=cu_seqlens,
|
||||
@@ -208,266 +116,81 @@ def compute_standard_reference(layer_id, chunk_idx, scale, debug=False):
|
||||
causal=True,
|
||||
)
|
||||
|
||||
# Extract output for current chunk only
|
||||
# Extract output for current chunk
|
||||
start_pos = sum(chunk_lengths[:-1])
|
||||
end_pos = sum(chunk_lengths)
|
||||
chunk_output = full_o[start_pos:end_pos]
|
||||
|
||||
if debug:
|
||||
print(f" full_o shape: {full_o.shape}")
|
||||
print(f" extracting positions [{start_pos}:{end_pos}]")
|
||||
print(f" chunk_output shape: {chunk_output.shape}, mean={chunk_output.mean().item():.4f}")
|
||||
|
||||
return chunk_output.cpu()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Test Runner
|
||||
# ============================================================
|
||||
|
||||
def run_test(verbose=True):
|
||||
"""Run the hook-based chunked prefill correctness test."""
|
||||
global captures
|
||||
captures = []
|
||||
|
||||
if verbose:
|
||||
print("=" * 70)
|
||||
print("Test: Hook-Based Chunked Prefill Correctness")
|
||||
print("=" * 70)
|
||||
print(f"Model: {MODEL_PATH}")
|
||||
print(f"Input length: {INPUT_LEN} tokens")
|
||||
print(f"Block size: {BLOCK_SIZE}")
|
||||
print(f"Expected chunks: {INPUT_LEN // BLOCK_SIZE}")
|
||||
print()
|
||||
|
||||
# Initialize LLM with CPU offload
|
||||
llm = LLM(
|
||||
MODEL_PATH,
|
||||
enforce_eager=True,
|
||||
max_model_len=MAX_MODEL_LEN,
|
||||
max_num_batched_tokens=MAX_MODEL_LEN,
|
||||
enable_cpu_offload=True,
|
||||
kvcache_block_size=BLOCK_SIZE,
|
||||
num_gpu_blocks=NUM_GPU_BLOCKS,
|
||||
)
|
||||
|
||||
# Get model info
|
||||
num_layers = len(llm.model_runner.model.model.layers)
|
||||
head_dim = llm.model_runner.model.model.layers[0].self_attn.attn.head_dim
|
||||
scale = head_dim ** -0.5
|
||||
|
||||
if verbose:
|
||||
print(f"Num layers: {num_layers}")
|
||||
print(f"Head dim: {head_dim}")
|
||||
print()
|
||||
|
||||
# Register hooks
|
||||
hooks = register_hooks(llm)
|
||||
if verbose:
|
||||
print(f"Registered {len(hooks)} hooks")
|
||||
|
||||
# Generate random prompt
|
||||
seed(42)
|
||||
prompt_token_ids = [[randint(0, 10000) for _ in range(INPUT_LEN)]]
|
||||
|
||||
# Run prefill only (max_tokens=1)
|
||||
if verbose:
|
||||
print("Running inference...")
|
||||
sampling_params = SamplingParams(temperature=0.6, max_tokens=1)
|
||||
outputs = llm.generate(prompt_token_ids, sampling_params, use_tqdm=False)
|
||||
|
||||
# Remove hooks
|
||||
for hook in hooks:
|
||||
hook.remove()
|
||||
|
||||
# Analyze captures
|
||||
if verbose:
|
||||
print(f"\nCaptured {len(captures)} attention calls")
|
||||
|
||||
# Group by layer and chunk
|
||||
chunks_per_layer = {}
|
||||
for c in captures:
|
||||
layer_id = c['layer_id']
|
||||
chunk_idx = c['chunk_idx']
|
||||
if layer_id not in chunks_per_layer:
|
||||
chunks_per_layer[layer_id] = set()
|
||||
chunks_per_layer[layer_id].add(chunk_idx)
|
||||
|
||||
if verbose:
|
||||
print("Chunks per layer:", {k: sorted(v) for k, v in chunks_per_layer.items()})
|
||||
print()
|
||||
|
||||
# First, verify CPU cache data integrity
|
||||
if verbose:
|
||||
print("\n--- CPU Cache Verification (Layer 0) ---")
|
||||
# Get original k from chunk 0 and chunk 1 captures
|
||||
chunk0_k = None
|
||||
chunk1_k = None
|
||||
chunk2_capture = None
|
||||
for c in captures:
|
||||
if c['layer_id'] == 0:
|
||||
if c['chunk_idx'] == 0:
|
||||
chunk0_k = c['k']
|
||||
elif c['chunk_idx'] == 1:
|
||||
chunk1_k = c['k']
|
||||
elif c['chunk_idx'] == 2:
|
||||
chunk2_capture = c
|
||||
|
||||
if chunk0_k is not None and chunk2_capture is not None and 'cpu_k0' in chunk2_capture:
|
||||
cpu_k0 = chunk2_capture['cpu_k0']
|
||||
diff_k0 = (chunk0_k - cpu_k0).abs().max().item()
|
||||
print(f"Chunk 0 k vs CPU cache block 0: max_diff={diff_k0:.6f}")
|
||||
if diff_k0 > 1e-3:
|
||||
print(f" WARNING: CPU cache block 0 differs from original chunk 0 k!")
|
||||
print(f" Original k[0,0,:5] = {chunk0_k[0,0,:5].tolist()}")
|
||||
print(f" CPU k0[0,0,:5] = {cpu_k0[0,0,:5].tolist()}")
|
||||
|
||||
if chunk1_k is not None and chunk2_capture is not None and 'cpu_k1' in chunk2_capture:
|
||||
cpu_k1 = chunk2_capture['cpu_k1']
|
||||
diff_k1 = (chunk1_k - cpu_k1).abs().max().item()
|
||||
print(f"Chunk 1 k vs CPU cache block 1: max_diff={diff_k1:.6f}")
|
||||
if diff_k1 > 1e-3:
|
||||
print(f" WARNING: CPU cache block 1 differs from original chunk 1 k!")
|
||||
print(f" Original k[0,0,:5] = {chunk1_k[0,0,:5].tolist()}")
|
||||
print(f" CPU k1[0,0,:5] = {cpu_k1[0,0,:5].tolist()}")
|
||||
|
||||
print()
|
||||
|
||||
# ================================================================
|
||||
# Test 1: Verify against merge-based reference (same algorithm)
|
||||
# ================================================================
|
||||
if verbose:
|
||||
print("--- Test 1: Merge-based Reference (verifies merge algorithm) ---")
|
||||
|
||||
all_passed_merge = True
|
||||
results_merge = []
|
||||
first_fail_debug = True
|
||||
|
||||
for c in captures:
|
||||
layer_id = c['layer_id']
|
||||
chunk_idx = c['chunk_idx']
|
||||
|
||||
if chunk_idx == 0:
|
||||
continue
|
||||
|
||||
debug_this = (chunk_idx >= 2 and layer_id == 0 and first_fail_debug)
|
||||
ref_output = compute_reference(layer_id, chunk_idx, scale, debug=debug_this)
|
||||
if ref_output is None:
|
||||
continue
|
||||
|
||||
actual_output = c['output']
|
||||
diff = (actual_output - ref_output).abs()
|
||||
max_diff = diff.max().item()
|
||||
mean_diff = diff.mean().item()
|
||||
|
||||
tol = 1e-2
|
||||
passed = max_diff < tol
|
||||
all_passed_merge = all_passed_merge and passed
|
||||
|
||||
status = "PASS" if passed else "FAIL"
|
||||
results_merge.append((layer_id, chunk_idx, passed, max_diff, mean_diff))
|
||||
|
||||
if verbose:
|
||||
print(f"[{status}] Layer {layer_id:2d}, Chunk {chunk_idx}: "
|
||||
f"max_diff={max_diff:.6f} mean_diff={mean_diff:.8f}")
|
||||
|
||||
if not passed and first_fail_debug:
|
||||
first_fail_debug = False
|
||||
print(f" Debug: actual_output shape={actual_output.shape}, mean={actual_output.mean().item():.4f}")
|
||||
print(f" Debug: ref_output shape={ref_output.shape}, mean={ref_output.mean().item():.4f}")
|
||||
max_idx = diff.argmax()
|
||||
flat_actual = actual_output.flatten()
|
||||
flat_ref = ref_output.flatten()
|
||||
print(f" Debug: max_diff at idx={max_idx.item()}, actual={flat_actual[max_idx].item():.4f}, ref={flat_ref[max_idx].item():.4f}")
|
||||
|
||||
print()
|
||||
|
||||
# ================================================================
|
||||
# Test 2: Verify against standard flash attention (single pass)
|
||||
# ================================================================
|
||||
if verbose:
|
||||
print("--- Test 2: Standard FlashAttn Reference (verifies correctness vs non-chunked) ---")
|
||||
|
||||
all_passed_standard = True
|
||||
results_standard = []
|
||||
first_fail_debug = True
|
||||
|
||||
for c in captures:
|
||||
layer_id = c['layer_id']
|
||||
chunk_idx = c['chunk_idx']
|
||||
|
||||
if chunk_idx == 0:
|
||||
continue
|
||||
|
||||
debug_this = (chunk_idx >= 2 and layer_id == 0 and first_fail_debug)
|
||||
std_ref_output = compute_standard_reference(layer_id, chunk_idx, scale, debug=debug_this)
|
||||
if std_ref_output is None:
|
||||
continue
|
||||
|
||||
actual_output = c['output']
|
||||
diff = (actual_output - std_ref_output).abs()
|
||||
max_diff = diff.max().item()
|
||||
mean_diff = diff.mean().item()
|
||||
|
||||
tol = 1e-2
|
||||
passed = max_diff < tol
|
||||
all_passed_standard = all_passed_standard and passed
|
||||
|
||||
status = "PASS" if passed else "FAIL"
|
||||
results_standard.append((layer_id, chunk_idx, passed, max_diff, mean_diff))
|
||||
|
||||
if verbose:
|
||||
print(f"[{status}] Layer {layer_id:2d}, Chunk {chunk_idx}: "
|
||||
f"max_diff={max_diff:.6f} mean_diff={mean_diff:.8f}")
|
||||
|
||||
if not passed and first_fail_debug:
|
||||
first_fail_debug = False
|
||||
print(f" Debug: actual_output shape={actual_output.shape}, mean={actual_output.mean().item():.4f}")
|
||||
print(f" Debug: std_ref_output shape={std_ref_output.shape}, mean={std_ref_output.mean().item():.4f}")
|
||||
max_idx = diff.argmax()
|
||||
flat_actual = actual_output.flatten()
|
||||
flat_ref = std_ref_output.flatten()
|
||||
print(f" Debug: max_diff at idx={max_idx.item()}, actual={flat_actual[max_idx].item():.4f}, ref={flat_ref[max_idx].item():.4f}")
|
||||
|
||||
print()
|
||||
print("=" * 70)
|
||||
|
||||
# Summary
|
||||
total_merge = len(results_merge)
|
||||
passed_merge = sum(1 for r in results_merge if r[2])
|
||||
total_standard = len(results_standard)
|
||||
passed_standard = sum(1 for r in results_standard if r[2])
|
||||
|
||||
print(f"Merge-based reference: {passed_merge}/{total_merge} tests passed")
|
||||
print(f"Standard FlashAttn ref: {passed_standard}/{total_standard} tests passed")
|
||||
|
||||
all_passed = all_passed_merge and all_passed_standard
|
||||
|
||||
if not all_passed_merge:
|
||||
print("\nFailed merge-based tests:")
|
||||
for layer_id, chunk_idx, passed, max_diff, mean_diff in results_merge:
|
||||
if not passed:
|
||||
print(f" - Layer {layer_id}, Chunk {chunk_idx}: max_diff={max_diff:.6f}")
|
||||
|
||||
if not all_passed_standard:
|
||||
print("\nFailed standard FlashAttn tests:")
|
||||
for layer_id, chunk_idx, passed, max_diff, mean_diff in results_standard:
|
||||
if not passed:
|
||||
print(f" - Layer {layer_id}, Chunk {chunk_idx}: max_diff={max_diff:.6f}")
|
||||
|
||||
print()
|
||||
return all_passed
|
||||
return full_o[start_pos:end_pos].cpu()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Main
|
||||
# ============================================================
|
||||
|
||||
if __name__ == "__main__":
|
||||
passed = run_test(verbose=True)
|
||||
llm = LLM(
|
||||
MODEL_PATH,
|
||||
enforce_eager=True,
|
||||
max_model_len=MAX_MODEL_LEN,
|
||||
max_num_batched_tokens=MAX_MODEL_LEN,
|
||||
enable_cpu_offload=True,
|
||||
kvcache_block_size=BLOCK_SIZE,
|
||||
num_gpu_blocks=NUM_GPU_BLOCKS,
|
||||
dtype="float16",
|
||||
)
|
||||
|
||||
if passed:
|
||||
print("test_chunked_prefill_hook: PASSED")
|
||||
else:
|
||||
print("test_chunked_prefill_hook: FAILED")
|
||||
exit(1)
|
||||
# Get model info
|
||||
num_layers = len(llm.model_runner.model.model.layers)
|
||||
head_dim = llm.model_runner.model.model.layers[0].self_attn.attn.head_dim
|
||||
scale = head_dim ** -0.5
|
||||
|
||||
# Register hooks
|
||||
hooks = []
|
||||
for layer_idx, decoder_layer in enumerate(llm.model_runner.model.model.layers):
|
||||
# Pre-hook: inject all ones for Q, K, V
|
||||
# pre_hook = decoder_layer.self_attn.attn.register_forward_pre_hook(make_ones_injection_hook())
|
||||
# hooks.append(pre_hook)
|
||||
# Post-hook: capture Q, K, V, output
|
||||
post_hook = decoder_layer.self_attn.attn.register_forward_hook(make_capture_hook(layer_idx))
|
||||
hooks.append(post_hook)
|
||||
|
||||
# Run inference
|
||||
seed(42)
|
||||
prompt_token_ids = [[randint(0, 10000) for _ in range(INPUT_LEN)]]
|
||||
outputs = llm.generate(prompt_token_ids, SamplingParams(temperature=0.6, max_tokens=1), use_tqdm=False)
|
||||
|
||||
# Remove hooks
|
||||
for hook in hooks:
|
||||
hook.remove()
|
||||
|
||||
# Get CPU cache reference
|
||||
offload_engine = llm.model_runner.kvcache_manager.offload_engine
|
||||
k_cache_cpu = offload_engine.k_cache_cpu.clone()
|
||||
v_cache_cpu = offload_engine.v_cache_cpu.clone()
|
||||
|
||||
# Verify: compare actual output with reference computed from CPU cache
|
||||
all_passed = True
|
||||
num_chunks = INPUT_LEN // BLOCK_SIZE
|
||||
|
||||
for idx,c in enumerate(captures):
|
||||
layer_id = c['layer_id']
|
||||
chunk_idx = c['chunk_idx']
|
||||
|
||||
# Skip chunk 0 (no previous KV to load)
|
||||
if chunk_idx == 0:
|
||||
continue
|
||||
|
||||
ref_output = compute_reference(layer_id, chunk_idx, scale, k_cache_cpu, v_cache_cpu, BLOCK_SIZE)
|
||||
if ref_output is None:
|
||||
continue
|
||||
|
||||
actual_output = c['output']
|
||||
diff = (actual_output - ref_output).abs()
|
||||
max_diff = diff.max().item()
|
||||
|
||||
passed = max_diff < 1e-1 # float16 tolerance
|
||||
all_passed = all_passed and passed
|
||||
|
||||
if not passed:
|
||||
print(f"[FAIL] Layer {layer_id}, Chunk {chunk_idx}: max_diff={max_diff:.6f}")
|
||||
__import__('pdb').set_trace()
|
||||
|
||||
print(f"test_chunked_prefill_hook: {'PASSED' if all_passed else 'FAILED'}")
|
||||
|
||||
Reference in New Issue
Block a user