rename fields

Author: brb-nv

cpp/tensorrt_llm/thop/alltoallOp.cpp

@@ -115,102 +115,102 @@ std::vector<torch::Tensor> alltoall_helix(
 /**
  * Helix All-to-All operation with two fields.
  *
- * Input tensors have shape [..., cp_size, kv_lora_rank] for field0 and [...,
- * cp_size, 2] for field1. The operation exchanges data along the cp_size
+ * Input tensors have shape [..., cp_size, kv_lora_rank] for partial_o and [...,
+ * cp_size, 2] for softmax_stats. The operation exchanges data along the cp_size
  * dimension across all ranks.
  *
- * @param field0 Field 0 tensor (half precision, shape [..., cp_size,
+ * @param partial_o Field 0 tensor (half precision, shape [..., cp_size,
  * kv_lora_rank])
- * @param field1 Field 1 tensor (float32, shape [..., cp_size, 2])
+ * @param softmax_stats Field 1 tensor (float32, shape [..., cp_size, 2])
  * @param workspace Workspace tensor (uint64, strided across ranks)
  * @param cp_rank Current context parallel rank
  * @param cp_size Total number of context parallel ranks
- * @return tuple of (field0_out, field1_out) with same shapes as inputs
+ * @return tuple of (partial_o_out, softmax_stats_out) with same shapes as inputs
  */
 std::tuple<torch::Tensor, torch::Tensor> alltoall_helix_native(
-    torch::Tensor field0, torch::Tensor field1, torch::Tensor workspace, int64_t cp_rank, int64_t cp_size)
+    torch::Tensor partial_o, torch::Tensor softmax_stats, torch::Tensor workspace, int64_t cp_rank, int64_t cp_size)
 {
 
     // Input validation
-    CHECK_TH_CUDA(field0);
-    CHECK_TH_CUDA(field1);
+    CHECK_TH_CUDA(partial_o);
+    CHECK_TH_CUDA(softmax_stats);
     CHECK_TH_CUDA(workspace);
-    CHECK_CONTIGUOUS(field0);
-    CHECK_CONTIGUOUS(field1);
+    CHECK_CONTIGUOUS(partial_o);
+    CHECK_CONTIGUOUS(softmax_stats);
 
     // Type checks
-    TORCH_CHECK(field0.scalar_type() == at::ScalarType::Half || field0.scalar_type() == at::ScalarType::BFloat16,
-        "field0 must be half or bfloat16");
-    CHECK_TYPE(field1, at::ScalarType::Float);
+    TORCH_CHECK(partial_o.scalar_type() == at::ScalarType::Half || partial_o.scalar_type() == at::ScalarType::BFloat16,
+        "partial_o must be half or bfloat16");
+    CHECK_TYPE(softmax_stats, at::ScalarType::Float);
     CHECK_TYPE(workspace, at::ScalarType::UInt64);
 
     // Shape validation
-    TORCH_CHECK(field0.dim() >= 2, "field0 must have at least 2 dimensions");
-    TORCH_CHECK(field1.dim() >= 2, "field1 must have at least 2 dimensions");
-    TORCH_CHECK(field0.dim() == field1.dim(), "field0 and field1 must have same number of dimensions");
+    TORCH_CHECK(partial_o.dim() >= 2, "partial_o must have at least 2 dimensions");
+    TORCH_CHECK(softmax_stats.dim() >= 2, "softmax_stats must have at least 2 dimensions");
+    TORCH_CHECK(partial_o.dim() == softmax_stats.dim(), "partial_o and softmax_stats must have same number of dimensions");
 
     // Get dimensions
-    int kv_lora_rank = field0.size(-1);
-    TORCH_CHECK(field0.size(-2) == cp_size && field1.size(-2) == cp_size,
-        "field0/1 second-to-last dimension must equal cp_size");
+    int kv_lora_rank = partial_o.size(-1);
+    TORCH_CHECK(partial_o.size(-2) == cp_size && softmax_stats.size(-2) == cp_size,
+        "partial_o/softmax_stats second-to-last dimension must equal cp_size");
     TORCH_CHECK(
-        field1.size(-1) % 2 == 0 && field1.size(-1) >= 2, "field1 last dimension must be divisible by 2 (float2)");
-    bool allowVariableField1 = field1.size(-1) > 2;
+        softmax_stats.size(-1) % 2 == 0 && softmax_stats.size(-1) >= 2, "softmax_stats last dimension must be divisible by 2 (float2)");
+    bool allowVariableField1 = softmax_stats.size(-1) > 2;
 
     // Check that leading dimensions match
