issue/1565 fix(nvidia): complete runtime support for Qwen MTP

Map E4M3 and BOOL through the existing ATen adaptor and preserve the
caller's CUDA device across NCCL communicator destruction.

Reuse the existing paged Prefill warp kernel for NVIDIA head size 256,
without changing other vendors' default dispatch. Extend existing
multi-page/long-context coverage and add finite FP8/mask cast checks.

Validation: fresh SM86 build, 88 paged Prefill cases, 2 cast tests,
and TP2 communicator teardown from both caller devices.

Closes #1565

Author: big-hip

include/infinicore/adaptor/aten_adaptor.hpp

@@ -45,6 +45,8 @@ inline at::ScalarType to_at_dtype(DataType dtype) {
         return at::kHalf;
     case DataType::BF16:
         return at::kBFloat16;
+    case DataType::F8:
+        return at::kFloat8_e4m3fn;
     case DataType::I8:
         return at::kChar;
     case DataType::U8:
@@ -53,6 +55,8 @@ inline at::ScalarType to_at_dtype(DataType dtype) {
         return at::kInt;
     case DataType::I64:
         return at::kLong;
+    case DataType::BOOL:
+        return at::kBool;
     default:
         throw std::runtime_error("Unsupported dtype for ATen");
     }

src/infiniccl/cuda/infiniccl_cuda.cu

@@ -119,7 +119,13 @@ infiniStatus_t commInitRank(
 }
 
 infiniStatus_t commDestroy(infinicclComm_t comm) {
-    CHECK_NCCL(ncclCommDestroy(getNcclComm(comm)));
+    // NCCL teardown may activate the communicator's device. Preserve the
+    // caller's device so its existing stream and runtime context stay valid.
+    int previous_device;
+    CHECK_INTERNAL(cudaGetDevice(&previous_device), cudaSuccess);
+    const auto status = ncclCommDestroy(getNcclComm(comm));
+    CHECK_INTERNAL(cudaSetDevice(previous_device), cudaSuccess);
+    CHECK_NCCL(status);
     delete comm;
     return INFINI_STATUS_SUCCESS;
 }

src/infiniop/ops/paged_attention_prefill/cuda/kernel_v2.cuh

@@ -280,7 +280,7 @@ __global__ void PagedAttentionPrefillWarpGlobalKernel(
     ptrdiff_t o_head_stride) {
 
     constexpr int kWarpSize = 32;
-    static_assert(HEAD_SIZE == 64 || HEAD_SIZE == 128 || HEAD_SIZE == 192, "Only head_size 64/128/192 supported in v0.4.");
+    static_assert(HEAD_SIZE == 64 || HEAD_SIZE == 128 || HEAD_SIZE == 192 || HEAD_SIZE == 256, "Unsupported head_size.");
     static_assert(HEAD_SIZE % kWarpSize == 0, "HEAD_SIZE must be divisible by 32.");
     constexpr int DIMS_PER_THREAD = HEAD_SIZE / kWarpSize;
 

src/infiniop/ops/paged_attention_prefill/nvidia/paged_attention_prefill_nvidia.cu

@@ -27,7 +27,11 @@ inline const char *default_prefill_kernel(const PagedAttentionPrefillInfo &info)
         return "ref";
     }
     if (info.head_size == 256) {
+#if defined(ENABLE_NVIDIA_API)
+        return "warp";
+#else
         return "ref";
+#endif
     }
     // Iluvatar/Hygon: use warp for the non-MLA shapes where it is the stable path.
 #if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API)
@@ -1020,6 +1024,17 @@ infiniStatus_t launch_prefill_warp(
                 v_batch_stride, v_row_stride, v_head_stride,
                 o_stride, o_head_stride);
         return INFINI_STATUS_SUCCESS;
+    case 256:
+        op::paged_attention_prefill::cuda::PagedAttentionPrefillWarpGlobalKernel<Tindex, Tdata, 256>
+            <<<grid, block, 0, stream>>>(
+                out, q, k_cache, v_cache, block_tables, total_kv_lens, cu_seqlens_q, alibi_slopes,
+                num_heads, num_seqs, num_kv_heads, total_q_tokens, scale, max_num_blocks_per_seq,
+                page_block_size, block_table_batch_stride,
+                q_stride, q_head_stride,
+                k_batch_stride, k_row_stride, k_head_stride,
+                v_batch_stride, v_row_stride, v_head_stride,
+                o_stride, o_head_stride);
+        return INFINI_STATUS_SUCCESS;
     default:
         return INFINI_STATUS_BAD_TENSOR_SHAPE;
     }

test/infinicore/ops/fp8_cast.py

@@ -0,0 +1,44 @@
+"""Check FP8 weights and boolean acceptance masks across the ATen cast bridge."""
+
+import torch
+from infinicore.lib import _infinicore
+
+import infinicore
+
+
+def test_fp8_cast():
+    bits = torch.arange(256, dtype=torch.int16)
+    bits = bits[(bits != 127) & (bits != 255)].to(torch.uint8)
+    source = bits.view(torch.float8_e4m3fn).cuda()
+    for dtype in (torch.float32, torch.bfloat16):
+        output = torch.empty(source.shape, dtype=dtype, device="cuda")
+        _infinicore.cast_(
+            infinicore.from_torch(output)._underlying,
+            infinicore.from_torch(source)._underlying,
+        )
+        infinicore.sync_device()
+        torch.testing.assert_close(output, source.to(dtype), rtol=0, atol=0)
+
+
+def test_bool_cast():
+    candidates = torch.tensor([3, 5, 7, 9], device="cuda")
+    expected = torch.tensor([4, 5, 7, 8], device="cuda")
+    source = infinicore.equal(
+        infinicore.from_torch(candidates), infinicore.from_torch(expected)
+    )
+    for dtype in (torch.float32, torch.int64):
+        output = torch.empty(source.shape, dtype=dtype, device="cuda")
+        _infinicore.cast_(
+            infinicore.from_torch(output)._underlying,
+            source._underlying,
+        )
+        infinicore.sync_device()
+        torch.testing.assert_close(
+            output, (candidates == expected).to(dtype), rtol=0, atol=0
+        )
+
+
+if __name__ == "__main__":
+    test_fp8_cast()
+    test_bool_cast()
+    print("Finite E4M3 and boolean acceptance mask casts passed")

test/infinicore/ops/paged_attention_prefill.py

@@ -26,6 +26,8 @@
     (1, 24, 4, 256, 8, 8, 1),
     (1, 12, 2, 256, 8, 8, 1),
     (1, 6, 1, 256, 8, 8, 1),
+    (2, 12, 2, 256, 64, 128, 2),
+    (1, 24, 4, 256, 64, 1024, 1),
     # New DeepSeek MLA wrapper case: verifies prefill supports q/k head
     # size 576 with value head size 512.
     (1, 16, 1, 576, 8, 8, 1, 512),
@@ -94,7 +96,9 @@ def parse_test_cases():
                 value_size,
             ) = case
         scale = head_size**-0.5
-        num_blocks = 8192
+        num_blocks = num_seqs * (
+            (max_step_len * num_rounds + block_size - 1) // block_size
+        )
         manager = SimpleCacheManager(num_blocks, block_size)
         kv_lens = torch.zeros(num_seqs, dtype=torch.int32)