[sglang] logits를 쓰지 않는 top-k: SGLang의 LiteTopK prefill indexer로 DSA 롱컨텍스트 TTFT 최대 2배 단축

PR 링크: sgl-project/sglang#47469 상태: Merged | 변경: +6436 / -8

들어가며

DSA(DeepSeek Sparse Attention) 계열 모델(GLM-5.x, DeepSeek-V4.1)은 prefill 시 각 query row마다 indexer score를 모든 visible key에 대해 계산한 뒤, 그중 상위 topk(512~2048)개의 key만 골라 attention을 수행합니다. 기존 SGLang 경로는 다음과 같았습니다.

1. DeepGEMM mqa_logits가 [rows, keys] 크기의 전체 logits 텐서를 global memory에 기록
2. 별도의 top-k 커널(DeepSelect, top-k v2, 혹은 legacy topk_transform_prefill_kernel)이 이 logits를 다시 읽어 상위 k개 선택

문제는 롱컨텍스트입니다. 4096 rows × 1M keys만 되어도 logits는 수 GB 단위의 쓰기/읽기 트래픽을 만들고, top-k 연산량도 rows × keys에 비례해 커집니다. 결국 indexer가 GEMM이 아니라 메모리 트래픽과 selection에 묶이게 됩니다.

이 PR은 LiteTopK 아이디어를 SM100(Blackwell)에 구현해, logits를 아예 쓰지 않는 opt-in 경로(SGLANG_OPT_USE_PREFILL_LITE_TOPK=1)를 추가합니다.

핵심 아이디어: Sampled Threshold

하나의 커널이 tensor core로 모든 key의 score를 계산하되, 결과를 저장하지 않고 row별 threshold T_r 이상인 (score, key) 쌍만 per-row slab에 내보냅니다. T_r은 그 row의 score 중 strided sample에서 뽑으며, row당 약 2.25 × topk개의 후보만 남도록 조정됩니다. 그 뒤 작은 top-k 커널이 후보 집합에서만 최종 선택을 합니다.

Before: Q·K → logits[rows, keys] (HBM write) → top-k (HBM read) → indices
After : Q·K → (TMEM) → head reduce → score >= T_r ? → cand[rows, ~2.25k] → small top-k → indices

코드 분석

1. fp4_indexer.cuh — 점수 계산과 필터링의 융합

파일 헤더 주석이 설계를 잘 요약합니다.

// The scoring core follows DeepGEMM's `sm100_mqa_logits` MX path: a 256-key KV tile is the A operand of two
// M128 x N128 x K64 `kind::mxf4` MMAs per math warpgroup ..., the
// accumulator lives in TMEM, every math thread owns one key of the tile and reduces the heads in registers.
// Instead of writing logits, each row sets a threshold T_r from a strided sample of its own scores and only the
// (score, key) pairs at or above T_r reach a per-row global slab; blocks of short rows store every score instead.

스코어링 코어는 DeepGEMM과 동일한 MMA 구조를 쓰기 때문에, 유닛 테스트에서 DeepGEMM logits와 bitwise equal을 보장할 수 있습니다. 차이는 epilogue뿐입니다. 출력 파라미터를 보면 logits 버퍼가 없고 후보 slab만 있습니다.

void* __restrict__ cand_scores;       // [rows, kCap] bf16
int32_t* __restrict__ cand_indices;   // [rows, kCap] key - ks[row]
int32_t* __restrict__ cand_count;     // [rows] keys at or above T (may exceed kCap)
int32_t* __restrict__ cand_overflow;  // [rows] 1 when count > kCap
float* __restrict__ row_thresholds;   // [rows] T_r

warp 특화(warp specialization)도 눈에 띕니다. math warp는 reduce 후 bf16 score를 kDrainSlots짜리 staging ring에 넣기만 하고, select warp 2개가 sample tile로부터 threshold를 구하며, drain warp 2개가 main tile을 필터링해 slab에 씁니다.

// Main tiles [0, kMaskStart) of a block are filtered by the drains from the staged scores; from tile kMaskStart
// on the math warps compute the hit masks (they wait for the block's thresholds there, once). ...
// Measured: 4 at 32 heads, 8 at 64 heads.
static constexpr uint32_t kMaskStart = kNumHeads == 64 ? 8 : 4;