-    for (int i = 0; i < field0.dim() - 2; i++)
+    for (int i = 0; i < partial_o.dim() - 2; i++)
     {
         TORCH_CHECK(
-            field0.size(i) == field1.size(i), "field0 and field1 must have matching dimensions except last two");
+            partial_o.size(i) == softmax_stats.size(i), "partial_o and softmax_stats must have matching dimensions except last two");
     }
-    TORCH_CHECK(field0.size(-1) * field0.element_size() % 16 == 0, "field0 must be aligned to 16 bytes");
+    TORCH_CHECK(partial_o.size(-1) * partial_o.element_size() % 16 == 0, "partial_o must be aligned to 16 bytes");
 
     TORCH_CHECK(workspace.dim() == 2, "workspace must be 2D (strided across ranks)");
     TORCH_CHECK(workspace.size(0) == cp_size, "workspace must have cp_size rows");
 
     // Calculate entry count (product of all dimensions before cp_size)
     // This is the number of entries to process per peer rank
     int entry_count = 1;
-    for (int i = 0; i < field0.dim() - 2; i++)
+    for (int i = 0; i < partial_o.dim() - 2; i++)
     {
-        entry_count *= field0.size(i);
+        entry_count *= partial_o.size(i);
     }
 
     // Reshape to 3D: [entry_count, cp_size, feature_dim]
-    torch::Tensor field0_3d = field0.reshape({entry_count, cp_size, kv_lora_rank});
-    torch::Tensor field1_3d = field1.reshape({entry_count, cp_size, field1.size(-1)});
+    torch::Tensor partial_o_3d = partial_o.reshape({entry_count, cp_size, kv_lora_rank});
+    torch::Tensor softmax_stats_3d = softmax_stats.reshape({entry_count, cp_size, softmax_stats.size(-1)});
 
     // Allocate output tensors (same shape as input)
-    torch::Tensor field0_out = torch::empty_like(field0);
-    torch::Tensor field1_out = torch::empty_like(field1);
+    torch::Tensor partial_o_out = torch::empty_like(partial_o);
+    torch::Tensor softmax_stats_out = torch::empty_like(softmax_stats);
 
-    torch::Tensor field0_out_3d = field0_out.reshape({entry_count, cp_size, kv_lora_rank});
-    torch::Tensor field1_out_3d = field1_out.reshape({entry_count, cp_size, field1.size(-1)});
+    torch::Tensor partial_o_out_3d = partial_o_out.reshape({entry_count, cp_size, kv_lora_rank});
+    torch::Tensor softmax_stats_out_3d = softmax_stats_out.reshape({entry_count, cp_size, softmax_stats.size(-1)});
 
     // Setup parameters
     tensorrt_llm::kernels::HelixAllToAllParams params;
 
     // Field 0 (variable size half)
-    params.sendFields[0].dataPtr = reinterpret_cast<uint8_t*>(field0_3d.data_ptr());
+    params.sendFields[0].dataPtr = reinterpret_cast<uint8_t*>(partial_o_3d.data_ptr());
     params.sendFields[0].elementCount = kv_lora_rank;
-    params.sendFields[0].elementSize = field0.element_size();
-    params.sendFields[0].stride = field0_3d.stride(1) * field0.element_size();
+    params.sendFields[0].elementSize = partial_o.element_size();
+    params.sendFields[0].stride = partial_o_3d.stride(1) * partial_o.element_size();
 
-    params.recvFields[0].dataPtr = reinterpret_cast<uint8_t*>(field0_out_3d.data_ptr());
+    params.recvFields[0].dataPtr = reinterpret_cast<uint8_t*>(partial_o_out_3d.data_ptr());
     params.recvFields[0].elementCount = kv_lora_rank;
-    params.recvFields[0].elementSize = field0.element_size();
-    params.recvFields[0].stride = field0_out_3d.stride(1) * field0.element_size();
+    params.recvFields[0].elementSize = partial_o.element_size();
+    params.recvFields[0].stride = partial_o_out_3d.stride(1) * partial_o.element_size();
 
     // Field 1 (single float2)
