tensorrt_llm/_torch/models/modeling_gpt_oss.py (1)

39-41: Make the token threshold configurable.

Allow override via env or model config to tune crossover without code changes.

-# Use TinyGEMM when the number of tokens is not larger than this threshold
-MIN_LATENCY_TINYGEMM_NUM_TOKENS = 128
+# Use TinyGEMM when the number of tokens is not larger than this threshold
+# Allow override via env var for tuning: TRTLLM_TINYGEMM_TOKEN_THRESHOLD
+MIN_LATENCY_TINYGEMM_NUM_TOKENS = int(os.environ.get("TRTLLM_TINYGEMM_TOKEN_THRESHOLD", 128))

cpp/tensorrt_llm/thop/gptossTinyGemm.cpp (2)

14-25: Strengthen input validation (device + layout).

Ensure all tensors are on the same CUDA device and have dense layout before launching CUDA.

 torch::Tensor tinygemm2_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias)
 {
     TORCH_CHECK(input.dim() == 2, "input must be 2D");
     TORCH_CHECK(weight.dim() == 2, "weight must be 2D");
     TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
     TORCH_CHECK(input.sizes()[1] == weight.sizes()[1], "input.size(1) must match weight.size(1)");
     TORCH_CHECK(weight.sizes()[0] == bias.sizes()[0], "weight.size(0) must match bias.size(0)");
+    TORCH_CHECK(input.is_cuda() && weight.is_cuda() && bias.is_cuda(), "All tensors must be CUDA tensors");
+    TORCH_CHECK(
+        input.get_device() == weight.get_device() && weight.get_device() == bias.get_device(),
+        "All tensors must be on the same CUDA device");
+    TORCH_CHECK(
+        input.is_non_overlapping_and_dense() && weight.is_non_overlapping_and_dense() && bias.is_non_overlapping_and_dense(),
+        "All tensors must be dense (non-overlapping) for the CUDA kernel");
     CHECK_INPUT(input, torch::kBFloat16);
     CHECK_INPUT(weight, torch::kBFloat16);
     CHECK_INPUT(bias, torch::kBFloat16);
     return tinygemm2_cuda_forward(input, weight, bias);
 }

---

28-33: Optional: rename schema arg to 'weight' for consistency.

Purely cosmetic; avoids confusion vs. conv usage elsewhere.

-    m.def(
-        "gptoss_tinygemm(Tensor input, Tensor conv_weight, "
-        "Tensor bias) -> Tensor");
+    m.def("gptoss_tinygemm(Tensor input, Tensor weight, Tensor bias) -> Tensor");

cpp/tensorrt_llm/kernels/gptoss_tinygemm2/tinygemm2_kernel.cuh (10)

17-22: Use C++ headers and drop unused include.

Replace C headers with C++ equivalents and remove (unused).

-#include "cuda_bf16.h"
-#include <assert.h>
-#include <stdint.h>
-#include <stdio.h>
-#include <vector>
+#include "cuda_bf16.h"
+#include <cassert>
+#include <cstdint>
+#include <cstdio>

---

132-134: Prefer #if defined over #ifdef.

Aligns with coding guideline.

-#ifdef INTERNAL
+#if defined(INTERNAL)
     asm volatile(".pragma \"set knob DontInsertYield\";\n" : : : "memory");
-#endif
+#endif

---

55-85: Potential alignment/aliasing risk when writing 32-bit regs into bf16 arrays.

Reinterpreting __nv_bfloat16* as int* assumes 4-byte alignment. Guard this by asserting/alignment hints, or store through a uint32_t temp and a properly aligned pointer.

Example fix for one site (repeat for others):

-    int* rvi = reinterpret_cast<int*>(&rv[0]);
-    rvi[0] = dst;
+    // rv is 2 x 16-bit; enforce 4-byte aligned store
+    auto* rvi = reinterpret_cast<uint32_t*>(__builtin_assume_aligned(&rv[0], alignof(uint32_t)));
+    *rvi = static_cast<uint32_t>(dst);

Alternatively, change rv to be an aligned buffer type at call sites or pass a uint32_t reference and reinterpret after.

---

179-181: Always use braces for if/else bodies.

Guideline requires braces even for single statements.

-    if (PROFILE && threadIdx.x == 0 && blockIdx.y == 0)
-        profile[blockIdx.x].start = gclock64();
+    if (PROFILE && threadIdx.x == 0 && blockIdx.y == 0)
+    {
+        profile[blockIdx.x].start = gclock64();
+    }

---

263-267: Braces for profiling branches.

Apply braces for consistency.

-            if (PROFILE && blockIdx.y == 0 && ki == 0 && weight_warp)
-                profile[blockIdx.x].weight_load_start = gclock64();
-            if (PROFILE && blockIdx.y == 0 && ki == 0 && !weight_warp)
-                profile[blockIdx.x].act_load_start = gclock64();
+            if (PROFILE && blockIdx.y == 0 && ki == 0 && weight_warp)
+            {
+                profile[blockIdx.x].weight_load_start = gclock64();
+            }
+            if (PROFILE && blockIdx.y == 0 && ki == 0 && !weight_warp)
+            {
+                profile[blockIdx.x].act_load_start = gclock64();
+            }

---

323-327: Eliminate commented-out code and ensure row offsets are equal.

Dead/commented lines should be removed per guidelines. If row_offset equality is required, assert it when debugging.

-        // int row_offset_act = (reinterpret_cast <uintptr_t>(ptr_act)/128)%8;
-        // assert(row_offset_wt==row_offset_act);
         int row_offset_act = row_offset_wt;

---

352-354: Braces for profiling write.

-            if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0 && ki == 0)
-                profile[blockIdx.x].compute_start = gclock64();
+            if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0 && ki == 0)
+            {
+                profile[blockIdx.x].compute_start = gclock64();
+            }

---

395-401: mbarrier arrive: consider .expect_tx or comment intent.

Compute path uses mbarrier.arrive.shared (no expect_tx) whereas loader path uses expect_tx. If intentional, add a short comment; otherwise, use the matching variant.

---

195-196: Type consistency: use __nv_bfloat16 consistently.

sh_bias declared as nv_bfloat16 while the rest uses __nv_bfloat16. Unify to a single type.

-    __shared__ nv_bfloat16 sh_bias[TILE_M];
+    __shared__ __nv_bfloat16 sh_bias[TILE_M];

---

1-1: Filename style.

C++ filenames should be lowerCamelCase without underscores. Consider renaming to tinyGemm2Kernel.cuh.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro