[pull] main from NVIDIA:main

cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.cu

@@ -16,6 +16,7 @@
 
 #include "tensorrt_llm/common/cudaUtils.h"
 #include "ulyssesPostUnscatterKernel.h"
+#include <algorithm>
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
 
@@ -27,48 +28,62 @@ namespace kernels
 namespace
 {
 
-// Each block handles one (p, b, sp) tile across all H heads.
-// Reads H*D bf16 contiguous from input; writes the same bytes scattered across
-// H rows of the output in NHD layout [B, P*Sp, H, D]. The caller (op wrapper)
-// returns this storage as a transpose-view when an HND-shape output is needed,
-// so the resulting tensor is HND-shape with NHD-stride (matching what the
-// sync `_forward_unfused` path produces via `q.transpose(1, 2)`).
-//
-// threads/block = H * (D / 8); each thread copies one uint4 (8 bf16).
+// Each block copies one (p, b, psp) tile of ONE tensor (H*D bf16) into NHD storage
+// [B, P*Sp, H, D]. Q/K/V may differ in (Sp, H) (cross-attn: Q=audio vs K/V=video): grid.x
+// is packed over all three tensors' tiles (P*Sp_q + P*Sp_k + P*Sp_v) and each block maps
+// blockIdx.x to its tensor, so no block is idle even when shapes differ. Block is sized to
+// max(H)*(D/8); threads with h >= H idle only when H differs (GQA).
 template <typename T>
 __global__ void ulyssesPostUnscatterKernel(T const* __restrict__ q_in, T const* __restrict__ k_in,
     T const* __restrict__ v_in, T* __restrict__ q_out, T* __restrict__ k_out, T* __restrict__ v_out, int const P,
-    int const B, int const Sp, int const H, int const D, int const vec_per_row)
+    int const B, int const D, int const vec_per_row, int const Sp_q, int const H_q, int const Sp_k, int const H_k,
+    int const Sp_v, int const H_v)
 {
     constexpr int VEC = 8;
 
-    int const h = threadIdx.x / vec_per_row;
-    int const vec_idx = threadIdx.x - h * vec_per_row;
-
-    int const psp = blockIdx.x; // 0 .. P*Sp-1
-    int const p = psp / Sp;
-    int const sp = psp - p * Sp;
-    int const b = blockIdx.y;
-    int const PSp = P * Sp;
+    // blockIdx.x packed over Q|K|V tiles: [0, P*Sp_q) -> Q, next P*Sp_k -> K, rest -> V.
+    int const nq = P * Sp_q;
+    int const nk = P * Sp_k;
+    int const bx = blockIdx.x;
 
     T const* in_ptr;
     T* out_ptr;
-    switch (blockIdx.z)
+    int Sp, H, psp;
+    if (bx < nq)
     {
-    case 0:
         in_ptr = q_in;
         out_ptr = q_out;
-        break;
-    case 1:
+        Sp = Sp_q;
+        H = H_q;
+        psp = bx;
+    }
+    else if (bx < nq + nk)
+    {
         in_ptr = k_in;
         out_ptr = k_out;
-        break;
-    default:
+        Sp = Sp_k;
+        H = H_k;
+        psp = bx - nq;
+    }
+    else
+    {
         in_ptr = v_in;
         out_ptr = v_out;
-        break;
+        Sp = Sp_v;
+        H = H_v;
+        psp = bx - nq - nk;
     }
 
+    int const h = threadIdx.x / vec_per_row;
+    if (h >= H) // block sized to max(H); mask extra heads only when H differs (GQA)
+        return;
+    int const vec_idx = threadIdx.x - h * vec_per_row;
+
+    int const p = psp / Sp;
+    int const sp = psp - p * Sp;
+    int const b = blockIdx.y;
+    int const PSp = P * Sp;
+
     // in[p, b, sp, h, d]: ((((p*B + b)*Sp + sp)*H + h)*D + vec_idx*VEC)
     // NHD out[b, p*Sp+sp, h, d]: (((b*PSp + psp)*H + h)*D + vec_idx*VEC)
     // int64_t: P*B*Sp*H*D can exceed 2^31 at large workloads.
@@ -83,16 +98,18 @@ __global__ void ulyssesPostUnscatterKernel(T const* __restrict__ q_in, T const*
 } // namespace
 
 void launchUlyssesPostUnscatter(void const* q_in, void const* k_in, void const* v_in, void* q_out, void* k_out,
-    void* v_out, int P, int B, int Sp, int H, int D, cudaStream_t stream)
+    void* v_out, int P, int B, int D, int Sp_q, int H_q, int Sp_k, int H_k, int Sp_v, int H_v, cudaStream_t stream)
 {
     constexpr int VEC = 8;
     TLLM_CHECK_WITH_INFO(D % VEC == 0, "ulyssesPostUnscatter: D must be a multiple of 8 (uint4 vec), got %d", D);
     int const vec_per_row = D / VEC;
-    int const threads = H * vec_per_row;
+    int const H_max = std::max(H_q, std::max(H_k, H_v));
+    int const threads = H_max * vec_per_row;
     TLLM_CHECK_WITH_INFO(threads <= 1024,
-        "ulyssesPostUnscatter: threads/block (H*D/8) must be <= 1024, got H=%d D=%d -> %d", H, D, threads);
+        "ulyssesPostUnscatter: threads/block (maxH*D/8) must be <= 1024, got maxH=%d D=%d -> %d", H_max, D, threads);
 