threshold가 결정되기 전의 앞쪽 타일은 staging된 score를 나중에 drain이 처리하고, 이후 타일은 math warp가 직접 hit mask를 계산합니다. threshold 계산 지연을 MMA 뒤에 숨기는 구조입니다. 레지스터 재분배를 위해 setmaxnreg.inc/dec PTX를 직접 노출하고, 분기 없는 emission을 위해 predicated store도 정의합니다.

// Predicated stores keep an emission branch-free.
SGL_DEVICE void st_global_if(bool pred, void* addr, uint32_t value) {
  asm volatile("{\n\t.reg .pred p;\n\tsetp.ne.u32 p, %0, 0;\n\t@p st.global.u32 [%1], %2;\n\t}" ...);
}

후보 필터링은 데이터 의존적이라 분기가 쉽게 diverge하는데, predicate store로 warp 내 제어 흐름을 균일하게 유지합니다.

2. 샘플링 방식: 8-key 블록 → 단일 key + 홀수 stride (리뷰 반영)

가장 흥미로운 기술적 논의는 리뷰어 Heisenberg-Yin의 제안입니다. 초기 구현은 8-key 블록 단위로 샘플했는데, 상위 key들은 연속된 run으로 몰려 있는 경향이 있어 블록이 run에 걸리거나 완전히 놓치면서 threshold가 noisy해졌습니다. 결과적으로 ~1M keys에서 일부 row가 slab cap에서 잘리거나 k개보다 적게 반환됐습니다.

제안은 단일 key를 홀수 stride로 샘플하고, strided tensor map으로 TMA 한 번에 tile을 로드하는 것이었고, 저자는 이를 채택했습니다. 그 흔적이 파라미터에 남아 있습니다.

// threshold samples: kv_fp4 as [keys / S, S, 64] (key = row * S + phase), 256 rows of one phase per box, for each
// odd sample stride S (see sample_stride)
CUtensorMap tmap_kv_strided[6];
...
uint32_t sample_limit[6];  // keys < limit[m] = (num_kv / S_m) * S_m are reachable through tmap_kv_strided[m]

KV를 [keys/S, S, 64] 3D 텐서로 재해석하면 phase 하나를 고정한 slice가 곧 stride-S 샘플이 되므로, gather 없이 3D TMA 한 번으로 샘플 tile을 가져올 수 있습니다. 결과는 FP8 short row 2→0, 최대 후보 10720→8422, 커널 시간 0.90~1.00x, FP4도 short row 12→0이었습니다. 홀수 stride와 hashed phase는 주기적 데이터와의 공명(resonance)을 줄이기 위한 선택입니다.

3. topk_page_table_transform.cuh — 1D grid로 수정

PAGED 출력을 위해 request-relative 위치를 page-size-64 테이블의 물리 KV slot으로 변환하는 작은 커널입니다.

const uint32_t q = blockIdx.x * blockDim.x + threadIdx.x;
const auto row = q / kNumQuads;
const auto q4 = q % kNumQuads;
PDLWaitPrimary<kUsePDL>();
if (q < params.total_quads) {
  ...
  AlignedVector<int32_t, 4> pos;
  pos.load(indices_row, q4);
#pragma unroll
  for (uint32_t i = 0; i < 4; ++i) {
    if (pos[i] >= 0) pos[i] = page_row[pos[i] / kPageSize] * kPageSize + pos[i] % kPageSize;
  }
  pos.store(indices_row, q4);
}
PDLTriggerSecondary<kUsePDL>();

리뷰에서 BBuf가 지적한 것은 초기 버전이 row를 grid.y에 두어 65536 rows(--chunked-prefill-size 65536)에서 launch 실패한다는 점이었습니다(grid.y 최대 65535). 수정 후 rows × topk/4 quad 전체를 flat 1D grid로 펴고 런타임 div/mod로 row를 복원합니다. B200에서 4096 rows 이상은 시간 변화가 0.7% 이하, 아주 작은 launch만 ~0.08us 비용이 듭니다. 또한 int4 단위 벡터 load/store, 16B alignment 검증, PDL(Programmatic Dependent Launch)로 앞 커널과의 겹침까지 챙겼습니다.

