File size: 6,820 Bytes
d91766b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | from __future__ import annotations
import torch
from diffulex.engine.request import DllmReq
from .core import SamplerBase
from .output import SampleOutputBase
class SamplerNoShiftLogits(SamplerBase):
pass
class DllmSamplerNoShiftBase(SamplerNoShiftLogits):
output_cls = SampleOutputBase
@staticmethod
def _split_logits_per_req(attn_metadata, reqs: list[DllmReq], logits: torch.Tensor) -> tuple[torch.Tensor, ...]:
cu = attn_metadata.cu_seqlens_q
if cu is not None and len(cu) == len(reqs) + 1:
split_sizes = [(int(cu[i + 1]) - int(cu[i])) for i in range(len(reqs))]
else:
split_sizes = [
len(req.running_sequence) if attn_metadata.is_prefill[idx] else req.chunk_size
for idx, req in enumerate(reqs)
]
return torch.split(logits, split_sizes, dim=0)
@staticmethod
def _prefill_mask_token_local_ids(req: DllmReq, block, req_logits: torch.Tensor) -> list[int]:
# Use contiguous cached prefix length for prefill-logits alignment.
# `in_cache_len` may include non-prefix cached blocks and can overshoot.
prefix_offset = int(req.contiguous_in_cache_prefix_len)
local_ids = [idx - prefix_offset for idx in block.mask_token_global_ids]
if not local_ids:
return local_ids
if min(local_ids) < 0 or max(local_ids) >= req_logits.shape[0]:
raise IndexError(
"Prefill mask-token logits index out of bounds: "
f"req_id={getattr(req, 'req_id', '?')}, "
f"block_id={getattr(block, 'block_id', '?')}, "
f"in_cache_len={prefix_offset}, "
f"global_ids={block.mask_token_global_ids}, "
f"local_ids={local_ids}, "
f"req_logits_len={req_logits.shape[0]}"
)
return local_ids
def forward(
self,
reqs: list[DllmReq],
logits: torch.Tensor,
temperatures: torch.Tensor,
top_p=None,
top_k=None,
margin_confidence=False,
neg_entropy=False,
**kwargs,
):
attn_metadata = self.fetch_attn_metadata()
split_logits = self._split_logits_per_req(attn_metadata, reqs, logits)
accepted_ids_map = {}
sampled_tokens_map = {}
true_local_ids_map = {}
mask_token_rel_ids_map = {}
confidence_map = {}
initial_confidence_map = {}
for idx, (temperature, req, req_logits) in enumerate(zip(temperatures, reqs, split_logits)):
true_local_ids_sub_map = {}
accepted_ids_sub_map = {}
sampled_tokens_sub_map = {}
mask_token_rel_ids_sub_map = {}
confidence_sub_map = {}
initial_confidence_sub_map = {}
for block_id, block in enumerate(req.dllm_blocks):
if not block.is_active or (block.num_mask_tokens == 0):
continue
if len(block.mask_token_global_ids) == 0:
continue
if attn_metadata.is_prefill[idx]:
# Prefix-cache prefill can produce q_len=0 for some requests in mixed batches.
# In that case there are no logits to sample for this req in this step.
if req_logits.shape[0] == 0:
continue
local_ids = self._prefill_mask_token_local_ids(req, block, req_logits)
mask_token_logits = req_logits[local_ids, ...]
else:
buf_offset = block.start - req.dllm_block_buffer.first_running_block.start
buf_ids = [buf_offset + i for i in block.mask_token_relative_ids]
mask_token_logits = req_logits[buf_ids, ...]
confidence, sampled_tokens, initial_confidence = self.sample_tokens(
mask_token_logits,
temperature,
top_p=top_p,
top_k=top_k,
neg_entropy=(neg_entropy == "neg_entropy"),
margin_confidence=(margin_confidence == "margin_confidence"),
forbidden_token_ids=[int(block.mask_token_id)],
)
accepted_ids = self._compute_accepted_ids(
block, confidence, initial_confidence, sampled_tokens, **kwargs
)
block_id_str = str(block_id)
accepted_ids_list = accepted_ids.to(device="cpu").tolist()
true_local_ids_sub_map[block_id_str] = [block.mask_token_relative_ids[i] for i in accepted_ids_list]
accepted_ids_sub_map[block_id_str] = accepted_ids_list
sampled_tokens_sub_map[block_id_str] = sampled_tokens.to(device="cpu").tolist()
mask_token_rel_ids_sub_map[block_id_str] = list(block.mask_token_relative_ids)
confidence_sub_map[block_id_str] = confidence.to(device="cpu").tolist()
initial_confidence_sub_map[block_id_str] = initial_confidence.to(device="cpu").tolist()
req_id_str = str(req.req_id)
true_local_ids_map[req_id_str] = true_local_ids_sub_map
accepted_ids_map[req_id_str] = accepted_ids_sub_map
sampled_tokens_map[req_id_str] = sampled_tokens_sub_map
mask_token_rel_ids_map[req_id_str] = mask_token_rel_ids_sub_map
confidence_map[req_id_str] = confidence_sub_map
initial_confidence_map[req_id_str] = initial_confidence_sub_map
sample_output = self.output_cls(
true_local_ids_map=true_local_ids_map,
accepted_ids_map=accepted_ids_map,
sampled_tokens_map=sampled_tokens_map,
mask_token_rel_ids_map=mask_token_rel_ids_map,
confidence_map=confidence_map,
initial_confidence_map=initial_confidence_map,
)
return self._postprocess_sample_output(
reqs=reqs,
split_logits=split_logits,
temperatures=temperatures,
sample_output=sample_output,
attn_metadata=attn_metadata,
**kwargs,
)
def _postprocess_sample_output(
self,
reqs: list[DllmReq],
split_logits: tuple[torch.Tensor, ...],
temperatures: torch.Tensor,
sample_output: SampleOutputBase,
attn_metadata,
**kwargs,
) -> SampleOutputBase:
del reqs, split_logits, temperatures, attn_metadata, kwargs
return sample_output
def _compute_accepted_ids(
self,
block,
confidence: torch.Tensor,
initial_confidence: torch.Tensor,
sampled_tokens: torch.Tensor,
**kwargs,
) -> torch.Tensor:
raise NotImplementedError
|