{"record":{"id":"b62325e96d849b25","repo":"sgl-project/sglang","slug":"force-flush-must-be-on-the-same-device-as-the-in","errorCode":null,"errorMessage":"`force_flush` must be on the same device as the inputs.","messagePattern":"`force_flush` must be on the same device as the inputs\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/helion/kda_replayssm.py","lineNumber":752,"sourceCode":"    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,\n        k_cache,\n        g_cache,","sourceCodeStart":734,"sourceCodeEnd":770,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/helion/kda_replayssm.py#L734-L770","documentation":"When force_flush is supplied (a per-row flush control tensor), it must sit on the same device as the other ReplaySSM inputs (mixed_qkv's device). The wrapper checks force_flush.device != device separately from the cache checks because force_flush is optional; a CPU-resident index tensor with CUDA caches is the classic mistake. Launching with a cross-device index would fault inside the Triton kernel.","triggerScenarios":"Calling helion_fused_recurrent_kda_replayssm_decode with force_flush created via torch.zeros(B, dtype=torch.int32) on CPU (no device= argument) while all caches are on CUDA; commonly exercised by test_replayssm_per_row_flush_contract.","commonSituations":"Constructing control tensors like torch.zeros(B, dtype=torch.int32) without device=; scheduling code that builds flush flags on the scheduler (CPU) process and passes them straight to the GPU kernel wrapper.","solutions":["Create force_flush with device=mixed_qkv.device (e.g. torch.zeros(B, dtype=torch.int32, device=qkv.device))","Or call force_flush = force_flush.to(mixed_qkv.device) before invoking the op","If flushing all rows, pass None and use the kernel's implicit flush path instead of a CPU tensor of ones"],"exampleFix":"# before\nforce_flush = torch.zeros(batch, dtype=torch.int32)\n# after\nforce_flush = torch.zeros(batch, dtype=torch.int32, device=mixed_qkv.device)","handlingStrategy":"validation","validationCode":"if force_flush is not None and force_flush.device != mixed_qkv.device:\\n    force_flush = force_flush.to(mixed_qkv.device)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always pass device= when creating index/control tensors","If flush-everything semantics are wanted, pass force_flush=None instead of a CPU ones tensor"],"tags":["helion","kda","replayssm","device-placement","force-flush"],"backgroundTag":"tensor-device-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}