Repository navigation
support indexer_n_heads = 16 - #369
Conversation
Signed-off-by: JaredforReal <w13431838023@gmail.com>
There was a problem hiding this comment.
Pull request overview
This PR extends the SM100 MQA logits implementations to support num_heads == 16 by adjusting TMEM load paths for smaller accumulator vectors and padding the shared-memory weight staging layout to satisfy TMA alignment requirements. It also updates the JIT-side shared-memory sizing and host-side input validation to allow the new head count.
Changes:
- Pad per-stage shared-memory stride for the weights pipeline in SM100 paged kernels to keep TMA destination addresses aligned.
- Extend TMEM load helper logic to support 8x/16x loads needed by
num_heads == 16(and FP4’snum_heads/2path). - Update JIT shared-memory size calculations and API assertions to accept
num_heads == 16.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| deep_gemm/include/deep_gemm/impls/sm100_fp8_paged_mqa_logits.cuh | Pads SMEM weight stage stride for TMA alignment; extends TMEM load sizes. |
| deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits.cuh | Extends TMEM load sizes to support 16-head variants. |
| deep_gemm/include/deep_gemm/impls/sm100_fp4_paged_mqa_logits.cuh | Pads SMEM weight stage stride for TMA alignment; extends TMEM load sizes (incl. 8x). |
| deep_gemm/include/deep_gemm/impls/sm100_fp4_mqa_logits.cuh | Extends TMEM load sizes to support 16-head FP4 paths (kNumHeads/2 == 8). |
| csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp | Updates SM100 shared-memory size accounting for padded weight stride. |
| csrc/apis/attention.hpp | Updates host assertions to allow num_heads == 16 for relevant entry points. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = kNextNAtom * kNumHeads * sizeof(float); | ||
| // TMA requires 128B - aligned shared memory addresses, so pad the per - stage stride | ||
| // (e.g. `kNextNAtom == 1` with 16 heads gives only 64B per stage) | ||
| static constexpr uint32_t SMEM_WEIGHT_STRIDE_PER_STAGE = SMEM_WEIGHT_SIZE_PER_STAGE < 128 ? 128 : SMEM_WEIGHT_SIZE_PER_STAGE; |
| static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = kNextNAtom * kNumHeads * sizeof(float); | ||
| // TMA requires 128B - aligned shared memory addresses, so pad the per - stage stride | ||
| // (e.g. `kNextNAtom == 1` with 16 heads gives only 64B per stage) | ||
| static constexpr uint32_t SMEM_WEIGHT_STRIDE_PER_STAGE = SMEM_WEIGHT_SIZE_PER_STAGE < 128 ? 128 : SMEM_WEIGHT_SIZE_PER_STAGE; |
🤖 ds-review-bot Code ReviewModel 1本 MR 为 MQA logits 相关 kernel 添加 num_heads == 16 的支持（迁移自 vllm-project/pull/1）。核心改动有三处且相互一致：1) 放宽 attention.hpp 中 num_heads 的 host 断言，允许 16/32/64；2) 为 SM100 分页 kernel（FP8 与 FP4）将每 stage 的 weight 共享内存步长补齐到 128B，以满足 TMA 目的地址 128B 对齐要求——这修复了 num_heads==16 且 next_n_per_cta==1 时每 stage 仅 64B 的对齐问题（host 分配与 device 布局同步对齐，arrive_and_expect_tx 仍按实际写入的 SMEM_WEIGHT_SIZE_PER_STAGE 计，正确）；3) 给 TMEM load 增加 N==8 分支（SM100_TMEM_LOAD_32dp32b8x，cutlass 中确实存在），FP4 reduce 路径 kNumHeads/2==8 时正需此分支。非分页路径因 block_q=128/num_heads 使 weight stage 恒为 512B（128B 对齐），无需改动；SM120 各 kernel 已通过 constexpr_align 将 weight stage 对齐到 512/1024，且 MMA_N=8 可覆盖 16 heads，故本次未触及。整体实现正确、改动范围聚焦，未发现 bug。 Model 2本 MR 为 MQA logits 系列内核新增了 indexer_n_heads = 16 的支持（从 DeepGEMM 上游 PR 迁移）。改动包含三部分：(1) 在 attention.hpp 的 4 处断言中放宽 num_heads 取值，允许 16；(2) 在 fp4/fp8、paged/非 paged 的 4 个内核的 tmem_load lambda 中扩展支持 N==8 与 N==16，复用已存在的 cute::SM100_TMEM_LOAD_32dp32b8x/16x；(3) 在两个 paged 内核中将 weight 的 per-stage 步长按 128B 对齐（16 heads 且 kNextNAtom==1 时仅 64B，不满足 TMA 对齐要求），并同步更新 host 端 smem 大小计算。 整体评估：改动逻辑自洽且边界处理正确。已核对 weight TMA 的 expected-tx 字节数仍使用未填充的 SMEM_WEIGHT_SIZE_PER_STAGE（TMA 只写入实际数据），而 smem_weights 步长与后续 barrier_ptr 偏移使用填充后的 SMEM_WEIGHT_STRIDE_PER_STAGE，host 端与 device 端的 smem 计算一致；非 paged 内核因 BLOCK_Q*kNumHeads 恒为 128（block_q=128/num_heads）始终 128B 对齐故无需填充，仅做了 tmem_load 扩展，处理正确。在所有合法的 next_n_per_cta∈{1,2} 与 heads∈{16,32,64} 组合下，host 端 align(SIZE,128) 与 device 端 (SIZE<128?128:SIZE) 结果完全一致，未发现 bug。 Files reviewed: 6 |
Brings deepseek-ai PR #377 (559d79f, including #364 and #343) onto nv_dev (a6b593d = #369 merge). Merge base is 88965b0. Conflict resolution highlights: * MQA: adopt #377's unified SM100 implementation and common MMA/PTX refactor while preserving nv_dev's SM90 next_n=4 multicast and block_kv=32 support, dedicated SM120 paths, FP16/BF16 weights, and 16-head indexer coverage. * Scheduler safety: retain zero-context and empty-varlen handling; drive producer prefetch from the scheduler's actual next query; bound the SM90 block_kv=32 odd-page tail before reading block tables. * FP16 weights: port the implementation to the new headers and preserve padded host offsets, per-warpgroup row indexing, and per-row compressed window bounds. * GEMM/layout: combine #377's grouped scale-factor, PSUM, zero-padding, and Mega-MoE changes with nv_dev's SM120 dense/grouped dispatch. * Tests: merge upstream coverage with nv_dev's SM90/SM100/SM120 cases and add regressions for empty metadata, skipped requests, FP16 compressed windows, and the SM90 odd physical-page tail. Validation: * All Git-tracked Python files pass py_compile. * No unresolved conflicts, conflict markers, missing project includes, or references to removed headers were found. * CUDA build/JIT and GPU numerical tests were not run because this host has no nvcc, PyTorch, or visible GPU tooling.
migrated from vllm-project#1