-    params.sendFields[1].dataPtr = reinterpret_cast<uint8_t*>(field1_3d.data_ptr<float>());
-    params.sendFields[1].elementCount = field1.size(-1);
-    params.sendFields[1].elementSize = field1.element_size();
-    params.sendFields[1].stride = field1_3d.stride(1) * field1.element_size();
+    params.sendFields[1].dataPtr = reinterpret_cast<uint8_t*>(softmax_stats_3d.data_ptr<float>());
+    params.sendFields[1].elementCount = softmax_stats.size(-1);
+    params.sendFields[1].elementSize = softmax_stats.element_size();
+    params.sendFields[1].stride = softmax_stats_3d.stride(1) * softmax_stats.element_size();
 
-    params.recvFields[1].dataPtr = reinterpret_cast<uint8_t*>(field1_out_3d.data_ptr<float>());
-    params.recvFields[1].elementCount = field1.size(-1);
-    params.recvFields[1].elementSize = field1.element_size();
-    params.recvFields[1].stride = field1_out_3d.stride(1) * field1.element_size();
+    params.recvFields[1].dataPtr = reinterpret_cast<uint8_t*>(softmax_stats_out_3d.data_ptr<float>());
+    params.recvFields[1].elementCount = softmax_stats.size(-1);
+    params.recvFields[1].elementSize = softmax_stats.element_size();
+    params.recvFields[1].stride = softmax_stats_out_3d.stride(1) * softmax_stats.element_size();
 
     // Entry count and workspace
     params.entryCount = entry_count;
@@ -227,7 +227,7 @@ std::tuple<torch::Tensor, torch::Tensor> alltoall_helix_native(
     auto stream = at::cuda::getCurrentCUDAStream();
     tensorrt_llm::kernels::launchHelixAllToAll(params, allowVariableField1, stream);
 
-    return std::make_tuple(field0_out, field1_out);
+    return std::make_tuple(partial_o_out, softmax_stats_out);
 }
 
 /**
@@ -267,7 +267,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
     m.def("alltoall_helix(Tensor[] input_list, int[] group, int? num_lists) -> Tensor[]");
     m.def(
-        "alltoall_helix_native(Tensor field0, Tensor field1, Tensor workspace, int "
+        "alltoall_helix_native(Tensor partial_o, Tensor softmax_stats, Tensor workspace, int "
         "cp_rank, int cp_size) -> (Tensor, Tensor)");
     m.def("get_helix_workspace_size_per_rank(Tensor __dummy__, int cp_size) -> int");
     m.def(

tensorrt_llm/_torch/modules/attention.py

@@ -1147,24 +1147,24 @@ def _attn_forward_gen(self, attn_backend: AttentionBackend, q: torch.Tensor,
                 # partial_o: [num_tokens, num_heads * kv_lora_rank] -> [num_tokens, cp_size, num_heads_tp_cp, kv_lora_rank]
                 # softmax_stats: [num_tokens, num_heads, 2] -> [num_tokens, cp_size, num_heads_tp_cp, 2]
 
-                field0 = partial_o.view(num_tokens, cp_size,
+                partial_o = partial_o.view(num_tokens, cp_size,
                                         self.num_heads_tp_cp,
                                         kv_lora_rank).transpose(1,
                                                                 2).contiguous()
-                field1 = softmax_stats.view(num_tokens, cp_size,
+                softmax_stats = softmax_stats.view(num_tokens, cp_size,
                                             self.num_heads_tp_cp,
                                             2).transpose(1, 2).contiguous()
 
                 # Call FIFO-based helixAllToAll.
-                field0_out, field1_out = helix.alltoall_native(field0, field1)
+                partial_o_out, softmax_stats_out = helix.alltoall_native(partial_o, softmax_stats)
 
-                # field0_out: [num_tokens, num_heads_tp_cp, cp_size, kv_lora_rank]
-                # field1_out: [num_tokens, num_heads_tp_cp, cp_size, 2]
+                # partial_o_out: [num_tokens, num_heads_tp_cp, cp_size, kv_lora_rank]
+                # softmax_stats_out: [num_tokens, num_heads_tp_cp, cp_size, 2]
                 # cp_dim = 2 (the dimension where cp_size is located)
 
                 # Call helix_post_process_native with cp_dim=2.
                 return torch.ops.trtllm.helix_post_process_native(
-                    field0_out, field1_out, 1.0, 2)
+                    partial_o_out, softmax_stats_out, 1.0, 2)
         else:
             attn_output = attn_backend.forward(q, k, v, attn_metadata, **kwargs)
             return attn_output