Triton Error [CUDA]: device kernel image is invalid

[rationalspark]

Thanks for the wonderful work.

When running Mamba2, I encountered the error "Triton Error [CUDA]: device kernel image is invalid".

Should you be so kind as to provide some advice?

My environment is torch 2.3.0+cu118 triton 2.3.0 The GPU is RTX3090.

The code is

''' import torch from mamba_ssm import Mamba2 batch, length, dim = 2, 512, 256 x = torch.randn(batch, length, dim).to("cuda:2") model = Mamba2(

This module uses roughly 3 expand d_model^2 parameters

d_model=dim, # Model dimension d_model
d_state=64,  # SSM state expansion factor, typically 64 or 128
d_conv=4,    # Local convolution width
expand=2,    # Block expansion factor
headdim=128

).to("cuda:2") y = model(x) assert y.shape == x.shape '''

The error messages are

---

RuntimeError Traceback (most recent call last) Cell In[3], line 13 4 x = torch.randn(batch, length, dim).to("cuda:2") 5 model = Mamba2( 6 # This module uses roughly 3 expand d_model^2 parameters 7 d_model=dim, # Model dimension d_model (...) 11 headdim=128 12 ).to("cuda:2") ---> 13 y = model(x) 14 assert y.shape == x.shape

File ~/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, kwargs) 1530 return self._compiled_call_impl(*args, *kwargs) # type: ignore[misc] 1531 else: -> 1532 return self._call_impl(args, kwargs)

File ~/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, *kwargs) 1536 # If we don't have any hooks, we want to skip the rest of the logic in 1537 # this function, and just call forward. 1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks 1539 or _global_backward_pre_hooks or _global_backward_hooks 1540 or _global_forward_hooks or _global_forward_pre_hooks): -> 1541 return forward_call(args, **kwargs) 1543 try: 1544 result = None

File ~/work/ts/mamba/mamba_ssm/modules/mamba2.py:176, in Mamba2.forward(self, u, seqlen, seq_idx, inference_params) 174 dt_limit_kwargs = {} if self.dt_limit == (0.0, float("inf")) else dict(dt_limit=self.dt_limit) 175 if self.use_mem_eff_path and inference_params is None: --> 176 out = mamba_split_conv1d_scan_combined( 177 zxbcdt, 178 rearrange(self.conv1d.weight, "d 1 w -> d w"), 179 self.conv1d.bias, 180 self.dt_bias, 181 A, 182 D=rearrange(self.D, "(h p) -> h p", p=self.headdim) if self.D_has_hdim else self.D, 183 chunk_size=self.chunk_size, 184 seq_idx=seq_idx, 185 activation=self.activation, 186 rmsnorm_weight=self.norm.weight if self.rmsnorm else None, 187 rmsnorm_eps=self.norm.eps if self.rmsnorm else 1e-6, 188 outproj_weight=self.out_proj.weight, 189 outproj_bias=self.out_proj.bias, 190 headdim=None if self.D_has_hdim else self.headdim, 191 ngroups=self.ngroups, 192 norm_before_gate=self.norm_before_gate, 193 **dt_limit_kwargs, 194 ) 195 if seqlen_og is not None: 196 out = rearrange(out, "b l d -> (b l) d")

