Skip to content

[SM100] Paged MQA logits: optional exact coarse score histogram for LiteTopK decode #94

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Heisenberg-Yin wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

[SM100] Paged MQA logits: optional exact coarse score histogram for LiteTopK decode #94

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 19 additions & 4 deletions csrc/apis/attention.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,8 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st
const torch::Tensor& block_table,
const torch::Tensor& schedule_meta,
const int& max_context_len,
const std::optional<torch::Tensor>& indices) {
const std::optional<torch::Tensor>& indices,
const std::optional<torch::Tensor>& histogram) {
const auto [q_fp, q_sf] = q;
const auto qk_dtype = q_fp.scalar_type();
const bool is_fp4 = qk_dtype == kPackedFP4;
Expand Down Expand Up @@ -239,6 +240,20 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st
DG_HOST_ASSERT(indices_tensor.scalar_type() == torch::kInt);
}

// Check the optional coarse histogram: every stored non-NaN BF16 score is added to int32 `[rows, 1024]`
int* histogram_ptr = nullptr;
if (histogram.has_value()) {
DG_HOST_ASSERT(arch_major == 10 and "The paged MQA logits histogram requires SM100");
DG_HOST_ASSERT((num_heads == 32 or num_heads == 64) and "The paged MQA logits histogram supports 32 or 64 heads");
DG_HOST_ASSERT(histogram->device() == q_fp.device());
DG_HOST_ASSERT(histogram->scalar_type() == torch::kInt and histogram->is_contiguous());
DG_HOST_ASSERT(histogram->dim() == 2 and histogram->size(0) == batch_size * next_n and
histogram->size(1) == kMQALogitsHistogramBins);
// Bin pairs are added as 64-bit words
DG_HOST_ASSERT(reinterpret_cast<uintptr_t>(histogram->data_ptr()) % 8 == 0);
histogram_ptr = histogram->data_ptr<int>();
}

// Check schedule metadata
auto [_schedule_meta_size, _meta_info_size] = get_shape<2>(schedule_meta);
DG_HOST_ASSERT(_schedule_meta_size == num_sms + 1 and _meta_info_size == 2);
Expand All @@ -265,10 +280,10 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st