저자는 이 lookup을 top-k gather에 융합하자는 제안(0.97~0.98x)은 거절했는데, top-k가 KV layout에 독립적이도록 유지하려는 의도입니다. 작은 성능보다 모듈 경계를 택한 합리적 판단입니다.

4. 통합 및 라우팅

• DSA: 배치의 가장 긴 row가 8 × index_topk 이상의 key를 볼 때만 사용
• DeepSeek-V4.1: 압축 key가 24576개 이상일 때만 사용

~16K keys 미만에서는 필터가 대부분 key를 통과시켜 이득이 없기 때문입니다(BBuf 측정에서 FP8 8K는 0.62x). 리뷰 중 --enable-fp4-indexer 조합에서 BF16 query가 FP8 전용 커널로 들어가는 버그가 발견되어 Fp4Indexer는 use_prefill_lite_topk = False로 고정되었습니다.

왜 이게 좋은가

성능

B300 커널 기준(FP8, 32 heads, topk 2048, vs DeepGEMM + DeepSelect):

context x rows	ours	baseline	ratio
32K x 4096	482us	1181us	2.45
128K x 4096	1533us	3307us	2.16
1M x 16384	50161us	84329us	1.68

BBuf의 독립 GB300 측정에서는 GLM-5.2-FP8 1M 토큰 TTFT가 129.57s → 63.33s(2.05x), 256K는 1.40x였습니다. 반면 DeepSeek-V4.1-Flash는 full top-k 레이어가 3개뿐이라 TTFT 차이가 noise 수준(0.947~1.037x)입니다. 리뷰어는 또 baseline 선택이 중요하다고 짚었습니다: FP8 32K에서 DeepSelect 대비 2.33x지만 기존 top-k v2 대비로는 1.43x입니다. 벤치마크의 ours()가 candidate top-k를 빠뜨렸다는 지적도 있었죠. 공정한 비교 기준에 대한 좋은 교훈입니다.

Trade-off: 손실 가능한 최적화

이 경로는 정확한 top-k를 보장하지 않습니다. slab(kCap) 초과 시 truncation, 동일 bin 후보가 256개를 넘으면 잘리는 문제를 리뷰어가 제기했지만, 저자는 fallback이 non-trivial overhead를 만들고 실제로 거의 발생하지 않는다며 기본값 off의 opt-in으로 두는 쪽을 택했습니다. 실측 recall은 GLM-5.3 0.9999995, DeepSeek-V4.1 0.994(bf16 tie)이고, ruler2 128K 등 정확도 지표도 유의한 차이가 없었습니다.

일반적 교훈

1. 중간 결과를 materialize하지 마라. 정답이 rows × k인데 rows × keys를 HBM에 쓰는 것은 낭비입니다. 생산자 커널의 epilogue에서 소비 조건을 미리 적용하면 트래픽이 사라집니다.
2. 정확한 selection 대신 근사 threshold + 정확한 소규모 selection. 샘플 기반 threshold로 후보를 ~2.25k로 줄이면 나머지는 싸게 정확히 풀 수 있습니다.
3. 샘플링은 데이터 분포를 고려하라. 연속 블록 샘플은 clustered 데이터에서 편향되고, 단일 key·홀수 stride가 훨씬 견고했습니다.
4. 라우팅 threshold로 손익분기점을 지켜라. 짧은 컨텍스트에서는 dense 경로가 더 빠릅니다.
5. 손실 최적화는 flag 뒤에. 측정된 recall과 정확도 평가를 근거로 opt-in 제공하는 것이 실용적 타협입니다.

참고 자료

• LiteTopK paper — PR이 구현한 sampled-threshold top-k 기법
• CUDA Driver API: cuTensorMapEncodeTiled — strided KV 샘플링에 쓰인 TMA tensor map 생성 API
• PTX ISA: tcgen05 instructions — SM100 TMEM/MMA(kind::mxf4) 명령
• CUDA Programming Guide: Programmatic Dependent Launch — page table transform 커널의 PDL 사용
• DeepGEMM — baseline mqa_logits 커널 및 스코어링 코어의 원형

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.