From 13ed18e39af637649d9d0d369712d1a0c920b42d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A4=A9=E9=82=91?= Date: Mon, 25 May 2026 14:41:11 +0800 Subject: [PATCH 1/4] [feat] gr position op: optional explicit query_time anchor for time bias add_timestamp_positional_embeddings derives the time-bias gap as `ts_gap = query_time - event_time`, where `query_time` was always the last in-sequence timestamp. That is correct for DLRM-HSTU (the candidate request time is concatenated last) but wrong for any UIH-only sequence, where the most-recent event always lands at gap 0. Add an optional per-row `query_time` tensor to the pytorch + triton kernels and the dispatcher. When passed, it anchors the gap for every position (including the last); when None, the last-timestamp gather is unchanged, so DLRM-HSTU and existing configs are byte-identical. The triton kernel gates the new `QueryTime` pointer behind a `HAS_QUERY_TIME` constexpr, mirroring the existing `NumTargets` / `HAS_MULTIPLE_TARGETS` pattern; backward is unaffected (it replays the stored time-bucket indices). Co-Authored-By: Claude Opus 4.7 (1M context) --- tzrec/ops/_pytorch/pt_position.py | 19 ++++++-- tzrec/ops/_triton/triton_position.py | 22 ++++++++- tzrec/ops/position.py | 3 ++ tzrec/ops/position_test.py | 70 ++++++++++++++++++++++++++++ 4 files changed, 108 insertions(+), 6 deletions(-) diff --git a/tzrec/ops/_pytorch/pt_position.py b/tzrec/ops/_pytorch/pt_position.py index 811fa9a4..a228b4e8 100644 --- a/tzrec/ops/_pytorch/pt_position.py +++ b/tzrec/ops/_pytorch/pt_position.py @@ -93,6 +93,7 @@ def pytorch_add_timestamp_positional_embeddings( interleave_targets: bool, time_bucket_fn: str, time_bucket_increments: float, + query_time: Optional[torch.Tensor] = None, ) -> torch.Tensor: max_pos_ind = pos_embeddings.size(0) # position encoding @@ -115,11 +116,19 @@ def pytorch_add_timestamp_positional_embeddings( max_lengths=[max_seq_len], padding_value=0.0, ).squeeze(-1) - query_time = torch.gather( - timestamps, - dim=1, - index=(seq_lengths - 1).unsqueeze(1).clamp(min=0).to(torch.int64), - ) + if query_time is None: + # No explicit anchor: use the last in-sequence timestamp. For + # DLRM-HSTU the candidate is concatenated last, so this is the + # request time; for any UIH-only sequence it is the most-recent event. + query_time = torch.gather( + timestamps, + dim=1, + index=(seq_lengths - 1).unsqueeze(1).clamp(min=0).to(torch.int64), + ) + else: + # Explicit per-row request time (HSTUMatch two-tower: no candidate is + # concatenated, so the anchor cannot be derived from the sequence). + query_time = query_time.view(-1, 1).to(timestamps.dtype) ts = query_time - timestamps ts = ts + time_delta ts = ts.clamp(min=1e-6) / time_bucket_increments diff --git a/tzrec/ops/_triton/triton_position.py b/tzrec/ops/_triton/triton_position.py index 9e632c8b..0750ea9c 100644 --- a/tzrec/ops/_triton/triton_position.py +++ b/tzrec/ops/_triton/triton_position.py @@ -287,6 +287,7 @@ def _add_timestamp_position_embeddings_kernel( TsEmb, Out, TS, + QueryTime, PosInds, TsInds, NumTargets, @@ -306,6 +307,7 @@ def _add_timestamp_position_embeddings_kernel( HAS_MULTIPLE_TARGETS: tl.constexpr, INTERLEAVE_TARGETS: tl.constexpr, TIME_BUCKET_FN: tl.constexpr, + HAS_QUERY_TIME: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_N: tl.constexpr, ): @@ -348,7 +350,12 @@ def _add_timestamp_position_embeddings_kernel( pos_emb_offsets = pos_inds[:, None] * stride_pn + offs_d[None, :] # timestamp encoding ts = tl.load(TS + seq_start + offs_n, mask=mask_n) - query_time = tl.load(TS + seq_end - 1) + if HAS_QUERY_TIME: + # Explicit per-row request time (HSTUMatch two-tower). + query_time = tl.load(QueryTime + off_b) + else: + # Last in-sequence timestamp (canonical / DLRM-HSTU request anchor). + query_time = tl.load(TS + seq_end - 1) ts = query_time - ts + time_delta ts = tl.where(ts > 1e-6, ts, 1e-6) / time_bucket_increments if TIME_BUCKET_FN == "log": @@ -468,6 +475,7 @@ def triton_add_timestamp_positional_embeddings_fwd( interleave_targets: bool, time_bucket_fn: str, time_bucket_increments: float, + query_time: Optional[torch.Tensor], ) -> Tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int ]: @@ -489,6 +497,10 @@ def triton_add_timestamp_positional_embeddings_fwd( out = torch.empty_like(seq_embeddings) timestamps = switch_to_contiguous_if_needed(timestamps) + if query_time is not None: + query_time = switch_to_contiguous_if_needed( + query_time.view(-1).to(timestamps.dtype) + ) ts_inds = torch.empty_like(seq_embeddings[:, 0], dtype=torch.int32) pos_inds = torch.empty_like(seq_embeddings[:, 0], dtype=torch.int32) ts_emb_size = ts_embeddings.shape[0] @@ -506,6 +518,7 @@ def triton_add_timestamp_positional_embeddings_fwd( TsEmb=ts_embeddings, Out=out, TS=timestamps, + QueryTime=query_time, PosInds=pos_inds, TsInds=ts_inds, NumTargets=num_targets, @@ -525,6 +538,7 @@ def triton_add_timestamp_positional_embeddings_fwd( HAS_MULTIPLE_TARGETS=num_targets is not None, INTERLEAVE_TARGETS=interleave_targets, TIME_BUCKET_FN=time_bucket_fn, + HAS_QUERY_TIME=query_time is not None, BLOCK_D=BLOCK_D, ) sorted_ts_key_inds, sorted_ts_value_inds = torch.sort(ts_inds) @@ -557,6 +571,7 @@ def forward( interleave_targets: bool, time_bucket_fn: str, time_bucket_increments: float, + query_time: Optional[torch.Tensor] = None, ): ( out, @@ -579,6 +594,7 @@ def forward( interleave_targets=interleave_targets, time_bucket_fn=time_bucket_fn, time_bucket_increments=time_bucket_increments, + query_time=query_time, ) ctx.save_for_backward( sorted_pos_key_inds, @@ -612,6 +628,7 @@ def backward( None, None, None, + None, ]: ( sorted_pos_key_inds, @@ -666,6 +683,7 @@ def backward( None, None, None, + None, ) @@ -697,6 +715,7 @@ def triton_add_timestamp_positional_embeddings( interleave_targets: bool, time_bucket_fn: str, time_bucket_increments: float, + query_time: Optional[torch.Tensor] = None, ) -> torch.Tensor: return _AddTimestampPositionEmbeddingsFunction.apply( seq_embeddings, @@ -711,4 +730,5 @@ def triton_add_timestamp_positional_embeddings( interleave_targets, time_bucket_fn, time_bucket_increments, + query_time, ) diff --git a/tzrec/ops/position.py b/tzrec/ops/position.py index 618831ef..99ef3965 100644 --- a/tzrec/ops/position.py +++ b/tzrec/ops/position.py @@ -104,6 +104,7 @@ def add_timestamp_positional_embeddings( time_bucket_fn: str = "sqrt", time_bucket_increments: float = 60.0, kernel: Kernel = Kernel.PYTORCH, + query_time: Optional[torch.Tensor] = None, ) -> torch.Tensor: assert time_bucket_fn in ["sqrt", "log"] if kernel == Kernel.CUTLASS: @@ -127,6 +128,7 @@ def add_timestamp_positional_embeddings( interleave_targets=interleave_targets, time_bucket_fn=time_bucket_fn, time_bucket_increments=time_bucket_increments, + query_time=query_time, ) else: return pytorch_add_timestamp_positional_embeddings( @@ -142,4 +144,5 @@ def add_timestamp_positional_embeddings( interleave_targets=interleave_targets, time_bucket_fn=time_bucket_fn, time_bucket_increments=time_bucket_increments, + query_time=query_time, ) diff --git a/tzrec/ops/position_test.py b/tzrec/ops/position_test.py index 1c53b8f8..aa4419da 100644 --- a/tzrec/ops/position_test.py +++ b/tzrec/ops/position_test.py @@ -388,6 +388,76 @@ def _test_add_timestamp_positional_embeddings( rtol=2e-2 if dtype != torch.float32 else None, ) + @unittest.skipIf(*gpu_unavailable) + def test_add_timestamp_positional_embeddings_explicit_query_time(self) -> None: + """Explicit per-row query_time anchors ts_gap. + + Contract checked here: + * passing the last in-sequence timestamp as ``query_time`` reproduces + the ``query_time=None`` (gather-last) path exactly, on both kernels; + * the pytorch and triton explicit paths agree; + * a later ``query_time`` (a request after the last event) shifts the + gaps and changes the output -- the bug this fixes: a UIH-only tower + could never represent staleness relative to the request time. + """ + from tzrec.ops.position import add_timestamp_positional_embeddings + + device = torch.device("cuda") + torch.manual_seed(0) + batch_size, d, max_seq_len, num_time_buckets = 4, 16, 48, 1000 + + seq_lengths = torch.randint(1, max_seq_len + 1, (batch_size,), device=device) + seq_offsets = torch.zeros(batch_size + 1, dtype=torch.int64, device=device) + seq_offsets[1:] = torch.cumsum(seq_lengths, dim=0) + total = int(seq_offsets[-1].item()) + + pos_w = torch.empty(max_seq_len, d, device=device).uniform_(-1.0, 1.0) + ts_w = torch.empty(num_time_buckets + 1, d, device=device).uniform_(-1.0, 1.0) + seq_emb = torch.empty(total, d, device=device).uniform_(-0.1, 0.1) + # strictly increasing per-row timestamps + timestamps = torch.cat( + [ + torch.cumsum( + torch.randint(1, 100, (int(n.item()),), device=device), dim=0 + ) + for n in seq_lengths + ] + ).to(torch.float32) + last_ts = timestamps[seq_offsets[1:] - 1] # [B], per-row last event + # HSTUMatch's UIH-only path passes a zeros num_targets (not None). + num_targets = torch.zeros(batch_size, dtype=torch.int64, device=device) + + def run(kernel: Kernel, query_time: object) -> torch.Tensor: + return add_timestamp_positional_embeddings( + alpha=1.0, + max_seq_len=max_seq_len, + max_contextual_seq_len=0, + position_embeddings_weight=pos_w, + timestamp_embeddings_weight=ts_w, + seq_offsets=seq_offsets, + seq_lengths=seq_lengths, + seq_embeddings=seq_emb, + timestamps=timestamps, + num_targets=num_targets, + interleave_targets=False, + time_bucket_fn="sqrt", + time_bucket_increments=60.0, + kernel=kernel, + query_time=query_time, # pyre-ignore[6] + ) + + for kernel in (Kernel.PYTORCH, Kernel.TRITON): + torch.testing.assert_close(run(kernel, None), run(kernel, last_ts)) + # explicit path agrees across kernels + torch.testing.assert_close( + run(Kernel.PYTORCH, last_ts), run(Kernel.TRITON, last_ts) + ) + # a request strictly after every event shifts the gaps -> output changes + future = last_ts + 100000.0 + self.assertFalse( + torch.allclose(run(Kernel.TRITON, None), run(Kernel.TRITON, future)) + ) + if __name__ == "__main__": unittest.main() From 5a59ad1e76c070e51c53608254f3c77d60775a75 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A4=A9=E9=82=91?= Date: Mon, 25 May 2026 14:41:18 +0800 Subject: [PATCH 2/4] [feat] HSTUMatchEncoder: thread per-row query_time through positional encoder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit HSTUPositionalEncoder forwards an optional `query_time` to the time-bias op. `_HSTUPipelineBase` reads it from `grouped_features[query_time_key]` (a `[B, 1]` raw request-time scalar) and passes it to the encoder — sourced from grouped_features rather than the shared preprocessor tuple so the ranking (ContextualInterleavePreprocessor) path is untouched. `query_time_key` defaults to "" (last-timestamp anchor, unchanged); HSTUMatchEncoder plumbs it to the base so the two-tower user side can opt into request-time anchoring. Co-Authored-By: Claude Opus 4.7 (1M context) --- tzrec/modules/gr/hstu_transducer.py | 19 +++++++++++++++++++ tzrec/modules/gr/positional_encoder.py | 5 +++++ 2 files changed, 24 insertions(+) diff --git a/tzrec/modules/gr/hstu_transducer.py b/tzrec/modules/gr/hstu_transducer.py index adaef252..b0ff9d5a 100644 --- a/tzrec/modules/gr/hstu_transducer.py +++ b/tzrec/modules/gr/hstu_transducer.py @@ -71,8 +71,13 @@ def __init__( attn_truncation_split_layer: int = 0, attn_truncation_tail_len: int = 0, name: str = "", + query_time_key: str = "", ) -> None: super().__init__(is_inference=is_inference) + # Grouped-feature key of the per-row request time used as the time-bias + # anchor. Empty -> anchor on the last in-sequence timestamp (canonical + # HSTU / DLRM-HSTU, which concatenates the candidate request time). + self._query_time_key: str = query_time_key self._input_preprocessor: InputPreprocessor = create_input_preprocessor( input_preprocessor, uih_embedding_dim=uih_embedding_dim, @@ -129,6 +134,13 @@ def _preprocess( output_num_targets, ) = self._input_preprocessor(grouped_features) + # Per-row request time anchor (HSTUMatch). Read from grouped_features + # rather than the preprocessor tuple so the shared ranking path is + # untouched. `[B, 1]` raw values -> the op reshapes to `[B]`. + query_time: Optional[torch.Tensor] = None + if self._query_time_key != "": + query_time = grouped_features[self._query_time_key] + with record_function("hstu_positional_encoder"): if self._positional_encoder is not None: output_seq_embeddings = self._positional_encoder( @@ -138,6 +150,7 @@ def _preprocess( seq_timestamps=output_seq_timestamps, seq_embeddings=output_seq_embeddings, num_targets=output_num_targets, + query_time=query_time, ) output_seq_embeddings = torch.nn.functional.dropout( @@ -468,6 +481,10 @@ class HSTUMatchEncoder(_HSTUPipelineBase): is_inference (bool): whether to run in inference mode. attn_truncation_split_layer (int): see `HSTUTransducer`. attn_truncation_tail_len (int): see `HSTUTransducer`. + query_time_key (str): grouped-feature key of the per-row request time + used as the time-bias anchor. Empty (default) anchors on the last + UIH timestamp; pass a scalar request-time group to anchor on the + actual request time (decoupled from UIH staleness). """ def __init__( @@ -487,6 +504,7 @@ def __init__( attn_truncation_split_layer: int = 0, attn_truncation_tail_len: int = 0, name: str = "", + query_time_key: str = "", ) -> None: super().__init__( uih_embedding_dim=uih_embedding_dim, @@ -504,6 +522,7 @@ def __init__( attn_truncation_split_layer=attn_truncation_split_layer, attn_truncation_tail_len=attn_truncation_tail_len, name=name, + query_time_key=query_time_key, ) self._output_postprocessor: OutputPostprocessor = create_output_postprocessor( output_postprocessor, embedding_dim=stu["embedding_dim"] diff --git a/tzrec/modules/gr/positional_encoder.py b/tzrec/modules/gr/positional_encoder.py index 4f3d4021..7cb9094f 100644 --- a/tzrec/modules/gr/positional_encoder.py +++ b/tzrec/modules/gr/positional_encoder.py @@ -76,6 +76,7 @@ def forward( seq_timestamps: torch.Tensor, seq_embeddings: torch.Tensor, num_targets: Optional[torch.Tensor], + query_time: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Forward the module. @@ -86,6 +87,9 @@ def forward( seq_timestamps (torch.Tensor): input sequence timestamps. seq_embeddings (torch.Tensor): input sequence embeddings. num_targets (int): number of targets. + query_time (torch.Tensor, optional): per-row request time used as + the time-bias anchor (``ts_gap = query_time - timestamp``). + When ``None``, the last in-sequence timestamp is used. Returns: torch.Tensor: output sequence embedding with position embedding. @@ -106,6 +110,7 @@ def forward( time_bucket_fn=self._time_bucket_fn, time_bucket_increments=self._time_bucket_increments, kernel=self.kernel(), + query_time=query_time, ) else: seq_embeddings = add_positional_embeddings( From 58c11ed1294867a665afb785fd18020be4a30586 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A4=A9=E9=82=91?= Date: Mon, 25 May 2026 14:41:26 +0800 Subject: [PATCH 3/4] [feat] HSTUMatch: query_time group -> request-time-anchored time bias MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit HSTUUserTower detects an optional `query_time` DEEP group (a single per-row request-time raw feature) among its feature groups and threads the key to HSTUMatchEncoder. With it, the HSTU time bias anchors on the request time instead of the last UIH event, so the user embedding reflects how stale the history is relative to the (serving) request — the signal time encoding is meant to provide, and which the exported two-tower user encoder otherwise had no way to receive. Absent the group, behavior is unchanged. No proto change: the group is routed by name, like contextual / uih_action / uih_timestamp. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/source/models/hstu_match.md | 1 + tzrec/models/hstu.py | 22 +++++++++++ tzrec/models/hstu_test.py | 64 ++++++++++++++++++++++++++++++-- 3 files changed, 84 insertions(+), 3 deletions(-) diff --git a/docs/source/models/hstu_match.md b/docs/source/models/hstu_match.md index 16dc46f9..ba6fc310 100644 --- a/docs/source/models/hstu_match.md +++ b/docs/source/models/hstu_match.md @@ -221,6 +221,7 @@ model_config { - uih_action: 用户历史交互的行为事件序列，注: 该行为事件按位存储，如 expr, click, add, buy 三个行为，则一般 expr=0, click=1, add=2, buy=4；类型为 JAGGED_SEQUENCE，当 `uih_preprocessor.action_encoder` 配置时必填 - uih_watchtime: 用户历史交互的行为时长序列；类型为 JAGGED_SEQUENCE，当 action encoder 需要 watchtime 时必填 - uih_timestamp: 用户历史交互的行为时间戳序列；类型为 JAGGED_SEQUENCE，当 `positional_encoder.use_time_encoding=true` 时必填 + - query_time: 每行一个标量的**请求时间** raw 特征；类型为 DEEP，可选。配置后，时间编码的相对时间差以请求时间为基准 (`ts_gap = query_time - 行为时间戳`)，而非默认的最后一个 UIH 行为时间。这样 user embedding 能反映行为序列相对于 (在线服务) 请求时刻的新鲜度；不配置时回退到以最后一个 UIH 行为时间为基准。请求时间需与 uih_timestamp 使用相同单位，在线服务时传入当前请求时刻 **group_name 不能变**，user_tower/item_tower 通过 group_name 索引对应的 feature_group diff --git a/tzrec/models/hstu.py b/tzrec/models/hstu.py index 2e3c6acd..27d117c7 100644 --- a/tzrec/models/hstu.py +++ b/tzrec/models/hstu.py @@ -96,6 +96,21 @@ def __init__( contextual_feature_dim = contextual_dims[0] max_contextual_seq_len = len(contextual_dims) + # Optional per-row request time anchor for the HSTU time bias. When a + # `query_time` DEEP group (a single scalar request-time raw feature) is + # present, the encoder anchors `ts_gap = query_time - event_time` on it + # instead of the last UIH timestamp — decoupling the time bias from how + # stale the history is relative to the (serving) request. Absent (the + # default), behavior is unchanged. + query_time_key = next( + ( + feature_group.group_name + for feature_group in feature_groups + if feature_group.group_name == "query_time" + ), + "", + ) + self._hstu_encoder: HSTUMatchEncoder = HSTUMatchEncoder( uih_embedding_dim=embedding_group.group_total_dim( f"{tower_config.input}.sequence" @@ -105,6 +120,7 @@ def __init__( contextual_group_name=contextual_group_name, scaling_seqlen=tower_config.max_seq_len, is_inference=False, + query_time_key=query_time_key, **config_to_kwargs(tower_config.hstu), ) if self._output_dim > 0: @@ -266,6 +282,12 @@ class HSTUMatch(MatchModel): UIHPreprocessor's action_encoder and the HSTU positional encoder's time bias. Required when `uih_preprocessor.action_encoder` is configured. + - "query_time" (optional, DEEP): a single per-row scalar request-time + raw feature. When present (and `positional_encoder.use_time_encoding` + is on), the time bias anchors `ts_gap = query_time - event_time` on + this request time instead of the last UIH timestamp — so the bias + reflects history staleness relative to the (serving) request. Absent, + the anchor falls back to the last UIH event. User tower returns the last-position UIH embedding per user; it is compared against candidate embeddings via the configured similarity at both train and diff --git a/tzrec/models/hstu_test.py b/tzrec/models/hstu_test.py index 4f6c6f57..3633c849 100644 --- a/tzrec/models/hstu_test.py +++ b/tzrec/models/hstu_test.py @@ -14,7 +14,7 @@ import torch from hypothesis import Verbosity, assume, given, settings from hypothesis import strategies as st -from torchrec import JaggedTensor, KeyedJaggedTensor +from torchrec import JaggedTensor, KeyedJaggedTensor, KeyedTensor from tzrec.datasets.utils import BASE_DATA_GROUP, CAND_POS_LENGTHS, Batch from tzrec.features.feature import create_features @@ -35,7 +35,7 @@ from tzrec.utils.test_util import TestGraphType, create_test_model, gpu_unavailable -def _build_model(device: torch.device) -> HSTUMatch: +def _build_model(device: torch.device, with_query_time: bool = False) -> HSTUMatch: """Build an HSTUMatch model with standard test configuration. Mirrors the production grouped-sequence pattern: `uih_seq` and @@ -43,6 +43,10 @@ def _build_model(device: torch.device) -> HSTUMatch: dim / `embedding_name` so the two flattened features share one embedding table. `uih_seq` also carries the `historical_ts` raw sub-feature for the timestamp dense path. + + When ``with_query_time`` is set, the user tower turns on time encoding + and a scalar ``request_time`` raw feature is exposed through a + ``query_time`` DEEP group — the per-row time-bias anchor. """ feature_cfgs = [ feature_pb2.FeatureConfig( @@ -84,6 +88,12 @@ def _build_model(device: torch.device) -> HSTUMatch: ) ), ] + if with_query_time: + feature_cfgs.append( + feature_pb2.FeatureConfig( + raw_feature=feature_pb2.RawFeature(feature_name="request_time") + ) + ) features = create_features(feature_cfgs) feature_groups = [ model_pb2.FeatureGroupConfig( @@ -102,6 +112,14 @@ def _build_model(device: torch.device) -> HSTUMatch: group_type=model_pb2.FeatureGroupType.JAGGED_SEQUENCE, ), ] + if with_query_time: + feature_groups.append( + model_pb2.FeatureGroupConfig( + group_name="query_time", + feature_names=["request_time"], + group_type=model_pb2.FeatureGroupType.DEEP, + ) + ) model_config = model_pb2.ModelConfig( feature_groups=feature_groups, hstu_match=match_model_pb2.HSTUMatch( @@ -120,6 +138,8 @@ def _build_model(device: torch.device) -> HSTUMatch: attn_num_layers=2, positional_encoder=module_pb2.GRPositionalEncoder( num_position_buckets=512, + num_time_buckets=512, + use_time_encoding=with_query_time, ), input_preprocessor=module_pb2.GRInputPreprocessor( uih_preprocessor=module_pb2.GRUIHPreprocessor(), @@ -153,13 +173,16 @@ def _build_model(device: torch.device) -> HSTUMatch: return hstu -def _build_batch(device: torch.device) -> Batch: +def _build_batch(device: torch.device, with_query_time: bool = False) -> Batch: """Build a test Batch with the row-(B-1) suffix candidate layout. UIH: user1 has 3 items, user2 has 4 items. Candidates: row 0 = [pos_0]; row 1 (last) = [pos_1, simple_neg_0, simple_neg_1] -- the shared simple-neg pool sits in the last row's suffix. pos_lengths = [1, 1]. + + When ``with_query_time`` is set, a per-row ``request_time`` dense scalar + (strictly after each user's last UIH event at ts 3 / 7) is included. """ sparse_feature = KeyedJaggedTensor.from_lengths_sync( keys=["uih_seq__video_id", "cand_seq__video_id"], @@ -172,7 +195,14 @@ def _build_batch(device: torch.device) -> Batch: lengths=torch.tensor([3, 4]), ), } + dense_features = {} + if with_query_time: + dense_features[BASE_DATA_GROUP] = KeyedTensor.from_tensor_list( + keys=["request_time"], + tensors=[torch.tensor([[100.0], [100.0]])], + ) return Batch( + dense_features=dense_features, sparse_features={BASE_DATA_GROUP: sparse_feature}, sequence_dense_features=sequence_dense_features, jagged_labels={ @@ -245,6 +275,34 @@ def test_hstu_match(self, graph_type, kernel, device_str) -> None: self.assertEqual(scalar_feature_groups[0].feature_names, ["video_id"]) self.assertEqual(scalar_feature_groups[0].group_name, "candidate") + @given( + graph_type=st.sampled_from([TestGraphType.NORMAL, TestGraphType.FX_TRACE]), + ) + @settings(verbosity=Verbosity.verbose, max_examples=2, deadline=None) + def test_hstu_match_query_time(self, graph_type) -> None: + """The optional `query_time` DEEP group is the time-bias anchor. + + Detected and threaded into the user tower so it reads a per-row + request time instead of deriving it from the last UIH event. + """ + device = torch.device("cpu") + # No query_time group -> encoder falls back to the last UIH timestamp. + self.assertEqual( + _build_model(device).user_tower._hstu_encoder._query_time_key, "" + ) + # With the group -> encoder reads it as the explicit per-row anchor. + hstu = _build_model(device, with_query_time=True) + self.assertEqual(hstu.user_tower._hstu_encoder._query_time_key, "query_time") + hstu.set_kernel(Kernel.PYTORCH) + batch = _build_batch(device, with_query_time=True) + + if graph_type == TestGraphType.FX_TRACE: + predictions = create_test_model(hstu, graph_type)(batch) + else: + wrapped = TrainWrapper(hstu, device=device).to(device) + _, (_, predictions, _) = wrapped(batch) + self.assertEqual(predictions["similarity"].size(0), 2) + if __name__ == "__main__": unittest.main() From 3515190e58ccaa93534db3b0164e1a7bfb32c551 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A4=A9=E9=82=91?= Date: Tue, 26 May 2026 16:05:22 +0800 Subject: [PATCH 4/4] [test] HSTUMatch: exercise query_time end-to-end on kuairand fixtures Turn on request-time anchoring in the integration path. The kuairand-1k-match train/eval parquets now carry a per-row `request_time` (the retrieval split boundary: first post-history event, >= max uih timestamp), surfaced via a `request_time` raw feature + `query_time` DEEP group in the config and the doc example; ci_data.sh points at the regenerated fixtures. Fold the standalone query_time unit test into test_hstu_match -- the model is always built with the query_time group and asserts the threaded key, across NORMAL/FX_TRACE/JIT_SCRIPT x PYTORCH/TRITON x cpu/cuda. Trim the inline/docstring comments to one line. Verified on A10: match_integration_test.test_hstu_with_fg_train_eval green (train + eval + AOT export of the request_time-bearing user tower + item/user predict); hstu_test + position_test green. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/source/models/hstu_match.md | 13 +++- scripts/ci/ci_data.sh | 4 +- tzrec/models/hstu.py | 15 ++--- tzrec/models/hstu_test.py | 75 +++++++-------------- tzrec/tests/configs/hstu_kuairand_1k.config | 11 +++ 5 files changed, 53 insertions(+), 65 deletions(-) diff --git a/docs/source/models/hstu_match.md b/docs/source/models/hstu_match.md index ba6fc310..79d8a1ae 100644 --- a/docs/source/models/hstu_match.md +++ b/docs/source/models/hstu_match.md @@ -104,6 +104,12 @@ feature_configs { } } } +feature_configs { + raw_feature { + feature_name: "request_time" + expression: "user:request_time" + } +} model_config { feature_groups { group_name: "contextual" @@ -140,6 +146,11 @@ model_config { feature_names: "uih_seq__action_timestamp" group_type: JAGGED_SEQUENCE } + feature_groups { + group_name: "query_time" + feature_names: "request_time" + group_type: DEEP + } hstu_match { user_tower { input: "uih" @@ -221,7 +232,7 @@ model_config { - uih_action: 用户历史交互的行为事件序列，注: 该行为事件按位存储，如 expr, click, add, buy 三个行为，则一般 expr=0, click=1, add=2, buy=4；类型为 JAGGED_SEQUENCE，当 `uih_preprocessor.action_encoder` 配置时必填 - uih_watchtime: 用户历史交互的行为时长序列；类型为 JAGGED_SEQUENCE，当 action encoder 需要 watchtime 时必填 - uih_timestamp: 用户历史交互的行为时间戳序列；类型为 JAGGED_SEQUENCE，当 `positional_encoder.use_time_encoding=true` 时必填 - - query_time: 每行一个标量的**请求时间** raw 特征；类型为 DEEP，可选。配置后，时间编码的相对时间差以请求时间为基准 (`ts_gap = query_time - 行为时间戳`)，而非默认的最后一个 UIH 行为时间。这样 user embedding 能反映行为序列相对于 (在线服务) 请求时刻的新鲜度；不配置时回退到以最后一个 UIH 行为时间为基准。请求时间需与 uih_timestamp 使用相同单位，在线服务时传入当前请求时刻 + - query_time: 每行一个标量的请求时间 raw 特征 (需与 uih_timestamp 同单位)；类型为 DEEP，可选。配置后时间编码以请求时间为基准 (`ts_gap = query_time - 行为时间戳`)，否则回退到最后一个 UIH 行为时间 **group_name 不能变**，user_tower/item_tower 通过 group_name 索引对应的 feature_group diff --git a/scripts/ci/ci_data.sh b/scripts/ci/ci_data.sh index bb9bff66..c90ca71b 100644 --- a/scripts/ci/ci_data.sh +++ b/scripts/ci/ci_data.sh @@ -10,7 +10,7 @@ wget https://tzrec.oss-accelerate.aliyuncs.com/data/test/kuairand-1k-rtp-eval-c4 wget https://tzrec.oss-accelerate.aliyuncs.com/data/test/kuairand-mot-1k-train-c4096-s100-e28061f3c88f543b9e18f40be6ddb94d.parquet -O data/test/kuairand-mot-1k-train-c4096-s100.parquet wget https://tzrec.oss-accelerate.aliyuncs.com/data/test/kuairand-mot-1k-eval-c4096-s100-f185f38e3b4a49cb791d2e4302087a1f.parquet -O data/test/kuairand-mot-1k-eval-c4096-s100.parquet # kuairand-1k-match (HSTUMatch integration test fixtures) -wget https://tzrec.oss-accelerate.aliyuncs.com/data/test/kuairand-1k-match-train-c4096-s100-f1892eabc70ae3407afe9ff5bca8cb5f.parquet -O data/test/kuairand-1k-match-train-c4096-s100.parquet -wget https://tzrec.oss-accelerate.aliyuncs.com/data/test/kuairand-1k-match-eval-c4096-s100-e4ca5e15d157efa723041cd05c127228.parquet -O data/test/kuairand-1k-match-eval-c4096-s100.parquet +wget https://tzrec.oss-accelerate.aliyuncs.com/data/test/kuairand-1k-match-train-c4096-s100-aa77964ed7f50ca30645f8dd08dbf10d.parquet -O data/test/kuairand-1k-match-train-c4096-s100.parquet +wget https://tzrec.oss-accelerate.aliyuncs.com/data/test/kuairand-1k-match-eval-c4096-s100-8678a3ac699fb08f0602f4c06cef2edf.parquet -O data/test/kuairand-1k-match-eval-c4096-s100.parquet wget https://tzrec.oss-accelerate.aliyuncs.com/data/test/kuairand-1k-match-item-gl-3d459148303acd9f838da108efcc40e5.txt -O data/test/kuairand-1k-match-item-gl.txt wget https://tzrec.oss-accelerate.aliyuncs.com/data/test/kuairand-1k-match-item-c1-8dcadabdc3e9049ed9c2250565b4b134.parquet -O data/test/kuairand-1k-match-item-c1.parquet diff --git a/tzrec/models/hstu.py b/tzrec/models/hstu.py index 27d117c7..f84d60cb 100644 --- a/tzrec/models/hstu.py +++ b/tzrec/models/hstu.py @@ -96,12 +96,8 @@ def __init__( contextual_feature_dim = contextual_dims[0] max_contextual_seq_len = len(contextual_dims) - # Optional per-row request time anchor for the HSTU time bias. When a - # `query_time` DEEP group (a single scalar request-time raw feature) is - # present, the encoder anchors `ts_gap = query_time - event_time` on it - # instead of the last UIH timestamp — decoupling the time bias from how - # stale the history is relative to the (serving) request. Absent (the - # default), behavior is unchanged. + # Optional `query_time` DEEP group: per-row request-time anchor for the + # HSTU time bias (absent -> anchor on the last UIH timestamp). query_time_key = next( ( feature_group.group_name @@ -283,11 +279,8 @@ class HSTUMatch(MatchModel): encoder's time bias. Required when `uih_preprocessor.action_encoder` is configured. - "query_time" (optional, DEEP): a single per-row scalar request-time - raw feature. When present (and `positional_encoder.use_time_encoding` - is on), the time bias anchors `ts_gap = query_time - event_time` on - this request time instead of the last UIH timestamp — so the bias - reflects history staleness relative to the (serving) request. Absent, - the anchor falls back to the last UIH event. + raw feature used as the HSTU time-bias anchor; absent, the anchor + falls back to the last UIH timestamp. User tower returns the last-position UIH embedding per user; it is compared against candidate embeddings via the configured similarity at both train and diff --git a/tzrec/models/hstu_test.py b/tzrec/models/hstu_test.py index 3633c849..b8784afe 100644 --- a/tzrec/models/hstu_test.py +++ b/tzrec/models/hstu_test.py @@ -35,7 +35,7 @@ from tzrec.utils.test_util import TestGraphType, create_test_model, gpu_unavailable -def _build_model(device: torch.device, with_query_time: bool = False) -> HSTUMatch: +def _build_model(device: torch.device) -> HSTUMatch: """Build an HSTUMatch model with standard test configuration. Mirrors the production grouped-sequence pattern: `uih_seq` and @@ -44,9 +44,9 @@ def _build_model(device: torch.device, with_query_time: bool = False) -> HSTUMat embedding table. `uih_seq` also carries the `historical_ts` raw sub-feature for the timestamp dense path. - When ``with_query_time`` is set, the user tower turns on time encoding - and a scalar ``request_time`` raw feature is exposed through a - ``query_time`` DEEP group — the per-row time-bias anchor. + Time encoding is on, with a scalar ``request_time`` raw feature exposed + through a ``query_time`` DEEP group — the per-row time-bias anchor + (mirrors the production config). """ feature_cfgs = [ feature_pb2.FeatureConfig( @@ -88,12 +88,11 @@ def _build_model(device: torch.device, with_query_time: bool = False) -> HSTUMat ) ), ] - if with_query_time: - feature_cfgs.append( - feature_pb2.FeatureConfig( - raw_feature=feature_pb2.RawFeature(feature_name="request_time") - ) + feature_cfgs.append( + feature_pb2.FeatureConfig( + raw_feature=feature_pb2.RawFeature(feature_name="request_time") ) + ) features = create_features(feature_cfgs) feature_groups = [ model_pb2.FeatureGroupConfig( @@ -112,14 +111,13 @@ def _build_model(device: torch.device, with_query_time: bool = False) -> HSTUMat group_type=model_pb2.FeatureGroupType.JAGGED_SEQUENCE, ), ] - if with_query_time: - feature_groups.append( - model_pb2.FeatureGroupConfig( - group_name="query_time", - feature_names=["request_time"], - group_type=model_pb2.FeatureGroupType.DEEP, - ) + feature_groups.append( + model_pb2.FeatureGroupConfig( + group_name="query_time", + feature_names=["request_time"], + group_type=model_pb2.FeatureGroupType.DEEP, ) + ) model_config = model_pb2.ModelConfig( feature_groups=feature_groups, hstu_match=match_model_pb2.HSTUMatch( @@ -139,7 +137,7 @@ def _build_model(device: torch.device, with_query_time: bool = False) -> HSTUMat positional_encoder=module_pb2.GRPositionalEncoder( num_position_buckets=512, num_time_buckets=512, - use_time_encoding=with_query_time, + use_time_encoding=True, ), input_preprocessor=module_pb2.GRInputPreprocessor( uih_preprocessor=module_pb2.GRUIHPreprocessor(), @@ -173,7 +171,7 @@ def _build_model(device: torch.device, with_query_time: bool = False) -> HSTUMat return hstu -def _build_batch(device: torch.device, with_query_time: bool = False) -> Batch: +def _build_batch(device: torch.device) -> Batch: """Build a test Batch with the row-(B-1) suffix candidate layout. UIH: user1 has 3 items, user2 has 4 items. @@ -181,8 +179,8 @@ def _build_batch(device: torch.device, with_query_time: bool = False) -> Batch: simple_neg_1] -- the shared simple-neg pool sits in the last row's suffix. pos_lengths = [1, 1]. - When ``with_query_time`` is set, a per-row ``request_time`` dense scalar - (strictly after each user's last UIH event at ts 3 / 7) is included. + A per-row ``request_time`` dense scalar (strictly after each user's last + UIH event at ts 3 / 7) is included as the time-bias anchor. """ sparse_feature = KeyedJaggedTensor.from_lengths_sync( keys=["uih_seq__video_id", "cand_seq__video_id"], @@ -195,12 +193,12 @@ def _build_batch(device: torch.device, with_query_time: bool = False) -> Batch: lengths=torch.tensor([3, 4]), ), } - dense_features = {} - if with_query_time: - dense_features[BASE_DATA_GROUP] = KeyedTensor.from_tensor_list( + dense_features = { + BASE_DATA_GROUP: KeyedTensor.from_tensor_list( keys=["request_time"], tensors=[torch.tensor([[100.0], [100.0]])], ) + } return Batch( dense_features=dense_features, sparse_features={BASE_DATA_GROUP: sparse_feature}, @@ -247,6 +245,9 @@ def test_hstu_match(self, graph_type, kernel, device_str) -> None: device = torch.device(device_str) hstu = _build_model(device=device) + # The query_time DEEP group is detected and threaded as the per-row + # time-bias anchor (request-time anchoring, not the last UIH event). + self.assertEqual(hstu.user_tower._hstu_encoder._query_time_key, "query_time") hstu.set_kernel(kernel) batch = _build_batch(device=device) @@ -275,34 +276,6 @@ def test_hstu_match(self, graph_type, kernel, device_str) -> None: self.assertEqual(scalar_feature_groups[0].feature_names, ["video_id"]) self.assertEqual(scalar_feature_groups[0].group_name, "candidate") - @given( - graph_type=st.sampled_from([TestGraphType.NORMAL, TestGraphType.FX_TRACE]), - ) - @settings(verbosity=Verbosity.verbose, max_examples=2, deadline=None) - def test_hstu_match_query_time(self, graph_type) -> None: - """The optional `query_time` DEEP group is the time-bias anchor. - - Detected and threaded into the user tower so it reads a per-row - request time instead of deriving it from the last UIH event. - """ - device = torch.device("cpu") - # No query_time group -> encoder falls back to the last UIH timestamp. - self.assertEqual( - _build_model(device).user_tower._hstu_encoder._query_time_key, "" - ) - # With the group -> encoder reads it as the explicit per-row anchor. - hstu = _build_model(device, with_query_time=True) - self.assertEqual(hstu.user_tower._hstu_encoder._query_time_key, "query_time") - hstu.set_kernel(Kernel.PYTORCH) - batch = _build_batch(device, with_query_time=True) - - if graph_type == TestGraphType.FX_TRACE: - predictions = create_test_model(hstu, graph_type)(batch) - else: - wrapped = TrainWrapper(hstu, device=device).to(device) - _, (_, predictions, _) = wrapped(batch) - self.assertEqual(predictions["similarity"].size(0), 2) - if __name__ == "__main__": unittest.main() diff --git a/tzrec/tests/configs/hstu_kuairand_1k.config b/tzrec/tests/configs/hstu_kuairand_1k.config index 89d99987..c6eaf0fd 100644 --- a/tzrec/tests/configs/hstu_kuairand_1k.config +++ b/tzrec/tests/configs/hstu_kuairand_1k.config @@ -136,6 +136,12 @@ feature_configs { } } } +feature_configs { + raw_feature { + feature_name: "request_time" + expression: "user:request_time" + } +} model_config { feature_groups { @@ -173,6 +179,11 @@ model_config { feature_names: "uih_seq__action_timestamp" group_type: JAGGED_SEQUENCE } + feature_groups { + group_name: "query_time" + feature_names: "request_time" + group_type: DEEP + } hstu_match { user_tower { input: "uih"