Skip to content
KernelIndex
Search⌘K
Sign in

gpt-o3 / cudad4241d

gpt-o3_cuda_d4241d · gpt-o3 · cuda · Apache-2.0
Kernel source · 60 lines ↓holds 1 record

Use it

Vendorable · source mirrored · Apache-2.0View source →

No package. Vendor the mirrored source: 60 lines, Apache-2.0, pinned at da91508.

main.cpp
curl "https://kernelindex.com/api/v1/implementations/flashinfer-gpt-o3-cuda-d4241d?include=source"
interfacecuda
revisionda915083d4c7
symbolrun
pathmain.cpp
Compatibility
measured onNVIDIA B200
declared hardwareNVIDIA B200
architecturessm_100
dtypesbf16, fp32, int32

Benchmark evidence

2 measurements across 1 GPU, fastest first.

Operation / workload
Hardware
Latency
Rank
Observed
GQA paged prefill causal h32 kv8 d128 ps1bf16 · [2, 32, 128]
NVIDIA B200
31.5µs
#1 of 7
2025-10-21
Run →
GQA paged prefill causal h32 kv8 d128 ps1bf16 · [1, 32, 128]
NVIDIA B200
160.6µs
#3 of 7
2025-10-21
Run →

Reported · How evidence levels are derived →

Source and license

sourcehttps://huggingface.co/datasets/flashinfer-ai/flashinfer-trace
commitda915083d4c7c5e61aa3005e3d17ae488e0fc71c
revision digestsha256:e13afb733ed2d9c9c16a265dde48f2b6a2f2c11d77272f68e73cd1c165e989c3
license declaredApache-2.0
license concludedApache-2.0
authorsgpt-o3
imported2026-08-20

Kernel source

main.cpp60 lines
#include "kernel.h"

#include <torch/extension.h>
#include <vector>
#include <cmath>

/* -------------------------------------------------------------------------- */
/*  PyTorch-facing function                                                   */
/* -------------------------------------------------------------------------- */
std::vector<torch::Tensor> run(
        torch::Tensor q,
        torch::Tensor k_cache,
        torch::Tensor v_cache,
        torch::Tensor qo_indptr,
        torch::Tensor kv_indptr,
        torch::Tensor kv_indices,
        double        sm_scale_double = 1.0 / std::sqrt(128.0))
{
    TORCH_CHECK(q.is_cuda(), "All tensors must be on the same CUDA device");
    auto device = q.device();

    const int64_t total_q = q.size(0);

    auto output = torch::empty({total_q, NUM_QO_HEADS, HEAD_DIM},
                               torch::TensorOptions()
                                   .dtype(torch::kBFloat16)
                                   .device(device));

    auto lse = torch::empty({total_q, NUM_QO_HEADS},
                            torch::TensorOptions()
                                .dtype(torch::kFloat32)
                                .device(device));

    /* Reference implementation initialises with zeros / -INF – replicate */
    output.zero_();
    lse.fill_(-INFINITY);

    gqa_paged_prefill_causal_h32_kv8_d128_ps1_launcher(
        q, k_cache, v_cache,
        qo_indptr, kv_indptr, kv_indices,
        static_cast<float>(sm_scale_double),
        output, lse);

    return {output, lse};
}

/* -------------------------------------------------------------------------- */
/*  PyBind11 module                                                           */
/* -------------------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run,
          "gqa_paged_prefill_causal_h32_kv8_d128_ps1 (B200-optimised)",
          pybind11::arg("q"),
          pybind11::arg("k_cache"),
          pybind11::arg("v_cache"),
          pybind11::arg("qo_indptr"),
          pybind11::arg("kv_indptr"),
          pybind11::arg("kv_indices"),
          pybind11::arg("sm_scale") = 1.0 / std::sqrt(128.0));
}
scrolls · 60 lines total

Source code from FlashInfer-Bench (flashinfer-ai/flashinfer-trace) · Apache-2.0

Best evidence level for this revision: reported

JSON