-    dim3 const grid(P * Sp, B, 3);
+    int const total_tiles = P * (Sp_q + Sp_k + Sp_v); // packed over Q/K/V, no idle blocks
+    dim3 const grid(total_tiles, B, 1);
     dim3 const block(threads);
 
     auto* q_in_typed = reinterpret_cast<__nv_bfloat16 const*>(q_in);
@@ -102,8 +119,8 @@ void launchUlyssesPostUnscatter(void const* q_in, void const* k_in, void const*
     auto* k_out_typed = reinterpret_cast<__nv_bfloat16*>(k_out);
     auto* v_out_typed = reinterpret_cast<__nv_bfloat16*>(v_out);
 
-    ulyssesPostUnscatterKernel<__nv_bfloat16><<<grid, block, 0, stream>>>(
-        q_in_typed, k_in_typed, v_in_typed, q_out_typed, k_out_typed, v_out_typed, P, B, Sp, H, D, vec_per_row);
+    ulyssesPostUnscatterKernel<__nv_bfloat16><<<grid, block, 0, stream>>>(q_in_typed, k_in_typed, v_in_typed,
+        q_out_typed, k_out_typed, v_out_typed, P, B, D, vec_per_row, Sp_q, H_q, Sp_k, H_k, Sp_v, H_v);
 }
 
 } // namespace kernels

cpp/tensorrt_llm/kernels/ulyssesPostUnscatterKernel.h

@@ -47,16 +47,18 @@ namespace kernels
 //   - Each block reads one fully contiguous (p, b, sp, :H, :D) tile of
 //     H*D bf16
 //   - H*(D/8) threads/block — each thread copies one uint4 (8 bf16)
-//   - Grid (P*Sp, B, 3): blockIdx.z selects Q / K / V
+//   - Grid (P*(Sp_q+Sp_k+Sp_v), B): blockIdx.x is packed over Q/K/V tiles and mapped
+//     to the owning tensor -> no idle blocks even when Q/K/V shapes differ
 //
 // Constraints:
 //   - dtype must be bf16
 //   - D must be a multiple of 8 (uint4 vector load/store, 8 bf16 per thread)
-//   - threads/block = H * (D / 8) must be <= 1024 (CUDA hw limit)
-void launchUlyssesPostUnscatter(void const* q_in, // [P, B, Sp, H, D]
-    void const* k_in, void const* v_in,
-    void* q_out,                                  // [B, P*Sp, H, D] NHD-contig
-    void* k_out, void* v_out, int P, int B, int Sp, int H, int D, cudaStream_t stream);
+//   - threads/block = max(H) * (D / 8) must be <= 1024 (CUDA hw limit)
+void launchUlyssesPostUnscatter(void const* q_in, // [P, B, Sp_q, H_q, D]
+    void const* k_in, void const* v_in,           // [P, B, Sp_{k,v}, H_{k,v}, D]
+    void* q_out,                                  // per-tensor [B, P*Sp, H, D] NHD-contig
+    void* k_out, void* v_out, int P, int B, int D, int Sp_q, int H_q, int Sp_k, int H_k, int Sp_v, int H_v,
+    cudaStream_t stream);
 
 } // namespace kernels
 

cpp/tensorrt_llm/thop/ulyssesPostUnscatterOp.cpp

@@ -27,19 +27,11 @@ namespace torch_ext
 
 // Post-Ulysses A2A unscatter: take Q/K/V tensors of shape [P, B, Sp, H, D]
 // (output of the head-dim -> seq-dim all-to-all) and produce SDPA-ready Q/K/V.
-// The kernel ALWAYS writes NHD-contig storage [B, P*Sp, H, D]. The returned
-// tensor shape depends on ``layout``:
-//   layout=0 (HND) → returns transpose-view [B, H, P*Sp, D]
-//                    (HND-shape, NHD-stride, NON-contig — mirrors the
-//                    `q.transpose(1, 2)` result in the sync `_forward_unfused`
-//                    path, which lets cudnn SDPA preserve NHD-stride through
-//                    its output and collapses the downstream
-//                    `_output_a2a.transpose(1, 2).contiguous()` to a no-op)
-//   layout=1 (NHD) → returns storage as-is [B, P*Sp, H, D] (NHD contig)
-// Replaces the eager chain
-//     t.permute(1, 0, 2, 3, 4).reshape(B, P * Sp, H, D).contiguous()           // NHD
-//     [.transpose(1, 2)]                                                       // HND: stride view only
-// for Q, K, V in one kernel launch.
+// The kernel writes NHD-contig storage [B, P*Sp, H, D]; the return depends on ``layout``:
+//   layout=0 (HND) → transpose-view [B, H, P*Sp, D] (HND-shape, NHD-stride, non-contig)
+//   layout=1 (NHD) → storage as-is [B, P*Sp, H, D] (NHD-contig)
+// Equivalent eager: t.permute(1,0,2,3,4).reshape(B, P*Sp, H, D).contiguous()[.transpose(1,2) if HND].
+// Q/K/V may differ in (Sp, H) (cross-attn: Q=audio vs K/V=video) — one packed launch handles both.
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> ulysses_post_unscatter_qkv(
     torch::Tensor& q_in, // [P, B, Sp, H, D]
     torch::Tensor& k_in, // [P, B, Sp, H, D]
@@ -48,38 +40,39 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> ulysses_post_unscatter_q
 {
     TORCH_CHECK(q_in.dim() == 5 && k_in.dim() == 5 && v_in.dim() == 5,
         "ulysses_post_unscatter_qkv expects 5D tensors [P, B, Sp, H, D]");
-    TORCH_CHECK(q_in.sizes() == k_in.sizes() && q_in.sizes() == v_in.sizes(), "Q/K/V must share the same shape");
     TORCH_CHECK(layout == 0 || layout == 1, "layout must be 0 (HND) or 1 (NHD), got ", layout);
 
     CHECK_INPUT(q_in, torch::kBFloat16);
     CHECK_INPUT(k_in, torch::kBFloat16);
     CHECK_INPUT(v_in, torch::kBFloat16);
 
-    // D % 8 enforced here at op boundary (mirrors sibling ulysses_permute_scatter).
-    // Without this, torch::empty allocates the three output tensors before the
-    // kernel launcher's TLLM_CHECK_WITH_INFO fires, producing a less-actionable
-    // error path. Vec width is 8 elements for bf16 (16-byte vectorized stores).
-    TORCH_CHECK(q_in.size(-1) % 8 == 0, "D (last dim) must be divisible by 8 (bf16 vec=8)");
-
+    // P (ulysses degree), B (batch), D (head dim) are shared across Q/K/V; Sp (seq shard)
+    // and H (heads/rank) may differ (cross-attn: Q=audio vs K/V=video). The launcher's
+    // packed grid handles equal or different shapes in one launch.
     int64_t const P = q_in.size(0);
     int64_t const B = q_in.size(1);
-    int64_t const Sp = q_in.size(2);
-    int64_t const H = q_in.size(3);
     int64_t const D = q_in.size(4);
+    TORCH_CHECK(k_in.size(0) == P && v_in.size(0) == P, "Q/K/V must share P (ulysses degree)");
+    TORCH_CHECK(k_in.size(1) == B && v_in.size(1) == B, "Q/K/V must share B (batch)");
+    TORCH_CHECK(k_in.size(4) == D && v_in.size(4) == D, "Q/K/V must share D (head dim)");
+    // D % 8 enforced at the op boundary: vec width is 8 bf16 (16-byte vectorized stores).
+    TORCH_CHECK(D % 8 == 0, "D (last dim) must be divisible by 8 (bf16 vec=8)");
+
+    int64_t const Sp_q = q_in.size(2), H_q = q_in.size(3);
+    int64_t const Sp_k = k_in.size(2), H_k = k_in.size(3);
+    int64_t const Sp_v = v_in.size(2), H_v = v_in.size(3);
 
     bool const is_hnd = (layout == 0);
     auto opts = q_in.options();
-    // Always allocate NHD-contig storage [B, P*Sp, H, D].
-    auto const storage_shape = std::vector<int64_t>{B, P * Sp, H, D};
-    auto q_out = torch::empty(storage_shape, opts);
-    auto k_out = torch::empty(storage_shape, opts);
-    auto v_out = torch::empty(storage_shape, opts);
-
-    // Empty-tensor no-op: P=0/B=0/Sp=0 produces zero grid extent in the
-    // kernel launcher (undefined cuLaunchKernel behavior across CUDA versions).
-    // The three output tensors above are already empty-shaped via P*Sp=0 or B=0,
-    // so returning them directly preserves the output-shape contract.
-    if (q_in.numel() == 0)
+    // Per-tensor NHD-contig storage [B, P*Sp, H, D].
+    auto q_out = torch::empty({B, P * Sp_q, H_q, D}, opts);
+    auto k_out = torch::empty({B, P * Sp_k, H_k, D}, opts);
+    auto v_out = torch::empty({B, P * Sp_v, H_v, D}, opts);
+
+    // Empty-tensor no-op: only when ALL three are empty (P=0 / B=0 / all Sp=0) is the grid
+    // extent zero (UB in the launcher). If any tensor is non-empty the launch is valid and
+    // the kernel bounds-checks each tensor's own (Sp, H), so a per-tensor empty is fine.
+    if (q_in.numel() == 0 && k_in.numel() == 0 && v_in.numel() == 0)
     {
         if (is_hnd)
         {
@@ -91,7 +84,8 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> ulysses_post_unscatter_q
     auto stream = at::cuda::getCurrentCUDAStream();
     tensorrt_llm::kernels::launchUlyssesPostUnscatter(q_in.data_ptr(), k_in.data_ptr(), v_in.data_ptr(),
         q_out.data_ptr(), k_out.data_ptr(), v_out.data_ptr(), static_cast<int>(P), static_cast<int>(B),
-        static_cast<int>(Sp), static_cast<int>(H), static_cast<int>(D), stream);
+        static_cast<int>(D), static_cast<int>(Sp_q), static_cast<int>(H_q), static_cast<int>(Sp_k),
+        static_cast<int>(H_k), static_cast<int>(Sp_v), static_cast<int>(H_v), stream);
 
     // HND callers get a transpose-view of the NHD storage (zero-copy stride
     // reinterpretation). NHD callers get the storage as-is.

docs/source/models/supported-models.md

@@ -19,6 +19,7 @@ The following is a table of supported models for the PyTorch backend:
 | `Gemma3ForCausalLM`                  | Gemma 3                            | `google/gemma-3-1b-it`                       |
 | `Gemma3nForConditionalGeneration` [^7]| Gemma 3n                           | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it` |
 | `Gemma4ForConditionalGeneration`     | Gemma 4                            | `google/gemma-4-E2B-it`, `google/gemma-4-E4B-it`, `google/gemma-4-26B-A4B-it` [^6], `google/gemma-4-31B-it` [^6] |
+| `Gemma4UnifiedForConditionalGeneration` | Gemma 4 12B Unified (encoder-free) | `google/gemma-4-12B`, `google/gemma-4-12B-it`              |
 | `Glm4MoeForCausalLM`                 | GLM-4.5, GLM-4.6, GLM-4.7          | `THUDM/GLM-4-100B-A10B`                      |
 | `Glm4MoeLiteForCausalLM` [^5]        | GLM-4.7-Flash                      | `zai-org/GLM-4.7-Flash`                      |
 | `GlmMoeDsaForCausalLM`               | GLM-5                              | `zai-org/GLM-5`                              |
@@ -74,6 +75,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `Glm4MoeLiteForCausalLM` [^5]    | Yes               | Yes        | Untested                   | Untested              | Yes             | No  | No               | No                | No     | Yes           | Untested         | Untested       | N/A                      | Untested              | Untested        |
 | `NemotronHForCausalLM`           | Yes               | Yes        | Yes                        | Yes                   | Yes             | Yes | No               | No                | No     | Yes           | Yes              | Yes            | N/A                      | Untested              | Untested        |
 | `Gemma4ForConditionalGeneration` | Untested          | Yes        | Untested                   | No                    | Yes             | No  | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
+| `Gemma4UnifiedForConditionalGeneration` | Untested          | Untested   | Untested                   | No                    | Yes             | No  | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
 | `Step3p7ForConditionalGeneration`| Yes               | Yes        | Yes                        | Untested              | Untested        | Yes | No               | No                | No     | Yes           | Untested         | Untested       | Yes                      | Untested              | Untested        |
 | `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | No  | No               | No                | No     | Yes           | Untested         | No             | N/A                      | Untested              | Untested        |
 
@@ -97,6 +99,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `Exaone4_5_ForConditionalGeneration` | Yes               | Yes        | Yes             | Yes           | Yes              | Yes            | Yes                   | No                        | L + I + V |
 | `Gemma3ForConditionalGeneration`     | Yes               | Yes        | N/A             | Yes           | Yes              | N/A            | Yes                   | No                        | L + I     |
 | `Gemma4ForConditionalGeneration`     | Untested          | Yes        | Yes             | Yes           | Untested         | No             | Untested              | No                        | L + I + V + A [^9] |
+| `Gemma4UnifiedForConditionalGeneration` | Untested          | Untested   | Untested        | Yes           | Untested         | No             | Untested              | No                        | L + I + A |
 | `HCXVisionForCausalLM`               | Yes               | Yes        | No              | Yes           | Yes              | Yes            | Yes                   | No                        | L + I     |
 | `LlavaLlamaModel (VILA)`             | Yes               | Yes        | No              | Yes           | Yes              | No             | Yes                   | No                        | L + I + V |
 | `LlavaNextForConditionalGeneration`  | Yes               | Yes        | Yes             | Yes           | Yes              | Yes            | Yes                   | Yes                       | L + I     |

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -1133,15 +1133,16 @@ def update_spec_dec_param(
                 # Dynamic draft length needs position offsets and packed mask to be shaped for each runtime draft length.
                 # So we create cache for position offsets and packed mask for each draft length to avoid reallocation.
                 assert max_draft_len == max_total_draft_tokens, "max_draft_len should be equal to max_total_draft_tokens for linear tree"
-                runtime_draft_len = (spec_metadata.runtime_draft_len
-                                     if spec_metadata is not None else
-                                     max_draft_len)
+                # For algos other than PARD, this equals runtime_draft_len (K); for PARD it's 2K-1.
+                runtime_draft_token_buffer_width = (
+                    spec_metadata.runtime_tokens_per_gen_step -
+                    1 if spec_metadata is not None else max_draft_len)
                 self.generate_spec_decoding_generation_length(
-                    runtime_draft_len=runtime_draft_len)
+                    runtime_draft_len=runtime_draft_token_buffer_width)
                 self.spec_decoding_position_offsets = generate_spec_decoding_position_offsets(
-                    self.max_num_requests, runtime_draft_len)
+                    self.max_num_requests, runtime_draft_token_buffer_width)
                 self.spec_decoding_packed_mask = generate_spec_decoding_packed_mask(
-                    self.max_num_requests, runtime_draft_len)
+                    self.max_num_requests, runtime_draft_token_buffer_width)
 
             self.update_position_offsets_for_cpp(cpp_query_len)