ITADN

Host TensorDescriptor vs. device tensor descriptor creation not producing similar outputs

#8700Openvymao 创建于 2026-03-13
V
vymaocommented
**Description** `triton.tools.tensor_descriptor import TensorDescriptor` does not seem fully compatible with H200/B200, nor `triton.language.tensor_descriptor`. If I omit the tensor descriptor and instead use `tl.make_tensor_descriptor` on the device, everything works fine. But if instead I pass in `TensorDescriptor` objects, things start to fail and I get NaNs in my kernel. Furthermore, this method: ``` @triton.jit def _maybe_make_tensor_desc(desc_or_ptr, shape, strides, block_shape): if isinstance(desc_or_ptr, tl.tensor_descriptor): return desc_or_ptr else: return tl.make_tensor_descriptor(desc_or_ptr, shape, strides, block_shape) ``` does not seem to register for `TensorDescriptor`s, meaining that `isinstance(desc_or_ptr, tl.tensor_descriptor):` fails for a `TensorDescriptor` object. **Triton Information** 3.6.0 **To Reproduce** I have a host function to either create host side TensorDescriptors or device side descriptors: ``` if supports_host_descriptor(): cache_desc_tensor = k_cache if y_dim_context > 0 else q.new_empty((1, d)) value_cache_desc_tensor = v_cache if y_dim_context > 0 else q.new_empty((1, d)) desc_q = make_host_desc(q, y_dim_q, d) desc_k = make_host_desc(k, y_dim_kv_main, d) desc_v = make_host_desc(v, y_dim_kv_main, d) desc_o = make_host_desc(out, y_dim_q, d) desc_k_cache = make_host_desc(cache_desc_tensor, max(y_dim_context, 1), d) desc_v_cache = make_host_desc(value_cache_desc_tensor, max(y_dim_context, 1), d) else: desc_q = q desc_k = k desc_v = v desc_o = out desc_k_cache = k_cache desc_v_cache = v_cache ``` On the B200, I've confirmed that all tensors become TensorDescriptors on host. If `supports_host_descriptor() == False` or I block the `if` branch, so the kernel produces descriptors, everything is fine. **Expected behavior** I would expect that creating the TensorDescriptors on host would be fine.
0 条评论