Fix hstu inference cudagraph capture on sm100 (#463)

* Fix hstu inference cudagraph capture on sm100

* Add fix for inference utest

* Clean up code format

Author: geoffreyQiu

.gitmodules

@@ -8,3 +8,6 @@
 [submodule "third_party/FlexKV"]
 	path = third_party/FlexKV
 	url = https://github.com/taco-project/FlexKV.git
+[submodule "third_party/nv-embedding-cache"]
+	path = third_party/nv-embedding-cache
+	url = https://github.com/NVIDIA/nv-embedding-cache.git

corelib/dynamicemb/dynamicemb/__init__.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 from .dump_load import DynamicEmbDump, DynamicEmbLoad
-from .incremental_dump import DeltaDumpResult, pop_evicted_keys
 from .dynamicemb_config import (
     BATCH_SIZE_PER_DUMP,
     DynamicEmbCheckMode,
@@ -33,6 +32,7 @@
     string_to_evict_strategy,
 )
 from .embedding_admission import FrequencyAdmissionStrategy, KVCounter
+from .incremental_dump import DeltaDumpResult, pop_evicted_keys
 from .optimizer import EmbOptimType, OptimizerArgs
 from .types import (
     BUCKET_ALIGNMENT,

corelib/dynamicemb/dynamicemb/jit/score_jit.py

@@ -94,8 +94,9 @@ def _remap_score_function(fn, perm):
     tree = ast.parse(src)
     func = tree.body[0]
     if not isinstance(func, ast.FunctionDef) or not func.args.args:
-        raise TypeError("score_function must be a plain function taking "
-                        "(scores, cur_timestamp).")
+        raise TypeError(
+            "score_function must be a plain function taking " "(scores, cur_timestamp)."
+        )
     func.decorator_list = []  # don't re-apply decorators when we recompile
     scores_name = func.args.args[0].arg
     n = len(perm)
@@ -154,9 +155,7 @@ def score_function_key(fn, perm, cc_major: int, cc_minor: int) -> int:
     return key or 1
 
 
-def register_score_function(
-    fn, score_strategy, cc_major: int, cc_minor: int
-) -> int:
+def register_score_function(fn, score_strategy, cc_major: int, cc_minor: int) -> int:
     """numba-compile fn -> LTO-IR, link into the custom cubin, cache under its
     key. Returns the key to pass as score_fn_key on inserts. Idempotent.
 

corelib/dynamicemb/dynamicemb/key_value_table.py

@@ -1113,9 +1113,7 @@ def _pop_state_evicted_keys(state: DynamicEmbTableState, table_id: int) -> torch
         # (== the table's index_type, which may be int32/uint32), so a hardcoded
         # int64 here would make callers that concat/compare across empty and
         # non-empty pops hit a dtype mismatch.
-        return torch.empty(
-            0, dtype=state.key_index_map.key_type, device=state.device
-        )
+        return torch.empty(0, dtype=state.key_index_map.key_type, device=state.device)
     keys = torch.cat(state.evicted_key_chunks)
     tids = torch.cat(state.evicted_tid_chunks)
     mask = tids == table_id

corelib/dynamicemb/setup.py

@@ -108,7 +108,9 @@ def get_version():
 # DEMB_ARCHS="80;90;100".
 DEMB_ARCHS = [
     a.strip()
-    for a in os.environ.get("DEMB_ARCHS", "75;80;90;100;120").replace(",", ";").split(";")
+    for a in os.environ.get("DEMB_ARCHS", "75;80;90;100;120")
+    .replace(",", ";")
+    .split(";")
     if a.strip()
 ]
 
@@ -188,7 +190,6 @@ def get_extensions():
     readme = f.read()
 import time
 
-
 EVICT_TU = "src/jit/evict_lrulfu.cu"
 
 
@@ -211,18 +212,31 @@ def compile_evict_fatbins():
     # compiles the same kernels.cuh / <cub/cub.cuh> / cooperative_groups headers,
     # so it needs the same relaxed-constexpr + extended-lambda support to stay
     # buildable across CUDA versions.
-    common = ["-std=c++17", "-O3", "--use_fast_math",
-              "--expt-relaxed-constexpr", "--expt-extended-lambda",
-              f"-I{root_path / 'src'}"]
+    common = [
+        "-std=c++17",
+        "-O3",
+        "--use_fast_math",
+        "--expt-relaxed-constexpr",
+        "--expt-extended-lambda",
+        f"-I{root_path / 'src'}",
+    ]
 
     variants = [
         ("LexFreqTsComparator", "evict_lrulfu_lex.fatbin", "sm"),
         ("UserFnComparator", "evict_lrulfu_custom.fatbin", "lto"),
     ]
     for comparator, out_name, code_kind in variants:
         out = str(out_dir / out_name)
-        cmd = ([nvcc, "--fatbin", *_gencode_flags(code_kind), *common,
-                f"-DDEMB_EVICT_COMPARATOR={comparator}", src, "-o", out])
+        cmd = [
+            nvcc,
+            "--fatbin",
+            *_gencode_flags(code_kind),
+            *common,
+            f"-DDEMB_EVICT_COMPARATOR={comparator}",
+            src,
+            "-o",
+            out,
+        ]
         print(f"[dynamicemb] nvcc evict fatbin ({comparator}): {' '.join(cmd)}")
         subprocess.run(cmd, check=True)
         print(f"[dynamicemb] {out_name}: {os.path.getsize(out)} bytes")

corelib/dynamicemb/test/unit_tests/table_operation/test_lru_lfu.py

@@ -450,12 +450,16 @@ def test_lru_lfu_custom_score_function_ranks_by_its_dimension(
     probes = {"OLD": (keys[:1], tids[:1]), "HF": (keys[1:2], tids[1:2])}
 
     table.lookup(
-        keys[1:], tids[1:], ScoreArg(name="frequency", value=ones[1:], policy=ScorePolicy.LRU_LFU)
+        keys[1:],
+        tids[1:],
+        ScoreArg(name="frequency", value=ones[1:], policy=ScorePolicy.LRU_LFU),
     )
     torch.cuda.synchronize()
     for _ in range(28):
         table.lookup(
-            keys[1:2], tids[1:2], ScoreArg(name="frequency", value=ones[1:2], policy=ScorePolicy.LRU_LFU)
+            keys[1:2],
+            tids[1:2],
+            ScoreArg(name="frequency", value=ones[1:2], policy=ScorePolicy.LRU_LFU),
         )
     torch.cuda.synchronize()
 
@@ -477,12 +481,18 @@ def test_lru_lfu_custom_score_function_ranks_by_its_dimension(
     ev_keys, ev_tids = probes[evicted]
     sv_keys, sv_tids = probes[survivor]
     _, ev_found, _ = table.lookup(
-        ev_keys, ev_tids, ScoreArg(name="frequency", value=None, policy=ScorePolicy.CONST)
+        ev_keys,
+        ev_tids,
+        ScoreArg(name="frequency", value=None, policy=ScorePolicy.CONST),
     )
     _, sv_found, _ = table.lookup(
-        sv_keys, sv_tids, ScoreArg(name="frequency", value=None, policy=ScorePolicy.CONST)
+        sv_keys,
+        sv_tids,
+        ScoreArg(name="frequency", value=None, policy=ScorePolicy.CONST),
     )
-    assert not torch.any(ev_found), f"{score_fn.__name__} should evict the {evicted} probe"
+    assert not torch.any(
+        ev_found
+    ), f"{score_fn.__name__} should evict the {evicted} probe"
     assert torch.all(sv_found), f"{score_fn.__name__} should keep the {survivor} probe"
 
 
@@ -506,7 +516,9 @@ def test_lru_lfu_decay_matches_python_oracle(current_device):
         one, one_tid, one_val = keys[i : i + 1], tids[i : i + 1], ones[i : i + 1]
         for _ in range(reps):
             table.lookup(
-                one, one_tid, ScoreArg(name="frequency", value=one_val, policy=ScorePolicy.LRU_LFU)
+                one,
+                one_tid,
+                ScoreArg(name="frequency", value=one_val, policy=ScorePolicy.LRU_LFU),
             )
         torch.cuda.synchronize()
 
@@ -551,8 +563,8 @@ def test_lru_lfu_score_function_logical_order_remap(current_device):
     physical frequency, so on identical tables they evict the SAME keys."""
     device = torch.cuda.current_device()
 
-    def _run(fn, strat):
-        table = _custom_table(fn, strat)
+    def _run(fn, strat):  # codespell:ignore strat
+        table = _custom_table(fn, strat)  # codespell:ignore strat
         n = 100
         keys = torch.arange(1, 1 + n, dtype=torch.int64, device=device)
         tids = torch.zeros(n, dtype=torch.int64, device=device)
@@ -561,7 +573,9 @@ def _run(fn, strat):
         # Give each key a distinct frequency (key i looked up n-1-i extra times).
         for r in range(1, n):
             table.lookup(
-                keys[:r], tids[:r], ScoreArg(name="frequency", value=ones[:r], policy=ScorePolicy.LRU_LFU)
+                keys[:r],
+                tids[:r],
+                ScoreArg(name="frequency", value=ones[:r], policy=ScorePolicy.LRU_LFU),
             )
         torch.cuda.synchronize()
         n_new = 40
@@ -758,7 +772,9 @@ def test_lru_lfu_default_evictor_timestamp_tiebreak(current_device):
     idx_parts = []
     for g in range(G):
         s = g * per
-        idx_g, _ = _insert(table, keys[s : s + per], tids[s : s + per], ones[s : s + per])
+        idx_g, _ = _insert(
+            table, keys[s : s + per], tids[s : s + per], ones[s : s + per]
+        )
         idx_parts.append(idx_g)
         torch.cuda.synchronize()
     idx = torch.cat(idx_parts)

docker/Dockerfile

@@ -256,14 +256,13 @@ RUN git clone https://github.com/jiayus-nvidia/flash-attention.git flash-attenti
 # -- Layer 4, 5: FlexKV and NVE. PYNVE_DISABLE_AVX512 only meaningful on x86;
 # leave it unset on arm64 (AVX512 doesn't exist there).
 COPY third_party/FlexKV /workspace/deps/FlexKV
+COPY third_party/nv-embedding-cache /workspace/deps/nve
 RUN cd FlexKV && mkdir -p build && cd build && \
     TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 9.0 10.0 12.0" cmake .. && \
     TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 9.0 10.0 12.0" cmake --build . -j8 && \
     cd .. && mkdir -p flexkv/lib && cp -P build/lib/*.so* flexkv/lib/ && \
     MAX_JOBS=8 FLEXKV_DEBUG=0 TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 9.0 10.0 12.0" python3 setup.py install && \
-    cd /workspace/deps && \
-    git clone -b v26.05 --recursive https://github.com/NVIDIA/nv-embedding-cache.git nve && \
-    cd nve && git submodule update --init --recursive && \
+    cd /workspace/deps/nve && \
     cd third_party/ && git clone https://github.com/NVIDIA/NVTX.git NVTX && cd .. && \
     if [ "${TARGETPLATFORM}" = "linux/arm64" ]; then AVX512_ENV=""; else AVX512_ENV="PYNVE_DISABLE_AVX512=1"; fi && \
     env TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 9.0 10.0 12.0" PYNVE_WITH_TORCH_BINDINGS=1 ${AVX512_ENV} \
@@ -351,4 +350,3 @@ RUN cd ./examples/hstu && \
       inference_aoti/triton_aoti/hstu_gr_ranking_kvcache/config.pbtxt
 
 WORKDIR /workspace/recsys-examples/
-

examples/hstu/inference/triton/hstu_export_aligned/model.py

@@ -38,8 +38,7 @@ def initialize(self, args):
 
         if self._max_batch_size < 1 or self._max_batch_size > 8:
             raise ValueError(
-                "HSTU_MAX_BATCH_SIZE must be in [1, 8], got "
-                f"{self._max_batch_size}"
+                "HSTU_MAX_BATCH_SIZE must be in [1, 8], got " f"{self._max_batch_size}"
             )
         for required_path in (hstu_root, gin_config, checkpoint_dir):
             if not required_path.exists():
@@ -82,9 +81,7 @@ def initialize(self, args):
             data_processor._item_feature_name,
             data_processor._action_feature_name,
         ]
-        self._contextual_feature_names = list(
-            data_processor._contextual_feature_names
-        )
+        self._contextual_feature_names = list(data_processor._contextual_feature_names)
         self._item_feature_name = data_processor._item_feature_name
         self._action_feature_name = data_processor._action_feature_name
         self._max_num_candidates = int(dataset_args.max_num_candidates)
@@ -195,9 +192,7 @@ def execute(self, requests):
                     device=self._device, dtype=torch.int64
                 )
 
-                batch = self._make_batch(
-                    values, lengths, num_candidates, batch_size
-                )
+                batch = self._make_batch(values, lengths, num_candidates, batch_size)
                 with torch.inference_mode():
                     logits = self._model(batch).float().cpu().numpy()
                 responses.append(

examples/hstu/inference_aoti/export_inference_gr_ranking_kvcache.py

@@ -603,9 +603,7 @@ def forward(
                 ):
                     rebuilt = self._rebuild_batch(values, lengths, num_candidates)
                     logits, auxiliary_output = _split_model_outputs(
-                        self.inner(
-                            rebuilt, user_ids.cpu(), total_history_lengths.cpu()
-                        )
+                        self.inner(rebuilt, user_ids.cpu(), total_history_lengths.cpu())
                     )
                     if auxiliary_output is not None:
                         raise RuntimeError(
@@ -748,9 +746,7 @@ def forward(
                             dump_dir, f"batch_{dump_idx:06d}_compiled_logits.pt"
                         ),
                     )
-                    print(
-                        f"    [Batch {dump_idx + 1}] Dumped C++ replay tensors"
-                    )
+                    print(f"    [Batch {dump_idx + 1}] Dumped C++ replay tensors")
                     dump_idx += 1
                     eval_module(logits.cuda(), batch.labels.values())
 

examples/hstu/inference_aoti/test_tritonserver_aoti_hstu_model.py

@@ -179,10 +179,7 @@ def _merge_samples(samples: Sequence[InputSample]) -> list[np.ndarray]:
     lengths = np.concatenate(
         [
             np.asarray(
-                [
-                    sample.feature_lengths[feature_index]
-                    for sample in samples
-                ],
+                [sample.feature_lengths[feature_index] for sample in samples],
                 dtype=samples[0].feature_lengths.dtype,
             )
             for feature_index in range(num_features)
@@ -489,9 +486,11 @@ def main() -> int:
         print(f"Slept {args.post_warmup_sleep_seconds:.3f} seconds after warmup")
 
     result = None
-    for cache_set_index, user_id_offset, cache_set_input_cases in (
-        cache_measurement_sets
-    ):
+    for (
+        cache_set_index,
+        user_id_offset,
+        cache_set_input_cases,
+    ) in cache_measurement_sets:
         print(
             f"Cache measurement set {cache_set_index}/"
             f"{CACHE_MEASUREMENT_SET_COUNT}: user_id_offset={user_id_offset}"