Fix DSpark reference generation block width

#2
by yisun0618 - opened
Files changed (1) hide show
  1. dflash.py +11 -5
dflash.py CHANGED
@@ -395,8 +395,9 @@ class DFlashDraftModel(Qwen3PreTrainedModel):
395
  max_length = num_input_tokens + max_new_tokens
396
 
397
  block_size = self.block_size
 
398
  output_ids = torch.full(
399
- (1, max_length + block_size),
400
  self.mask_token_id,
401
  dtype=torch.long,
402
  device=target.device,
@@ -430,9 +431,14 @@ class DFlashDraftModel(Qwen3PreTrainedModel):
430
  acceptance_lengths = []
431
  start = input_ids.shape[1]
432
  while start < max_length:
433
- block_output_ids = output_ids[:, start : start + block_size].clone()
434
- block_position_ids = position_ids[:, start : start + block_size]
435
- noise_embedding = target.model.embed_tokens(block_output_ids)
 
 
 
 
 
436
  draft_logits = target.lm_head(
437
  self(
438
  target_hidden=target_hidden,
@@ -443,7 +449,7 @@ class DFlashDraftModel(Qwen3PreTrainedModel):
443
  past_key_values=past_key_values_draft,
444
  use_cache=True,
445
  is_causal=False,
446
- )[:, -block_size + 1 :, :]
447
  )
448
  past_key_values_draft.crop(start)
449
  block_output_ids[:, 1:] = sample(draft_logits)
 
395
  max_length = num_input_tokens + max_new_tokens
396
 
397
  block_size = self.block_size
398
+ verify_width = block_size + 1
399
  output_ids = torch.full(
400
+ (1, max_length + verify_width),
401
  self.mask_token_id,
402
  dtype=torch.long,
403
  device=target.device,
 
431
  acceptance_lengths = []
432
  start = input_ids.shape[1]
433
  while start < max_length:
434
+ # The draft has ``block_size`` noise rows. Row j predicts the token at
435
+ # start+j+1, so target verification consumes the anchor plus all draft
436
+ # proposals (``block_size + 1`` tokens).
437
+ block_output_ids = output_ids[:, start : start + verify_width].clone()
438
+ block_position_ids = position_ids[:, start : start + verify_width]
439
+ noise_embedding = target.model.embed_tokens(
440
+ block_output_ids[:, :block_size]
441
+ )
442
  draft_logits = target.lm_head(
443
  self(
444
  target_hidden=target_hidden,
 
449
  past_key_values=past_key_values_draft,
450
  use_cache=True,
451
  is_causal=False,
452
+ )[:, -block_size:, :]
453
  )
454
  past_key_values_draft.crop(start)
455
  block_output_ids[:, 1:] = sample(draft_logits)