fix(kernel): correct jit tensor matcher checks and pdl arch guards

[jason-fxz]

A smaller version of #133 by @zihaomu, fixed in the shared headers instead of at the index/store call sites. Supersedes #133.

@zihaomu could you test this on AMD? Only NVIDIA was tested here.

• with_dtype<Ts...>(sym) / with_device<Codes...>(sym) applied no options, so index/store accepted float indices and CPU tensors. Index and store now list kDLROCM.
• A zero dimension counted as unset, so empty inputs failed and a 0-row tensor missed a size mismatch.
• griddepcontrol is emitted only for !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 900. use_pdl=True failed to build below sm_90 and on HIP.
• From fix(rocm): make TVM-FFI index and store JIT kernels portable to HIP #133: kWarpThreads in index/store, and the launch attribute cache only on CUDA.

Tested on RTX PRO 6000 (sm_120): pytest tests/ -m "not slow" gives the same result as main, and gpt-oss-20b e2e output is identical to main.

[zihaomu]

AMD validation for 8fab7c4: the main HIP changes pass the targeted non-empty index/store checks. There is one remaining empty-batch issue, exercised by four parameter combinations. A small local follow-up patch fixes all four; it has now been tested on both runtime stacks below.

Test machine and environments

• AMD Radeon AI PRO R9700, 32 GiB, RDNA4 / gfx1201. and W7900, RDNA3 Only one GPU was exposed to each container, and the two environments were tested sequentially on that same GPU.
• Host: Ubuntu 24.04.1 LTS, kernel/amdgpu module 7.0.0-31-generic, 2 x AMD EPYC 9334, 503 GiB usable system RAM.
• ROCm 7.14 image: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.11.0; PyTorch 2.11.0+rocm7.14.0, HIP 7.14.60850, Triton 3.7.1.
• Second image: freetoken:rocm10.0.0-dev-pytorch2.11-triton3.8-ubuntu24.04-v1; PyTorch 2.11.0+rocm10.0.0, HIP 7.15.26333, Triton 3.8.0. This is a local development-image tag, not a claim about an official ROCm release.
• apache-tvm-ffi==0.1.13.post3 in both environments. FreeToken's prebuilt kernel cache was disabled, and the patched runs used fresh, separate TVM-FFI JIT cache directories.

Before / after results

These are targeted kernel/adjacent tests, not the entire repository test suite.

Environment	Unmodified PR: original acceptance set	Local patch: same acceptance set	Additional empty-input regression checks	Patched total
ROCm 7.14	90 passed, 4 failed	94 passed, 0 failed	22 passed	116 passed, 0 failed
ROCm 10 dev image	60 passed, 4 failed	64 passed, 0 failed	22 passed	86 passed, 0 failed

The original four failing combinations were index/store x int32/int64 indices, all with a zero-length batch. Each was confirmed in a fresh Python process to avoid a failed HIP launch affecting the next test. They are additional boundary tests, not four previously passing upstream tests regressing.

Coverage includes FP32/FP16/BF16, split factors 1/2/4, masked vocabulary indexing, strided store buffers and untouched padding, invalid tensor contracts, and non-default-stream graph replay. The larger 7.14 set also includes pinned/fast-index-copy, fused multi-bank copy, launch-kwargs, cache-version, and ROCm-architecture tests.

The two old #133 tests expecting an explicit PDL rejection on HIP were deselected: this PR intentionally changes that contract to ordinary stream ordering. Instead, 12 explicit use_pdl=True combinations passed eager execution and 100 graph replays per combination, with changing data/indices/slots and exact Torch-reference comparisons, on both stacks.

Why the four empty-batch cases fail

For bank.shape=(8,64), indices.shape=(0,), and output/k/v.shape=(0,64), both kernels report:

RuntimeError: Runtime check failed at .../include/freetoken/utils.cuh:119:
CUDA error: invalid configuration argument

The message says CUDA because of the shared error wrapper; this was HIP execution.

Changing the unbound-size sentinel from 0 to -1 correctly allows L=0 to bind. However, IndexKernel::run and StoreKernel::run then calculate num_blocks=0 and still call LaunchKernel. The remaining failure is therefore the zero-sized launch, not the new dtype/device constraints or PDL fallback.

Importantly, the same four inputs also fail before this PR:

• Baseline 377a3bc4284f032e9d30777ce099828a158e7e4c: Size value is not set.
• fix(rocm): make TVM-FFI index and store JIT kernels portable to HIP #133 at 5309ca655c2c7b21e09bbb91a88b94d51b941b3d: the same error.