File ~/work/ts/mamba/mamba_ssm/ops/triton/ssd_combined.py:908, in mamba_split_conv1d_scan_combined(zxbcdt, conv1d_weight, conv1d_bias, dt_bias, A, D, chunk_size, initial_states, seq_idx, dt_limit, return_final_states, activation, rmsnorm_weight, rmsnorm_eps, outproj_weight, outproj_bias, headdim, ngroups, norm_before_gate) 889 def mamba_split_conv1d_scan_combined(zxbcdt, conv1d_weight, conv1d_bias, dt_bias, A, D, chunk_size, initial_states=None, seq_idx=None, dt_limit=(0.0, float("inf")), return_final_states=False, activation="silu", rmsnorm_weight=None, rmsnorm_eps=1e-6, outproj_weight=None, outproj_bias=None, headdim=None, ngroups=1, norm_before_gate=True): 890 """ 891 Argument: 892 zxbcdt: (batch, seqlen, 2 dim + 2 ngroups dstate + nheads) where dim == nheads headdim (...) 906 out: (batch, seqlen, dim) 907 """ --> 908 return MambaSplitConv1dScanCombinedFn.apply(zxbcdt, conv1d_weight, conv1d_bias, dt_bias, A, D, chunk_size, initial_states, seq_idx, dt_limit, return_final_states, activation, rmsnorm_weight, rmsnorm_eps, outproj_weight, outproj_bias, headdim, ngroups, norm_before_gate)

File ~/anaconda3/lib/python3.9/site-packages/torch/autograd/function.py:598, in Function.apply(cls, *args, *kwargs) 595 if not torch._C._are_functorch_transforms_active(): 596 # See NOTE: [functorch vjp and autograd interaction] 597 args = _functorch.utils.unwrap_dead_wrappers(args) --> 598 return super().apply(args, **kwargs) # type: ignore[misc] 600 if not is_setup_ctx_defined: 601 raise RuntimeError( 602 "In order to use an autograd.Function with functorch transforms " 603 "(vmap, grad, jvp, jacrev, ...), it must override the setup_context " 604 "staticmethod. For more details, please see " 605 "https://pytorch.org/docs/master/notes/extending.func.html" 606 )

File ~/anaconda3/lib/python3.9/site-packages/torch/cuda/amp/autocast_mode.py:115, in custom_fwd..decorate_fwd(*args, *kwargs) 113 if cast_inputs is None: 114 args[0]._fwd_used_autocast = torch.is_autocast_enabled() --> 115 return fwd(args, **kwargs) 116 else: 117 autocast_context = torch.is_autocast_enabled()

File ~/work/ts/mamba/mamba_ssm/ops/triton/ssd_combined.py:773, in MambaSplitConv1dScanCombinedFn.forward(ctx, zxbcdt, conv1d_weight, conv1d_bias, dt_bias, A, D, chunk_size, initial_states, seq_idx, dt_limit, return_final_states, activation, rmsnorm_weight, rmsnorm_eps, outproj_weight, outproj_bias, headdim, ngroups, norm_before_gate) 771 out = torch.cat([_swiglu_fwd(zx0), out], dim=-1) 772 else: --> 773 outx, , dt_out, dA_cumsum, states, final_states = _mamba_chunk_scan_combined_fwd(x, dt, A, B, C, chunk_size=chunk_size, D=D, z=None, dt_bias=dt_bias, initial_states=initial_states, seq_idx=seq_idx, dt_softplus=True, dt_limit=dt_limit) 774 # reshape input data into 2D tensor 775 x_rms = rearrange(out_x, "b s h p -> (b s) (h p)")

File ~/work/ts/mamba/mamba_ssm/ops/triton/ssd_combined.py:307, in _mamba_chunk_scan_combined_fwd(x, dt, A, B, C, chunk_size, D, z, dt_bias, initial_states, seq_idx, dt_softplus, dt_limit) 302 assert initial_states.shape == (batch, nheads, headdim, dstate) 303 # # (batch, nchunks, chunk_size, chunk_size) or (batch, nchunks, nheads, chunk_size, chunk_size) 304 # dA_cumsum_tmp0, dt_tmp0 = _chunk_cumsum_fwd(dt[:, :147], A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus) 305 # dA_cumsum_tmp1, dt_tmp1 = _chunk_cumsum_fwd(dt[:, 147:], A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus) 306 # dA_cumsum_tmp2, dt_tmp2 = _chunk_cumsum_fwd(dt[:, 147:256], A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus) --> 307 dA_cumsum, dt = _chunk_cumsum_fwd(dt, A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus, dt_limit=dt_limit) 308 states = _chunk_state_fwd(B, x, dt, dA_cumsum, seq_idx=seq_idx, states_in_fp32=True) 309 # states_tmp0 = _chunk_state_fwd(B[:, :147], x[:, :147], dt_tmp0, dA_cumsum_tmp0, states_in_fp32=True) 310 # states_tmp1 = _chunk_state_fwd(B[:, 147:], x[:, 147:], dt_tmp1, dA_cumsum_tmp1, states_in_fp32=True) 311 # states_tmp2 = _chunk_state_fwd(B[:, 147:256], x[:, 147:256], dt_tmp2, dA_cumsum_tmp2, states_in_fp32=True)

