fix(kernel): encode tensor maps with the typed op for cuda_host builds - #3
Open
spectrometerHBH wants to merge 1 commit into
Open
spectrometerHBH wants to merge 1 commit into
spectrometerHBH wants to merge 1 commit into
Conversation
Add txl.cu_tensor_map_encode_tiled, which keeps the cuTensorMapEncodeTiled operand order of the existing packed calls but emits TVM's typed tensormap_encode_tiled op. CUDA-host code generation encodes that op directly and rejects the packed runtime encoder; ordinary hosts lower it back to the same packed call. Switch the fp16_bf16_gemm and mla_dsv4_multishape host preludes to the helper.
tqchen
pushed a commit
to apache/tvm
that referenced
this pull request
Sep 29, 2026
…ls in CUDA-host bundles (#20490) A CUDA-host bundle (`tvm.backend.cuda.export_cuda_host`, #20395) compiles device and host code as one NVCC C++ translation unit. Building the native kernels of mlc-ai/TIRx-kernels this way exposed several patterns that the bundle could not compile or launch: - **Low-precision types in the host pass.** The CUDA device header included `cuda_fp16.h`, `cuda_bf16.h`, `cuda_fp8.h`, `cuda_fp6.h` and `cuda_fp4.h` (and defined the `fp8_e4_t`-style aliases) only under `defined(__CUDA_ARCH__)`. NVCC's host pass then fails on kernel signatures that use `half` or `nv_bfloat16`. The guards now also admit the host pass (`!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= N`); NVRTC and device-only NVCC compilation always define `__CUDA_ARCH__` and are unchanged. - **Tensor-map parameters.** The host wrapper bound a `T.TensorMap()` parameter as `CUtensorMap* x = ((void*)...)`, which C accepts and C++ rejects. It now casts explicitly. - **Argument types at the launch.** Host C code spells some device types differently (bfloat16 is `uint16_t*` on the host, `nv_bfloat16*` in the kernel), so the direct `<<<>>>` launch did not type-check. Launches now go through a small helper that converts each argument to the kernel's own parameter type and calls `cudaLaunchKernelEx`. - **Launch attributes.** `clusterCtaIdx.*`, `preferredClusterCtaIdx.*`, programmatic dependent launch and cooperative launch were rejected. They now become the corresponding `cudaLaunchAttribute`s, following `CUDAWrappedFunc` in the CUDA runtime module (including `cudaFuncAttributeNonPortableClusterSizeAllowed` for cluster launches and omitting a unit preferred cluster). Required block dimensions remain unsupported and are still diagnosed. ### Testing On B200 (sm_100a), CUDA 13.2, this branch: - `tests/python/codegen/test_target_codegen_cuda.py -k cuda_host`: 8 passed (4 tests, each under NVCC and NVRTC device compilation). New tests: - bfloat16 buffers with a two-CTA cluster launch: runs and checks both the values and each CTA's cluster rank (`[0, 1, 0, 1]`); - a programmatic dependent launch flag: runs and checks the launch attribute; - a tensor-map parameter: compiles the bundle with NVCC. - `tests/python/codegen/test_target_codegen_cuda.py`, `tests/python/tirx/codegen/test_codegen_cuda.py`, `tests/python/tirx-transform/test_tir_transform_split_host_device.py`, `tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py`: 487 passed, 8 skipped. The remaining 120 failures are all parametrizations of `test_ptx_cp_async`, which fails while parsing its TVMScript body (`prim._OpEQ` receives a string), before any code generation; this change does not touch that path. The same change is proposed for the v0.27.0 release branch in #20489. There it was also verified end to end with mlc-ai/TIRx-kernels (whose kernels currently target the v0.27.0 script APIs): with mlc-ai/TIRx-kernels#3, all 220 correctness configs of its nine single-GPU native kernels pass when every CUDA compile is built with `host="cuda_host"`, bundled with `export_cuda_host`, compiled by `tvm_ffi.cpp.build_inline` and loaded with `tvm_ffi.load_module`.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
fp16_bf16_gemmandmla_dsv4_multishapeencode their TMA descriptors in the host prelude with a hand-writtentxl.call_packed("runtime.cuTensorMapEncodeTiled", ...). TVM'scuda_hosttarget (tvm.backend.cuda.export_cuda_host) encodes descriptors itself from the typedtensormap_encode_tiledop and rejects the packed runtime encoder, so these two kernels could not be built as a CUDA-host bundle.This adds
txl.cu_tensor_map_encode_tiled(descriptor, dtype, rank, data, *operands). It takes the same operands in the same order as the packed call (shape operands, then the interleave, swizzle, L2-promotion and OOB-fill modes and an optional forced CUDA dtype) and emitstensormap_encode_tiled. The five call sites switch to it by name only. Ordinary hosts lower the typed op back to the same packedruntime.cuTensorMapEncodeTiledcall, so the default build path is unchanged.The CUDA-host side needs apache/tvm#20489 (C++ host-pass types, typed launch arguments and cluster/PDL launch attributes).
Testing
On B200 with TVM v0.27.0 plus apache/tvm#20489:
tests/test_tirx_lite_api.py: 37 passed, including two new tests for the helper (emitted op, attributes and operand count; operand-count validation).host="cuda_host", bundled byexport_cuda_host, compiled bytvm_ffi.cpp.build_inlineand loaded bytvm_ffi.load_module: all correctness configs pass for fp16_bf16_gemm (10) and mla_dsv4_multishape (94), along with every config of nvfp4_gemm, rmsnorm, alphamoe_fp8_blockscale_qwen3next, kda_decode_multishape, msa_prefill_multishape, msa_decode_multishape and vsa_multishape (220 in total).