Spaces:
Paused
Expose query_position and query_token_id in /attention/row response
Browse filesThe cached attention[-1] is the row of the *predicting* query position
(the position whose logits selected the step-s token), not the
predicted token's self-attention. The frontend was labelling this row
as "attention influences on this token", which implies a self-attention
relationship the data does not represent.
Record (query_position, query_token_id) per generation step on the
MatrixCache during analysis (both SSE and non-SSE paths). Surface them
in the /attention/row response so the Why this token? panel can show
explicit provenance ("Attention used to predict this token, from query
position N — `<token>`") instead of the misleading framing.
This is a documentation/honesty fix, not a data fix — the underlying
attention values are unchanged. It supports the project's stated
"grounded hypotheses, not causal claims" framing.
- backend/model_service.py +71 -3
|
@@ -140,10 +140,35 @@ class MatrixCache:
|
|
| 140 |
for older requests.
|
| 141 |
"""
|
| 142 |
with self._lock:
|
| 143 |
-
self._request_meta
|
|
|
|
| 144 |
"num_heads": num_heads,
|
| 145 |
"num_layers": num_layers,
|
| 146 |
"model_id": model_id,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 147 |
}
|
| 148 |
|
| 149 |
def get_request_metadata(self, request_id: str) -> Optional[dict]:
|
|
@@ -152,6 +177,18 @@ class MatrixCache:
|
|
| 152 |
meta = self._request_meta.get(request_id)
|
| 153 |
return dict(meta) if meta else None
|
| 154 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 155 |
def get(self, request_id: str, step: int, layer: int, head: int) -> Optional[dict]:
|
| 156 |
"""Retrieve matrix data, returning None if expired or not found."""
|
| 157 |
key = f"{request_id}:{step}:{layer}:{head}"
|
|
@@ -2048,6 +2085,16 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
|
|
| 2048 |
output_hidden_states=True
|
| 2049 |
)
|
| 2050 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 2051 |
# Get logits for next token
|
| 2052 |
raw_logits = outputs.logits[0, -1, :].clone() # Clone raw logits before any scaling
|
| 2053 |
|
|
@@ -2952,6 +2999,14 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
|
|
| 2952 |
output_hidden_states=True
|
| 2953 |
)
|
| 2954 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 2955 |
# Get logits for next token
|
| 2956 |
raw_logits = outputs.logits[0, -1, :].clone() # Clone raw logits before any scaling
|
| 2957 |
|
|
@@ -3913,12 +3968,21 @@ async def get_attention_row(
|
|
| 3913 |
- layer: Layer index
|
| 3914 |
- head: Head index (null if aggregated)
|
| 3915 |
- aggregate_mode: Mode used if aggregated (null otherwise)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3916 |
"""
|
| 3917 |
# Aggregation parameters (num_heads in particular) come from per-request
|
| 3918 |
# metadata recorded at analysis time, not from the live model config.
|
| 3919 |
# This keeps aggregate rows correct even after the user switches models
|
| 3920 |
# mid-session — the row is averaged over the heads of the model that
|
| 3921 |
# actually produced the cached attention.
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3922 |
if head is not None:
|
| 3923 |
# Fetch specific head — no aggregation, no metadata lookup needed.
|
| 3924 |
attention_row = matrix_cache.get_attention_row(request_id, step, layer, head)
|
|
@@ -3934,7 +3998,9 @@ async def get_attention_row(
|
|
| 3934 |
"seq_len": len(attention_row),
|
| 3935 |
"layer": layer,
|
| 3936 |
"head": head,
|
| 3937 |
-
"aggregate_mode": None
|
|
|
|
|
|
|
| 3938 |
}
|
| 3939 |
else:
|
| 3940 |
# Aggregate across all heads. num_heads is derived inside
|
|
@@ -3954,7 +4020,9 @@ async def get_attention_row(
|
|
| 3954 |
"seq_len": len(attention_row),
|
| 3955 |
"layer": layer,
|
| 3956 |
"head": None,
|
| 3957 |
-
"aggregate_mode": aggregate_mode
|
|
|
|
|
|
|
| 3958 |
}
|
| 3959 |
|
| 3960 |
|
|
|
|
| 140 |
for older requests.
|
| 141 |
"""
|
| 142 |
with self._lock:
|
| 143 |
+
existing = self._request_meta.get(request_id, {})
|
| 144 |
+
existing.update({
|
| 145 |
"num_heads": num_heads,
|
| 146 |
"num_layers": num_layers,
|
| 147 |
"model_id": model_id,
|
| 148 |
+
})
|
| 149 |
+
self._request_meta[request_id] = existing
|
| 150 |
+
|
| 151 |
+
def set_step_query_info(self, request_id: str, step: int,
|
| 152 |
+
query_position: int, query_token_id: int):
|
| 153 |
+
"""
|
| 154 |
+
Record the predicting query position for a given generation step.
|
| 155 |
+
|
| 156 |
+
At step `s`, the model's forward pass runs over the prompt + first
|
| 157 |
+
`s` generated tokens; the logits at the LAST query position select
|
| 158 |
+
the step-`s` token. The cached attention[step][layer][head][-1] is
|
| 159 |
+
therefore the attention pattern of *that* query position, not of
|
| 160 |
+
the predicted token (which doesn't exist in the sequence yet).
|
| 161 |
+
Storing the query position and token id makes the provenance of
|
| 162 |
+
attention rows explicit when the Why this token? panel surfaces
|
| 163 |
+
them — supporting the "grounded hypotheses, not causal claims"
|
| 164 |
+
framing required by the project.
|
| 165 |
+
"""
|
| 166 |
+
with self._lock:
|
| 167 |
+
meta = self._request_meta.setdefault(request_id, {})
|
| 168 |
+
steps = meta.setdefault("steps", {})
|
| 169 |
+
steps[step] = {
|
| 170 |
+
"query_position": query_position,
|
| 171 |
+
"query_token_id": query_token_id,
|
| 172 |
}
|
| 173 |
|
| 174 |
def get_request_metadata(self, request_id: str) -> Optional[dict]:
|
|
|
|
| 177 |
meta = self._request_meta.get(request_id)
|
| 178 |
return dict(meta) if meta else None
|
| 179 |
|
| 180 |
+
def get_step_query_info(self, request_id: str, step: int) -> Optional[dict]:
|
| 181 |
+
"""Return {query_position, query_token_id} for a step, or None."""
|
| 182 |
+
with self._lock:
|
| 183 |
+
meta = self._request_meta.get(request_id)
|
| 184 |
+
if not meta:
|
| 185 |
+
return None
|
| 186 |
+
steps = meta.get("steps")
|
| 187 |
+
if not steps:
|
| 188 |
+
return None
|
| 189 |
+
entry = steps.get(step)
|
| 190 |
+
return dict(entry) if entry else None
|
| 191 |
+
|
| 192 |
def get(self, request_id: str, step: int, layer: int, head: int) -> Optional[dict]:
|
| 193 |
"""Retrieve matrix data, returning None if expired or not found."""
|
| 194 |
key = f"{request_id}:{step}:{layer}:{head}"
|
|
|
|
| 2085 |
output_hidden_states=True
|
| 2086 |
)
|
| 2087 |
|
| 2088 |
+
# Record the query position whose logits will select this
|
| 2089 |
+
# step's token. attention[layer][head][-1] is the row of
|
| 2090 |
+
# this query — exposing it in /attention/row makes the
|
| 2091 |
+
# provenance of "Why this token?" attention explicit.
|
| 2092 |
+
query_position = current_ids.shape[1] - 1
|
| 2093 |
+
query_token_id = current_ids[0, -1].item()
|
| 2094 |
+
matrix_cache.set_step_query_info(
|
| 2095 |
+
request_id, step, query_position, query_token_id
|
| 2096 |
+
)
|
| 2097 |
+
|
| 2098 |
# Get logits for next token
|
| 2099 |
raw_logits = outputs.logits[0, -1, :].clone() # Clone raw logits before any scaling
|
| 2100 |
|
|
|
|
| 2999 |
output_hidden_states=True
|
| 3000 |
)
|
| 3001 |
|
| 3002 |
+
# Record the query position whose logits will select this
|
| 3003 |
+
# step's token (see non-SSE site for rationale).
|
| 3004 |
+
query_position = current_ids.shape[1] - 1
|
| 3005 |
+
query_token_id = current_ids[0, -1].item()
|
| 3006 |
+
matrix_cache.set_step_query_info(
|
| 3007 |
+
request_id, step, query_position, query_token_id
|
| 3008 |
+
)
|
| 3009 |
+
|
| 3010 |
# Get logits for next token
|
| 3011 |
raw_logits = outputs.logits[0, -1, :].clone() # Clone raw logits before any scaling
|
| 3012 |
|
|
|
|
| 3968 |
- layer: Layer index
|
| 3969 |
- head: Head index (null if aggregated)
|
| 3970 |
- aggregate_mode: Mode used if aggregated (null otherwise)
|
| 3971 |
+
- query_position: Sequence index whose attention this row represents.
|
| 3972 |
+
The model selects the step-N token from logits at this position;
|
| 3973 |
+
the cached attention[-1] is therefore *this* position's attention,
|
| 3974 |
+
not the predicted token's self-attention.
|
| 3975 |
+
- query_token_id: Token id at query_position (last context token).
|
| 3976 |
"""
|
| 3977 |
# Aggregation parameters (num_heads in particular) come from per-request
|
| 3978 |
# metadata recorded at analysis time, not from the live model config.
|
| 3979 |
# This keeps aggregate rows correct even after the user switches models
|
| 3980 |
# mid-session — the row is averaged over the heads of the model that
|
| 3981 |
# actually produced the cached attention.
|
| 3982 |
+
query_info = matrix_cache.get_step_query_info(request_id, step)
|
| 3983 |
+
query_position = query_info.get("query_position") if query_info else None
|
| 3984 |
+
query_token_id = query_info.get("query_token_id") if query_info else None
|
| 3985 |
+
|
| 3986 |
if head is not None:
|
| 3987 |
# Fetch specific head — no aggregation, no metadata lookup needed.
|
| 3988 |
attention_row = matrix_cache.get_attention_row(request_id, step, layer, head)
|
|
|
|
| 3998 |
"seq_len": len(attention_row),
|
| 3999 |
"layer": layer,
|
| 4000 |
"head": head,
|
| 4001 |
+
"aggregate_mode": None,
|
| 4002 |
+
"query_position": query_position,
|
| 4003 |
+
"query_token_id": query_token_id,
|
| 4004 |
}
|
| 4005 |
else:
|
| 4006 |
# Aggregate across all heads. num_heads is derived inside
|
|
|
|
| 4020 |
"seq_len": len(attention_row),
|
| 4021 |
"layer": layer,
|
| 4022 |
"head": None,
|
| 4023 |
+
"aggregate_mode": aggregate_mode,
|
| 4024 |
+
"query_position": query_position,
|
| 4025 |
+
"query_token_id": query_token_id,
|
| 4026 |
}
|
| 4027 |
|
| 4028 |
|