File ~/work/ts/mamba/mamba_ssm/ops/triton/ssd_chunk_state.py:582, in _chunk_cumsum_fwd(dt, A, chunk_size, dt_bias, dt_softplus, dt_limit) 580 grid_chunk_cs = lambda META: (batch, nchunks, triton.cdiv(nheads, META['BLOCK_SIZE_H'])) 581 with torch.cuda.device(dt.device.index): --> 582 _chunk_cumsum_fwd_kernel[grid_chunk_cs]( 583 dt, A, dt_bias, dt_out, dA_cumsum, 584 batch, seqlen, nheads, chunk_size, 585 dt_limit[0], dt_limit[1], 586 dt.stride(0), dt.stride(1), dt.stride(2), 587 A.stride(0), 588 dt_bias.stride(0) if dt_bias is not None else 0, 589 dt_out.stride(0), dt_out.stride(2), dt_out.stride(1), dt_out.stride(3), 590 dA_cumsum.stride(0), dA_cumsum.stride(2), dA_cumsum.stride(1), dA_cumsum.stride(3), 591 dt_softplus, 592 HAS_DT_BIAS=dt_bias is not None, 593 BLOCK_SIZE_CHUNK=triton.next_power_of_2(chunk_size), 594 ) 595 return dA_cumsum, dt_out

File ~/anaconda3/lib/python3.9/site-packages/triton/runtime/jit.py:167, in KernelInterface.getitem..(*args, kwargs) 161 def getitem(self, grid) -> T: 162 """ 163 A JIT function is launched with: fn[grid](*args, *kwargs). 164 Hence JITFunction.getitem returns a callable proxy that 165 memorizes the grid. 166 """ --> 167 return lambda args, kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)

File ~/anaconda3/lib/python3.9/site-packages/triton/runtime/autotuner.py:143, in Autotuner.run(self, *args, *kwargs) 141 pruned_configs = self.prune_configs(kwargs) 142 bench_start = time.time() --> 143 timings = {config: self._bench(args, config=config, **kwargs) for config in pruned_configs} 144 bench_end = time.time() 145 self.bench_time = bench_end - bench_start

File ~/anaconda3/lib/python3.9/site-packages/triton/runtime/autotuner.py:143, in (.0) 141 pruned_configs = self.prune_configs(kwargs) 142 bench_start = time.time() --> 143 timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs} 144 bench_end = time.time() 145 self.bench_time = bench_end - bench_start

File ~/anaconda3/lib/python3.9/site-packages/triton/runtime/autotuner.py:122, in Autotuner._bench(self, config, *args, **meta) 119 self.post_hook(args) 121 try: --> 122 return do_bench(kernel_call, warmup=self.warmup, rep=self.rep, quantiles=(0.5, 0.2, 0.8)) 123 except OutOfResources: 124 return [float("inf"), float("inf"), float("inf")]

File ~/anaconda3/lib/python3.9/site-packages/triton/testing.py:102, in do_bench(fn, warmup, rep, grad_to_none, quantiles, fast_flush, return_mode) 83 import torch 84 """ 85 Benchmark the runtime of the provided function. By default, return the median runtime of :code:fn along with 86 the 20-th and 80-th performance percentile. (...) 99 :type fast_flush: bool 100 """ --> 102 fn() 103 torch.cuda.synchronize() 105 # We maintain a buffer of 256 MB that we clear 106 # before each kernel call to make sure that the L2 107 # doesn't contain any input data before the run