// Dispatch implementation
if (arch_major == 10) {
const MQALogitsConfig config(num_heads, head_dim, qk_dtype, true);
const MQALogitsConfig config(num_heads, head_dim, qk_dtype, true, histogram.has_value());
sm100_paged_mqa_logits(config, q_fp, q_sf.value(), kv_cache, kv_cache_sf, weights, context_lens, logits,
block_table, indices_tensor, schedule_meta, batch_size, num_kv_blocks, block_kv,
aligned_max_context_len, block_table_stride);
aligned_max_context_len, block_table_stride, histogram_ptr);
} else if (arch_major == 9) {
DG_HOST_ASSERT(qk_dtype == torch::kFloat8_e4m3fn);
DG_HOST_ASSERT(num_heads == 32 or num_heads == 64);
Expand Down Expand Up @@ -459,7 +474,7 @@ static void register_apis(pybind11::module_& m) {
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("indices") = std::nullopt, py::arg("histogram") = std::nullopt);
m.def("get_sparse_mqa_logits_metadata", &get_sparse_mqa_logits_metadata,
py::arg("cu_seq_len_k_start"), py::arg("cu_seq_len_k_end"), py::arg("num_kv_tokens"),
py::arg("sparse_kv_block_indices"), py::arg("qk_dtype"), py::arg("sparse_block_kv"),
Expand Down
44 changes: 33 additions & 11 deletions csrc/jit_kernels/impls/sm100_mqa_logits.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ static constexpr int kMQALogitsSplitKV = 384;
static constexpr int kMQALogitsNumTmemStages = 3;
static constexpr int kMQALogitsSplitsPerChunk = 8;
static constexpr int kMQALogitsNumSpecializedThreads = 128;
// Bins per row of the optional paged histogram (`epilogue::kNumCoarseHistogramBins`)
static constexpr int kMQALogitsHistogramBins = 1024;

// The Q tile always covers 128 (token, head) rows
static int get_mqa_logits_block_q(const int& num_heads) {
Expand Down Expand Up @@ -117,17 +119,18 @@ static void __instantiate_kernel() {{
// Kernel geometry shared by the contiguous and paged runtimes: every math warpgroup consumes one UMMA_M = 128 KV
// tile, the Q tile always covers 128 (token, head) rows, and the FP4 pipelines are deeper (half-size stages)
struct MQALogitsConfig {
bool is_fp4, is_paged;
bool is_fp4, is_paged, with_histogram;
int num_heads, head_dim, block_q, umma_n;
int num_q_stages, num_kv_stages;
static constexpr int split_kv = kMQALogitsSplitKV;
static constexpr int num_tmem_stages = kMQALogitsNumTmemStages;
static constexpr int num_specialized_threads = kMQALogitsNumSpecializedThreads;
static constexpr int num_math_threads = (kMQALogitsSplitKV / 128) * 128;

MQALogitsConfig(const int& num_heads, const int& head_dim, const at::ScalarType& qk_dtype, const bool& is_paged):
is_fp4(qk_dtype == kPackedFP4), is_paged(is_paged), num_heads(num_heads), head_dim(head_dim),
block_q(get_mqa_logits_block_q(num_heads)) {
MQALogitsConfig(const int& num_heads, const int& head_dim, const at::ScalarType& qk_dtype, const bool& is_paged,
const bool& with_histogram = false):
is_fp4(qk_dtype == kPackedFP4), is_paged(is_paged), with_histogram(with_histogram),
num_heads(num_heads), head_dim(head_dim), block_q(get_mqa_logits_block_q(num_heads)) {
DG_HOST_ASSERT(qk_dtype == kPackedFP4 or qk_dtype == torch::kFloat8_e4m3fn);
if (is_fp4)
DG_HOST_ASSERT(head_dim == 64 or head_dim == 128);
Expand All @@ -136,6 +139,13 @@ struct MQALogitsConfig {
umma_n = align(block_q * num_heads, 8);
num_q_stages = is_fp4 ? 2 : 3;
num_kv_stages = is_fp4 ? 8 : 3;
if (with_histogram) {
// The window follows the shared storage: keep the deepest KV ring that leaves room for it (MXFP4 D128: 7)
while (smem_size() + histogram_smem_size() > SM100ArchSpec::smem_capacity) {
DG_HOST_ASSERT(num_kv_stages > 2);
-- num_kv_stages;
}
}
}

const char* qk_dtype_name() const {
Expand All @@ -157,6 +167,16 @@ struct MQALogitsConfig {
return static_cast<int>(align(num_heads * static_cast<uint32_t>(sizeof(nv_bfloat16)), 16u) / sizeof(nv_bfloat16));
}

// FP4 chunks span the KV ring, so a chunk stays resident for the other Q blocks of its request
int splits_per_chunk() const {
return is_fp4 ? num_kv_stages : kMQALogitsSplitsPerChunk;
}

// Histogram window: two Q blocks of rows, 1024 16-bit counters each (`epilogue::RowWindowHistogram`)
int histogram_smem_size() const {
return with_histogram ? 2 * block_q * kMQALogitsHistogramBins * static_cast<int>(sizeof(uint16_t)) : 0;
}

// Template arguments of `layout::MQALogitsSharedStorage`; the generated source asserts `sizeof` == `smem_size()`
std::string shared_storage_args() const {
return std::format("{}, {}, {}, {}, {}, {}, {}, {}, {}",
Expand Down Expand Up @@ -309,8 +329,10 @@ static void sm100_paged_mqa_logits(const MQALogitsConfig& config,
const int& num_q_tokens_total,
const int& num_kv_blocks, const int& page_kv,
const int& logits_stride,
const int& block_table_stride) {
const int& block_table_stride,
int* histogram) {
DG_HOST_ASSERT(config.is_paged);
DG_HOST_ASSERT(config.with_histogram == (histogram != nullptr));
DG_HOST_ASSERT(logits_stride % config.split_kv == 0);
const int num_sms = runtime->get_num_sms();
DG_HOST_ASSERT(schedule_meta.size(0) == num_sms + 1);
Expand All @@ -333,7 +355,7 @@ static void sm100_paged_mqa_logits(const MQALogitsConfig& config,
const auto tensor_map_weights = config.make_tensor_map_weights(weights, num_q_tokens_total);
const int smem_size = config.smem_size();

// Compile
// Compile (the histogram variant appends `kWithHistogram`)
const auto kernel = MQALogitsConfig::compile("sm100_paged_mqa_logits", std::format(R"(
#include <deep_gemm/impls/sm100_mqa_logits.cuh>

Expand All @@ -350,7 +372,7 @@ static void __instantiate_kernel() {{
{}, {}, {},
{}, {},
{}, {},
{}
{}{}
>);
}};
)",
Expand All @@ -359,14 +381,14 @@ static void __instantiate_kernel() {{
page_kv,
config.block_q, config.umma_n,
config.num_q_stages, config.num_kv_stages, config.num_tmem_stages,
config.split_kv, kMQALogitsSplitsPerChunk,
config.split_kv, config.splits_per_chunk(),
config.num_specialized_threads, config.num_math_threads,
config.qk_dtype_name()));
config.qk_dtype_name(), config.with_histogram ? ", true" : ""));

// Launch
jit->launch(
kernel, {
.num_smem_bytes = smem_size,
.num_smem_bytes = smem_size + config.histogram_smem_size(),
.grid_dim = dim3(num_sms, 1, 1),
.block_dim = dim3(config.num_specialized_threads + config.num_math_threads, 1, 1),
},
Expand All @@ -376,7 +398,7 @@ static void __instantiate_kernel() {{
block_table.data_ptr<int>(), indices.data_ptr<int>(), schedule_meta.data_ptr<int>(),
tensor_map_q, tensor_map_sf_q,
tensor_map_kv, tensor_map_sf_kv,
tensor_map_weights
tensor_map_weights, histogram
);
}

Expand Down
7 changes: 4 additions & 3 deletions csrc/torch_library_api.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -420,10 +420,10 @@ Tensor dg_fp8_fp4_mqa_logits(const Tensor& q, const OptTensor& q_sf, const Tenso
Tensor dg_fp8_fp4_paged_mqa_logits(const Tensor& q, const OptTensor& q_sf, const Tensor& fused_kv_cache,
const Tensor& weights, const Tensor& context_lens,
const Tensor& block_table, const Tensor& schedule_meta,
int64_t max_context_len, const OptTensor& indices) {
int64_t max_context_len, const OptTensor& indices, const OptTensor& histogram) {
return attention::fp8_fp4_paged_mqa_logits(
std::make_tuple(q, q_sf), fused_kv_cache, weights, context_lens, block_table, schedule_meta,
static_cast<int>(max_context_len), indices);
static_cast<int>(max_context_len), indices, histogram);
}

Tensor dg_get_mqa_logits_metadata(const Tensor& cu_seq_len_k_start, const Tensor& cu_seq_len_k_end,
Expand Down Expand Up @@ -785,7 +785,8 @@ TORCH_LIBRARY(deep_gemm, m) {
m.def("fp8_fp4_mqa_logits(Tensor q, Tensor? q_sf, Tensor kv_data, Tensor kv_sf, Tensor weights, "
"Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, int max_seqlen_k, Tensor? schedule_meta) -> Tensor");
m.def("fp8_fp4_paged_mqa_logits(Tensor q, Tensor? q_sf, Tensor fused_kv_cache, Tensor weights, "
"Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, Tensor? indices) -> Tensor");
"Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, Tensor? indices, "
"Tensor(a!)? histogram) -> Tensor");
m.def("get_mqa_logits_metadata(Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, int num_kv_tokens, "
"int num_heads) -> Tensor");
m.def("get_sparse_mqa_logits_metadata(Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, int num_kv_tokens, "
Expand Down
4 changes: 2 additions & 2 deletions deep_gemm/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,9 +142,9 @@ def fp8_einsum(expr, a, b, d, c=None, recipe=(1, 128, 128), epilogue=None):
d_data, d_sf = d if isinstance(d, (tuple, list)) else (d, None)
_C.fp8_einsum(expr, a_data, a_sf, b_data, b_sf, d_data, d_sf, c, recipe, *_encode_epilogue(epilogue))

def fp8_fp4_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, indices=None):
def fp8_fp4_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, indices=None, histogram=None):
(q, q_sf) = q if isinstance(q, (tuple, list)) else (q, None)
return _C.fp8_fp4_paged_mqa_logits(q, q_sf, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, indices)
return _C.fp8_fp4_paged_mqa_logits(q, q_sf, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, indices, histogram)

def fp8_fp4_mqa_logits(q, kv, weights, cu_seq_len_k_start, cu_seq_len_k_end, max_seqlen_k, schedule_meta=None):
(q, q_sf), (kv_data, kv_sf) = (q if isinstance(q, (tuple, list)) else (q, None)), _parse_tensor_or_tuple(kv)
Expand Down
Loading