{"record":{"id":"f49f8db5bca97da6","repo":"sgl-project/sglang","slug":"replayssm-inputs-must-be-on-the-same-device","errorCode":null,"errorMessage":"ReplaySSM inputs must be on the same device.","messagePattern":"ReplaySSM inputs must be on the same device\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/helion/kda_replayssm.py","lineNumber":750,"sourceCode":"        raise ValueError(\"`force_flush` must be a length-B int32 tensor or None.\")\n\n    cache_length = d_cache.size(2)\n    if cache_length < 1:\n        raise ValueError(\"ReplaySSM cache length must be at least 1.\")\n    if d_cache.shape[1:] != (num_v_heads, cache_length, value_dim):\n        raise ValueError(\"`d_cache` must have shape [slots, HV, L, V].\")\n    if k_cache.shape[1:] != (num_q_heads, cache_length, key_dim):\n        raise ValueError(\"`k_cache` must have shape [slots, H, L, K].\")\n    if g_cache.shape[1:] != (num_v_heads, cache_length, key_dim):\n        raise ValueError(\"`g_cache` must have shape [slots, HV, L, K].\")\n    if g_cache.dtype is not torch.float32:\n        raise ValueError(\"`g_cache` must have dtype torch.float32.\")\n\n    device = mixed_qkv.device\n    if any(\n        tensor.device != device for tensor in (d_cache, k_cache, g_cache, write_pos)\n    ):\n        raise ValueError(\"ReplaySSM inputs must be on the same device.\")\n    if force_flush is not None and force_flush.device != device:\n        raise ValueError(\"`force_flush` must be on the same device as the inputs.\")\n\n    cache_block = helion.next_power_of_2(max(16, cache_length))\n    use_lower_bound = lower_bound is not None\n    kernel = _select_replayssm_decode_kernel(\n        is_bf16_state=initial_state.dtype is torch.bfloat16,\n        num_v_heads=num_v_heads,\n    )\n    result = kernel(\n        mixed_qkv,\n        flat_a,\n        b,\n        A_log,\n        flat_dt_bias,\n        scale,\n        initial_state,\n        d_cache,","sourceCodeStart":732,"sourceCodeEnd":768,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/helion/kda_replayssm.py#L732-L768","documentation":"All ReplaySSM decode inputs (d_cache, k_cache, g_cache, write_pos) must live on the same CUDA device as mixed_qkv. The wrapper compares each tensor's .device against mixed_qkv.device because the Helion/Triton kernel cannot follow pointers across devices and would silently corrupt or crash. This fires when caches are on cuda:0 but the activations are on cuda:1, or caches are left on CPU.","triggerScenarios":"Calling helion_fused_recurrent_kda_replayssm_decode in a multi-GPU (TP/EP) setup where the state pool was allocated on a different rank/device than mixed_qkv, or caches created on CPU ('cpu' device) for testing and never moved.","commonSituations":"Tensor-parallel inference where the memory pool is allocated once on the primary device; unit tests that build caches with device='cpu' by default; device migration code that moves qkv but not the state pool.","solutions":["Move all caches and write_pos to mixed_qkv.device before the call (or allocate them on mixed_qkv.device in the first place)","In TP runs, allocate per-rank state pools on torch.cuda.current_device() rather than a fixed cuda:0","Add a debug assert: assert all(t.device == mixed_qkv.device for t in (...)) in your wrapper"],"exampleFix":"# before\nout = helion_fused_recurrent_kda_replayssm_decode(mixed_qkv, d_cache_cpu, k_cache_cpu, g_cache_cpu, write_pos, ...)\n# after\ndev = mixed_qkv.device\nout = helion_fused_recurrent_kda_replayssm_decode(mixed_qkv, d_cache_cpu.to(dev), k_cache_cpu.to(dev), g_cache_cpu.to(dev), write_pos.to(dev), ...)","handlingStrategy":"validation","validationCode":"dev = mixed_qkv.device\nassert all(t.device == dev for t in (d_cache, k_cache, g_cache, write_pos))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Allocate caches on torch.cuda.current_device() in multi-GPU setups","Move CPU-built control tensors to the activation device before every call"],"tags":["helion","kda","replayssm","multi-gpu","device-placement"],"backgroundTag":"tensor-device-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}