File ~/anaconda3/lib/python3.9/site-packages/triton/runtime/autotuner.py:110, in Autotuner._bench..kernel_call() 108 config.pre_hook(full_nargs) 109 self.pre_hook(args) --> 110 self.fn.run( 111 *args, 112 num_warps=config.num_warps, 113 num_stages=config.num_stages, 114 num_ctas=config.num_ctas, 115 enable_warp_specialization=config.enable_warp_specialization, 116 # enable_persistent=False, 117 **current, 118 ) 119 self.post_hook(args)

File ~/anaconda3/lib/python3.9/site-packages/triton/runtime/jit.py:425, in JITFunction.run(self, grid, warmup, *args, *kwargs) 423 if not warmup: 424 args = [arg.value for arg in args if not arg.param.is_constexpr] --> 425 kernel.run(grid_0, grid_1, grid_2, kernel.num_warps, kernel.num_ctas, # number of warps/ctas per instance 426 kernel.cluster_dims[0], kernel.cluster_dims[1], kernel.cluster_dims[2], # cluster 427 kernel.shared, stream, kernel.function, CompiledKernel.launch_enter_hook, 428 CompiledKernel.launch_exit_hook, kernel, 429 driver.assemble_tensormap_to_arg(kernel.metadata["tensormaps_info"], args)) 430 return kernel

File ~/anaconda3/lib/python3.9/site-packages/triton/compiler/compiler.py:255, in CompiledKernel.getattribute(self, name) 253 def getattribute(self, name): 254 if name == 'run': --> 255 self._init_handles() 256 return super().getattribute(name)

File ~/anaconda3/lib/python3.9/site-packages/triton/compiler/compiler.py:250, in CompiledKernel._init_handles(self) 248 raise OutOfResources(self.shared, max_shared, "shared memory") 249 # TODO: n_regs, n_spills should be metadata generated when calling ptxas --> 250 self.module, self.function, self.n_regs, self.n_spills = driver.utils.load_binary( 251 self.name, self.kernel, self.shared, device)

RuntimeError: Triton Error [CUDA]: device kernel image is invalid

Thank you very much for all your assistance.

[tridao]

It's a triton error, idk how to fix it but you can search triton repo issues

[catalpaaa]

try to build mamba and causal conv1d yourself with pip install -e, iirc

or maybe build triton from source

[jsie7]

I encountered this before as well. It seems as if mamba/triton is using binaries for a different cuda version, hence the invalidity error. Either build from source to ensure the correct version or manually select the correct binaries.

[catalpaaa]

I encountered this before as well. It seems as if mamba/triton is using binaries for a different cuda version, hence the invalidity error. Either build from source to ensure the correct version or manually select the correct binaries.

Yes, check if your CUDA_HOME is pointing to other cuda installations

[JHChen1]

我以前也遇到过这种情况。似乎 mamba/triton 使用的是不同 cuda 版本的二进制文件，因此出现无效错误。要么从源代码构建以确保正确的版本，要么手动选择正确的二进制文件。

Hello, I am currently using torch==2.0.1; triton==2.3.0; mamba_ssm==2.0.3. I am using cuda v11.8 V100 and have this problem: it can run correctly on cuda:0, but reports an error on cuda:1: "untimeError: Triton Error [CUDA]: context is destroyed". Can you give me some advice?

[tridao]

You can try upgrading pytorch, though I don't think Triton support V100 very well in general

[rationalspark]

Thanks for all the replies. I tried to build Mamba from source, or install the latest triton 2.3.0, but the "Triton Error [CUDA]: device kernel image is invalid" still exists. I also tried to build Trition from souce, but the compilation fails.