Repository navigation
Sparse MQA logits follow-ups for nv_dev: 16-byte paged KV stride, workspace reuse under graph capture - #459
Conversation
| if (not workspace.defined()) { | ||
| // Warm up each stream before capture so one-time zeroing is not replayed with the graph. | ||
| DG_HOST_ASSERT(c10::cuda::currentStreamCaptureStatusMayInitCtx() == c10::cuda::CaptureStatus::None); | ||
| if (c10::cuda::currentStreamCaptureStatusMayInitCtx() != c10::cuda::CaptureStatus::None) { |
There was a problem hiding this comment.
🔵 suggestion: The capture-time fallback scans the static workspaces map and returns the first defined workspace on the same device. This matches the single warmed stream per device used by TensorRT-LLM, but if more than one stream on the device has warmed a workspace, the choice is unordered_map-order dependent and may not be the stream the captured graph later replays on. Since the workspace is shared per-launch scratch, that could alias scratch across concurrently-executing streams. Consider documenting the 'one warmed workspace per device' assumption or tracking the warm-up stream so the fallback is deterministic. Not a blocker for the validated path.
🤖 v4
There was a problem hiding this comment.
Kept the simple scan (TensorRT-LLM warms exactly one stream per device); the comment now states that assumption.
| DG_HOST_ASSERT(fused_kv_cache.scalar_type() == torch::kUInt8 and fused_kv_cache.stride(1) == head_dim_with_sf and | ||
| fused_kv_cache.stride(3) == 1 and fused_kv_cache.stride(0) <= std::numeric_limits<int>::max() and | ||
| fused_kv_cache.stride(0) % 512 == 0); | ||
| fused_kv_cache.stride(0) % 16 == 0); // pages are read with 16-byte cp.async chunks |
There was a problem hiding this comment.
🔵 suggestion: Relaxing the page stride to a multiple of 16 is correct for the 16-byte cp.async copies, but the paged entry point's 16-byte pointer-alignment loop covers {q_fp, q_sf, weights, metadata} and not fused_kv_cache. Page starts are base + page_idx * stride(0), so 16-byte KV access still requires the cache's data_ptr() to be 16-byte aligned. Torch allocations normally satisfy this, but an explicit assert or comment would make the contract complete and prevent a future strided/offset view from silently producing misaligned cp.async addresses.
🤖 v4
There was a problem hiding this comment.
Done: fused_kv_cache is now in the 16-byte pointer-alignment loop.
| } | ||
|
|
||
| static pybind11::dict get_bf16_mega_gate_config(const int& num_tokens, const int& hidden, | ||
| static std::map<std::string, int> get_bf16_mega_gate_config(const int& num_tokens, const int& hidden, |
There was a problem hiding this comment.
🔵 suggestion: The std::mapstd::string,int return type depends on the pybind11 STL type caster. It currently arrives transitively through torch/python.h, but adding an explicit #include <pybind11/stl.h> (next to the existing <map>/<string> includes) would make the requirement explicit and avoid a fragile transitive dependency. Otherwise the change is correct and yields the desired dict[str, int] stub.
🤖 v4
There was a problem hiding this comment.
Done: explicit #include <pybind11/stl.h>.
There was a problem hiding this comment.
Dropped from this PR: the stub typing is handled on the TensorRT-LLM side instead (mypy override for the generated stub), so mega_gate.hpp is untouched here.
🤖 ds-review-bot Code Reviewv6v5v4This MR makes three well-scoped, low-risk follow-ups to the DSv4.1 sparse MQA logits without touching kernel algorithms. (1) The paged sparse MQA logits host/device asserts relax the KV page-stride requirement from 512-byte to 16-byte alignment, which matches the PagedSparseKVAccessor's 16-byte cp.async chunk copies and admits native 4352-byte page-footer caches. (2) get_sparse_mqa_logits_workspace now reuses an already-warmed, same-device workspace when called while a CUDA graph capture is active instead of tripping the cold-start assert; the genuinely cold path still asserts. (3) get_bf16_mega_gate_config returns std::map<std::string,int> so generated stubs become dict[str,int] rather than a bare dict. The changes are internally consistent, preserve existing behavior for padded (512-byte) caches, and align with the described TensorRT-LLM integration and validation (453 CSA2 cases, paged sparse decode with CUDA-graph capture/replay against 4352-byte pages). I found no blocking correctness issues. The notes below are non-blocking robustness/documentation suggestions. Files reviewed: 3 |
0f4df5b to
1191ab8
Compare
lfr-0531
changed the title
Frameworks such as TensorRT-LLM capture CUDA graphs on a dedicated side stream after warming up eagerly on the launching stream. The sparse MQA logits workspace is keyed by stream, so the first call on the capture stream used to trip the pre-capture assertion. Captured work replays on the launching stream, so fall back to that stream's warmed-up workspace while capturing instead of recording a new allocation into the graph. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
The paged sparse kernel addresses pages directly and copies them with 16-byte cp.async chunks, so a page stride only has to be a multiple of 16 bytes. Caches whose page-footer layout is not padded to 512 bytes (for example 64 x (64 + 4) = 4352-byte pages) can now be read in place. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
1191ab8 to
337b1f8
Compare
3 participants
Summary
Two small follow-ups to #447 for the DSv4.1 sparse MQA logits, found while integrating
nv_devinto TensorRT-LLM (index cache pages are 64 x (64 + 4) = 4352-byte page-footer pages, and CUDA graphs are captured on a side stream). No kernel algorithm changes.fp8_fp4_paged_sparse_mqa_logitsrequiredfused_kv_cache.stride(0) % 512 == 0(host and device asserts). The paged accessor only computespage = base + page_idx * kv_page_stride_bytesand copies the KV / SF blocks with 16-bytecp.asyncchunks, so 16-byte alignment is the actual requirement (the cache base pointer is now checked for it as well). Caches whose page stride is not padded to 512 bytes can be read in place.Validation
nv_devhead alone the 7 paged decode cases fail on the 512-byte stride assert.Supersedes #458 (the nv_dev + main 26/09 merge), which #447 has landed.
🤖 Generated with Claude Code