The relevant C++ kernel sources and index/store Python wrappers are unchanged between that baseline and checked main 078132480ea4a0424f677a27eb22085624d331e6. This is an incomplete fix for an existing empty-input limitation, not evidence of four new normal-workload regressions. The new launch-time error can also surface in the next GPU operation, so avoiding the invalid launch is worthwhile.

Locally tested follow-up diff

Return after the existing tensor/row-size validation, but before constructing/launching a zero-sized grid. This preserves invalid-input rejection and does not change the non-empty path. No ROCm-specific conditional or change to the shared LaunchKernel is needed.

diff --git a/python/freetoken/kernel/csrc/jit/index.cu b/python/freetoken/kernel/csrc/jit/index.cu
--- a/python/freetoken/kernel/csrc/jit/index.cu
+++ b/python/freetoken/kernel/csrc/jit/index.cu
@@ -134,6 +134,11 @@ struct IndexKernel {
     RuntimeCheck(entry_size == element_size,
                  "IndexKernel: element_size mismatch.");
 
+    // Empty batches are valid no-ops; do not launch a zero-sized grid.
+    if (num_indices == 0) {
+      return;
+    }
+
     constexpr auto kWarpPerBlock = num_threads / device::kWarpThreads;
     const auto num_warps = num_splits * num_indices;
     const auto num_blocks = div_ceil(num_warps, kWarpPerBlock);
diff --git a/python/freetoken/kernel/csrc/jit/store.cu b/python/freetoken/kernel/csrc/jit/store.cu
--- a/python/freetoken/kernel/csrc/jit/store.cu
+++ b/python/freetoken/kernel/csrc/jit/store.cu
@@ -96,6 +96,11 @@ struct StoreKernel {
     const auto kv_cache_stride = X.unwrap() * dtype_size;
     const auto kv_input_stride = Y.unwrap() * dtype_size;
 
+    // Empty batches are valid no-ops; do not launch a zero-sized grid.
+    if (length == 0) {
+      return;
+    }
+
     const auto params = StoreKernelParams{
         .k_cache = k_cache.data_ptr(),
         .v_cache = v_cache.data_ptr(),

The 22 extra local regression cases cover:

• 12 cases: index/masked-index/store x int32/int64 x eager/graph. Cache contents stay unchanged, and a following GPU operation executes correctly. Graph tests use a non-default stream and five replays with a changing marker value.
• 8 cases: index/store must still reject CPU tensors, mixed devices, float indices, and inconsistent row counts even when one input has zero rows.
• 2 cases: directly calling a JIT module compiled for the wrong row-byte size must still fail for empty inputs.

Observed patched pytest summaries:

ROCm 7.14: 116 passed, 2 deselected in 61.37s
ROCm 10 dev: 86 passed, 2 deselected in 40.40s

Minimal reproducer for the four original failures

Save as repro_empty.py in a configured FreeToken checkout. Each invocation below should be a separate process.

import sys
import torch
from freetoken.kernel import indexing, store_cache

kernel, index_dtype = sys.argv[1:]
bank = torch.ones((8, 64), device="cuda")
other = bank.clone()
indices = torch.empty(0, dtype=getattr(torch, index_dtype), device="cuda")
empty = torch.empty((0, 64), device="cuda")
if kernel == "index":
    assert indexing(bank, indices, output=empty) is empty
else:
    store_cache(bank, other, indices, empty, empty)
torch.cuda.synchronize()
torch.testing.assert_close(bank, torch.ones_like(bank), rtol=0, atol=0)
torch.testing.assert_close(other, torch.ones_like(other), rtol=0, atol=0)
print("PASS")

FREETOKEN_DISABLE_KERNEL_CACHE=1 python repro_empty.py index int32
FREETOKEN_DISABLE_KERNEL_CACHE=1 python repro_empty.py index int64
FREETOKEN_DISABLE_KERNEL_CACHE=1 python repro_empty.py store int32
FREETOKEN_DISABLE_KERNEL_CACHE=1 python repro_empty.py store int64

Use a fresh TVM_FFI_CACHE_DIR when switching between the unmodified and patched code.

Recommendation: keep the shared-header fix, add the two local no-op guards and regression coverage. This resolves the remaining four cases on the tested AMD machine. The patch above is locally validated but has not been committed/pushed to this PR. No NVIDIA, RDNA3/CDNA, multi-GPU, full-model serving, or performance-regression claim is being made here.

[jason-fxz]

Added in 1f31fbc, thanks! Merging.