[SM100] Add exact BF16 paged MQA histograms for next-n 1–6

README.md

@@ -184,6 +184,9 @@ Each `DG_JIT_*` variable falls back to the corresponding global `DJ_JIT_*` varia
 
 For additional examples and details, please refer to [the test code](tests) or review the corresponding Python documentation.
 
+The optional SM100 MXFP4 producer with exact BF16 score histograms and automatic
+Q4/Q6 routing is documented in [Paired BF16 paged MQA](docs/mqa_bf16.md).
+
 ## Acknowledgement
 
 DeepGEMM is inspired by the [CUTLASS](https://github.com/nvidia/cutlass) project. Thanks and respect to the developers!

csrc/apis/attention.hpp

@@ -17,6 +17,7 @@
 #include "../jit_kernels/impls/sm100_sparse_mqa_logits.hpp"
 #include "../jit_kernels/impls/sm90_fp8_mqa_logits.hpp"
 
+#include "attention_bf16.hpp"
 #include "layout.hpp"
 #include "../jit_kernels/impls/smxx_clean_logits.hpp"
 
@@ -624,6 +625,14 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q,
 #ifndef DG_USE_TVM_FFI
 
 static void register_apis(pybind11::module_& m) {
+    m.def("get_paged_mqa_logits_bf16_metadata", &get_paged_mqa_logits_bf16_metadata,
+          py::arg("context_lens"), py::arg("block_kv"), py::arg("num_sms"),
+          py::arg("indices") = std::nullopt, py::arg("tokens_per_request") = 1);
+    m.def("fp4_paged_mqa_logits_bf16", &fp4_paged_mqa_logits_bf16,
+          py::arg("q"), py::arg("kv_cache"), py::arg("weights"),
+          py::arg("context_lens"), py::arg("block_table"), py::arg("schedule_meta"),
+          py::arg("max_context_len"), py::arg("indices") = std::nullopt,
+          py::arg("histogram") = std::nullopt, py::arg("tokens_per_request") = 1);
     m.def("fp8_gemm_nt_skip_head_mid", &fp8_gemm_nt_skip_head_mid,
           py::arg("a"), py::arg("b"), py::arg("d"), py::arg("head_splits"),
           py::arg("recipe") = std::nullopt,

csrc/apis/attention_bf16.hpp

@@ -0,0 +1,99 @@
+#pragma once
+
+#include "../jit_kernels/impls/sm100_mqa_logits_bf16.hpp"
+
+namespace deep_gemm::attention {
+
+static void check_bf16_int_tensor(const torch::Tensor& t, const torch::Device& device) {
+    DG_HOST_ASSERT(t.is_cuda() and t.device() == device);
+    DG_HOST_ASSERT(t.scalar_type() == torch::kInt and t.is_contiguous());
+    DG_HOST_ASSERT(reinterpret_cast<uintptr_t>(t.data_ptr()) % 8 == 0);
+}
+
+static torch::Tensor bf16_request_indices(const torch::Tensor& lens,
+                                          const std::optional<torch::Tensor>& indices) {
+    if (indices.has_value()) {
+        check_bf16_int_tensor(*indices, lens.device());
+        DG_HOST_ASSERT(indices->dim() == 1 and indices->size(0) == lens.size(0));
+        return *indices;
+    }
+    // Without request IDs every row is independent, regardless of the tile hint.
+    return torch::arange(lens.size(0), lens.options());
+}
+
+static torch::Tensor get_paged_mqa_logits_bf16_metadata(const torch::Tensor& context_lens,
+                                                       int block_kv, int num_sms,
+                                                       const std::optional<torch::Tensor>& indices,
+                                                       int tokens_per_request) {
+    DG_HOST_ASSERT(jit->device.get_arch_major() == 10 and block_kv == 128);
+    DG_HOST_ASSERT(num_sms > 0 and num_sms <= runtime->get_num_sms());
+    check_bf16_int_tensor(context_lens, context_lens.device());
+    DG_HOST_ASSERT(context_lens.dim() == 2 and context_lens.size(1) == 1);
+    auto ids = bf16_request_indices(context_lens, indices);
+    auto schedule = torch::empty({num_sms + 1, 2}, context_lens.options());
+    mqa_bf16::metadata(context_lens, ids, schedule, num_sms, tokens_per_request);
+    return schedule;
+}
+
+static torch::Tensor fp4_paged_mqa_logits_bf16(const std::pair<torch::Tensor, torch::Tensor>& q,
+                                              const torch::Tensor& fused_kv_cache,
+                                              const torch::Tensor& weights,
+                                              const torch::Tensor& context_lens,
+                                              const torch::Tensor& block_table,
+                                              const torch::Tensor& schedule_meta,
+                                              int max_context_len,
+                                              const std::optional<torch::Tensor>& indices,
+                                              const std::optional<torch::Tensor>& histogram,
+                                              int tokens_per_request) {
+    const auto& [q_fp, q_sf] = q;
+    const auto device = q_fp.device();
+    const mqa_bf16::Config config(tokens_per_request, histogram.has_value());
+    DG_HOST_ASSERT(jit->device.get_arch_major() == 10 and q_fp.is_cuda());
+    DG_HOST_ASSERT(q_fp.dim() == 4 and q_fp.size(1) == 1 and q_fp.size(2) == 32 and q_fp.size(3) == 64);
+    DG_HOST_ASSERT(q_fp.scalar_type() == kPackedFP4 and q_fp.is_contiguous());
+    const int rows = q_fp.size(0);
+    check_bf16_int_tensor(q_sf, device);
+    DG_HOST_ASSERT(q_sf.dim() == 3 and q_sf.size(0) == rows and q_sf.size(1) == 1 and q_sf.size(2) == 32);
+    DG_HOST_ASSERT(weights.is_cuda() and weights.device() == device and weights.scalar_type() == torch::kBFloat16);
+    DG_HOST_ASSERT(weights.dim() == 2 and weights.size(0) == rows and weights.size(1) == 32);
+    DG_HOST_ASSERT(weights.stride(1) == 1 and weights.stride(0) % 8 == 0);
+    DG_HOST_ASSERT(reinterpret_cast<uintptr_t>(weights.data_ptr()) % 16 == 0);
+
+    check_bf16_int_tensor(context_lens, device);
+    DG_HOST_ASSERT(context_lens.dim() == 2 and context_lens.size(0) == rows and context_lens.size(1) == 1);
+    DG_HOST_ASSERT(max_context_len >= 0);
+    DG_HOST_ASSERT(block_table.is_cuda() and block_table.device() == device and block_table.scalar_type() == torch::kInt);
+    DG_HOST_ASSERT(block_table.dim() == 2 and block_table.size(0) == rows and block_table.stride(1) == 1);
+    DG_HOST_ASSERT(block_table.size(1) >= math::ceil_div(max_context_len, 128));
+    check_bf16_int_tensor(schedule_meta, device);
+    DG_HOST_ASSERT(schedule_meta.dim() == 2 and schedule_meta.size(0) == runtime->get_num_sms() + 1 and schedule_meta.size(1) == 2);
+    auto ids = bf16_request_indices(context_lens, indices);
+
+    DG_HOST_ASSERT(fused_kv_cache.is_cuda() and fused_kv_cache.device() == device);
+    DG_HOST_ASSERT(fused_kv_cache.scalar_type() == torch::kByte and fused_kv_cache.dim() == 4);
+    DG_HOST_ASSERT(fused_kv_cache.size(1) == 128 and fused_kv_cache.size(2) == 1 and fused_kv_cache.size(3) == 68);
+    DG_HOST_ASSERT(fused_kv_cache.stride(1) == 68 and fused_kv_cache.stride(3) == 1);
+    DG_HOST_ASSERT(fused_kv_cache.stride(0) >= 128 * 68 and fused_kv_cache.stride(0) % 16 == 0);
+    DG_HOST_ASSERT(reinterpret_cast<uintptr_t>(fused_kv_cache.data_ptr()) % 16 == 0);
+    const int blocks = fused_kv_cache.size(0), kv_stride = fused_kv_cache.stride(0);
+    auto kv = torch::from_blob(fused_kv_cache.data_ptr(), {blocks, 128, 64}, {kv_stride, 64, 1},
+                                q_fp.options());
+    auto sf_kv = torch::from_blob(fused_kv_cache.data_ptr<uint8_t>() + 128 * 64,
+                                   {blocks, 128}, {kv_stride / 4, 1}, q_sf.options());
+
+    int* hist_ptr = nullptr;
+    if (histogram.has_value()) {
+        check_bf16_int_tensor(*histogram, device);
+        DG_HOST_ASSERT(histogram->dim() == 2 and histogram->size(0) == rows and histogram->size(1) == 1024);
+        hist_ptr = histogram->data_ptr<int>();
+    }
+    // LCM(384, 512): whole KV splits and 1024-byte aligned BF16 rows.
+    const int stride = align(max_context_len, 1536);
+    auto logits = torch::empty({rows, stride}, weights.options()).slice(1, 0, max_context_len);
+    if (rows > 0 and max_context_len > 0)
+        mqa_bf16::paged(config, q_fp, q_sf, kv, sf_kv, weights, context_lens, block_table, ids,
+                         schedule_meta, logits, hist_ptr);
+    return logits;
+}
+
+} // namespace deep_gemm::attention

csrc/jit_kernels/impls/sm100_mqa_logits_bf16.hpp

@@ -0,0 +1,111 @@
+#pragma once
+
+#include <format>
+
+#include "../../runtime/runtime.hpp"
+#include "../heuristics/sm100.hpp"
+#include "runtime_utils.hpp"
+#include <deep_gemm/mqa_bf16/layout.cuh>
+
+namespace deep_gemm::mqa_bf16 {
+
+// A paired API: the BF16 schedule uses split 384 and must never consume the
+// legacy paged API's split-256 metadata. The hint is CPU configuration; request
+// runs and causal lengths remain on the device.
+struct Config {
+    int block_q, num_tmem_stages, num_kv_stages;
+    bool histogram, swizzle;
+    static constexpr int split_kv = 384;
+
+    Config(int tokens_per_request, bool histogram):
+        block_q(tokens_per_request >= 5 ? 6 : 4),
+        num_tmem_stages(tokens_per_request >= 5 ? 2 : 3),
+        num_kv_stages(tokens_per_request >= 5 or histogram ? 7 : 8),
+        histogram(histogram), swizzle(tokens_per_request == 6) {
+        DG_HOST_ASSERT(tokens_per_request >= 1 and tokens_per_request <= 6);
+    }
+
+    int storage_size() const {
+        using FP4 = cutlass::float_e2m1_t;
+        if (block_q == 6)
+            return sizeof(layout::MQALogitsSharedStorage<32, 128, 6, 384, 192, 2, 7, 2, FP4>);
+        if (num_kv_stages == 7)
+            return sizeof(layout::MQALogitsSharedStorage<32, 128, 4, 384, 128, 2, 7, 3, FP4>);
+        return sizeof(layout::MQALogitsSharedStorage<32, 128, 4, 384, 128, 2, 8, 3, FP4>);
+    }
+
+    int smem_size() const {
+        const int bytes = storage_size() + (histogram ? 8 * 1024 * sizeof(uint16_t) : 0);
+        DG_HOST_ASSERT(bytes <= SM100ArchSpec::smem_capacity);
+        return bytes;
+    }
+};
+
+static void metadata(const torch::Tensor& lens, const torch::Tensor& indices,
+                     const torch::Tensor& schedule, int num_sms, int tokens_per_request) {
+    const int rows = lens.size(0);
+    const int threads = rows <= 512 ? 256 : (rows <= 2048 ? 512 : 1024);
+    const int smem = (2 * rows + threads / 32 + 1) * sizeof(int);
+    DG_HOST_ASSERT(smem <= SM100ArchSpec::smem_capacity);
+    const int block_q = Config(tokens_per_request, false).block_q;
+    const auto kernel = jit->compile("sm100_paged_mqa_logits_bf16_metadata", std::format(R"(
+#include <deep_gemm/mqa_bf16/scheduler.cuh>
+static void __instantiate_kernel() {{
+    auto ptr = reinterpret_cast<void*>(&deep_gemm::mqa_bf16::sched::sm100_paged_mqa_logits_metadata<384, {}, {}, {}>);
+}}
+)", num_sms, block_q, threads));
+    jit->launch(
+        kernel, {
+            .num_smem_bytes = smem,
+            .grid_dim = dim3(1, 1, 1),
+            .block_dim = dim3(threads, 1, 1),
+        },
+        rows, lens.data_ptr<int>(), indices.data_ptr<int>(), schedule.data_ptr<int>()
+    );
+}
+
+static void paged(const Config& c, const torch::Tensor& q, const torch::Tensor& sf_q,
+                  const torch::Tensor& kv, const torch::Tensor& sf_kv, const torch::Tensor& weights,
+                  const torch::Tensor& lens, const torch::Tensor& table, const torch::Tensor& indices,
+                  const torch::Tensor& schedule, const torch::Tensor& logits, int* histogram) {
+    const int rows = q.size(0), num_sms = runtime->get_num_sms();
+    const auto tensor_map_q = make_tma_2d_desc(
+        q, 128, rows * 32, 128, c.block_q * 32,
+        static_cast<int>(q.stride(2)), 64, 0, false, false);
+    const auto tensor_map_sf_q = make_tma_2d_desc(
+        sf_q, 32, rows, 32, c.block_q, static_cast<int>(sf_q.stride(1)), 0);
+    const auto tensor_map_kv = make_tma_3d_desc(
+        kv, 128, 128, kv.size(0), 128, 128, 1,
+        static_cast<int>(kv.stride(1)), static_cast<int>(kv.stride(0)),
+        64, 0, false, false);
+    const auto tensor_map_sf_kv = make_tma_2d_desc(
+        sf_kv, 128, sf_kv.size(0), 128, 1, static_cast<int>(sf_kv.stride(0)), 0);
+    const auto tensor_map_weights = make_tma_2d_desc(
+        weights, 32, rows, 32, c.block_q, static_cast<int>(weights.stride(0)), 0);
+    const auto kernel = jit->compile("sm100_paged_mqa_logits_bf16", std::format(R"(
+#include <deep_gemm/mqa_bf16/kernel.cuh>
+using namespace deep_gemm::mqa_bf16;
+static_assert(sizeof(layout::MQALogitsSharedStorage<32, 128, {}, 384, {}, 2, {}, {}, cutlass::float_e2m1_t>) == {});
+static void __instantiate_kernel() {{
+    auto ptr = reinterpret_cast<void*>(&sm100_paged_mqa_logits<
+        32, 128, 128, {}, {}, 2, {}, {}, 384, {}, 128, 384, cutlass::float_e2m1_t, {}, {}
+    >);
+}}
+)", c.block_q, c.block_q * 32, c.num_kv_stages, c.num_tmem_stages, c.storage_size(),
+    c.block_q, c.block_q * 32, c.num_kv_stages, c.num_tmem_stages, c.num_kv_stages,
+    c.histogram, c.swizzle));
+    jit->launch(
+        kernel, {
+            .num_smem_bytes = c.smem_size(),
+            .grid_dim = dim3(num_sms, 1, 1),
+            .block_dim = dim3(512, 1, 1),
+        },
+        rows, static_cast<int>(logits.stride(0)), static_cast<int>(table.stride(0)),
+        lens.data_ptr<int>(), logits.data_ptr(), table.data_ptr<int>(),
+        indices.data_ptr<int>(), schedule.data_ptr<int>(),
+        tensor_map_q, tensor_map_sf_q, tensor_map_kv, tensor_map_sf_kv,
+        tensor_map_weights, histogram
+    );
+}
+
+} // namespace deep_gemm::mqa_bf16

csrc/tvm_ffi_api.cpp

@@ -601,6 +601,16 @@ Tensor dg_get_paged_mqa_logits_metadata(TensorView context_lens, int64_t block_k
     return Tensor::FromDLPack(at::toDLPack(result));
 }
 
+Tensor dg_get_paged_mqa_logits_bf16_metadata(TensorView context_lens, int64_t block_kv,
+                                             int64_t num_sms, Optional<TensorView> indices,
+                                             int64_t tokens_per_request) {
+    auto result = attention::get_paged_mqa_logits_bf16_metadata(
+        convert_to_torch_tensor(context_lens), static_cast<int>(block_kv),
+        static_cast<int>(num_sms), to_optional_tensor(indices),
+        static_cast<int>(tokens_per_request));
+    return Tensor::FromDLPack(at::toDLPack(result));
+}
+
 Tensor dg_fp8_paged_mqa_logits(TensorView q, TensorView fused_kv_cache,
                               TensorView weights, TensorView context_lens,
                               TensorView block_table, TensorView schedule_meta,
@@ -615,6 +625,23 @@ Tensor dg_fp8_paged_mqa_logits(TensorView q, TensorView fused_kv_cache,
     return Tensor::FromDLPack(at::toDLPack(result));
 }
 
+Tensor dg_fp4_paged_mqa_logits_bf16(TensorView q, TensorView q_sf,
+                                    TensorView fused_kv_cache, TensorView weights,
+                                    TensorView context_lens, TensorView block_table,
+                                    TensorView schedule_meta, int64_t max_context_len,
+                                    Optional<TensorView> indices,
+                                    Optional<TensorView> histogram,
+                                    int64_t tokens_per_request) {
+    auto result = attention::fp4_paged_mqa_logits_bf16(
+        std::make_pair(convert_to_torch_tensor(q), convert_to_torch_tensor(q_sf)),
+        convert_to_torch_tensor(fused_kv_cache), convert_to_torch_tensor(weights),
+        convert_to_torch_tensor(context_lens), convert_to_torch_tensor(block_table),
+        convert_to_torch_tensor(schedule_meta), static_cast<int>(max_context_len),
+        to_optional_tensor(indices), to_optional_tensor(histogram),
+        static_cast<int>(tokens_per_request));
+    return Tensor::FromDLPack(at::toDLPack(result));
+}
+
 Tensor dg_fp8_fp4_mqa_logits(TensorView q, Optional<TensorView> q_sf, TensorView kv_data, TensorView kv_sf,
                             TensorView weights, TensorView cu_seq_len_k_start,
                             TensorView cu_seq_len_k_end, bool clean_logits, int64_t max_seqlen_k,
@@ -649,7 +676,9 @@ Tensor dg_fp8_fp4_paged_mqa_logits(TensorView q, Optional<TensorView> q_sf, Tens
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_gemm_nt_skip_head_mid, dg_fp8_gemm_nt_skip_head_mid);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_mqa_logits, dg_fp8_mqa_logits);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_paged_mqa_logits_metadata, dg_get_paged_mqa_logits_metadata);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_paged_mqa_logits_bf16_metadata, dg_get_paged_mqa_logits_bf16_metadata);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_paged_mqa_logits, dg_fp8_paged_mqa_logits);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp4_paged_mqa_logits_bf16, dg_fp4_paged_mqa_logits_bf16);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_mqa_logits, dg_fp8_fp4_mqa_logits);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_paged_mqa_logits, dg_fp8_fp4_paged_mqa_logits);
 

deep_gemm/__init__.py

@@ -149,6 +149,18 @@ def fp8_fp4_mqa_logits(q, kv, weights, cu_seq_len_k_start, cu_seq_len_k_end, cle
     def get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms, indices=None):
         return _from_dlpack_if_needed(_C.get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms, indices))
 
+    def get_paged_mqa_logits_bf16_metadata(context_lens, block_kv, num_sms, indices=None, tokens_per_request=1):
+        return _from_dlpack_if_needed(_C.get_paged_mqa_logits_bf16_metadata(
+            context_lens, block_kv, num_sms, indices, tokens_per_request))
+
+    def fp4_paged_mqa_logits_bf16(q, kv_cache, weights, context_lens, block_table,
+                                  schedule_meta, max_context_len, indices=None,
+                                  histogram=None, tokens_per_request=1):
+        q_data, q_sf = q
+        return _from_dlpack_if_needed(_C.fp4_paged_mqa_logits_bf16(
+            q_data, q_sf, kv_cache, weights, context_lens, block_table,
+            schedule_meta, max_context_len, indices, histogram, tokens_per_request))
+
     def tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits=None):
         _C.tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits)