[Feature] Add SM100 MQA logits kernels and native lowering - #2774
[Feature] Add SM100 MQA logits kernels and native lowering#2774Rachmanino wants to merge 18 commits into
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds FP8/FP4 SM100 MQA kernels and DeepGEMM benchmarks, improves SM100 vector-reduction and CUDA shuffle lowering, and updates TCGEN05/TMEM addressing with regression tests. ChangesSM100 MQA logits
SM100 vector reduction lowering
TCGEN05 and TMEM lowering
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant MQAData
participant TileLang
participant DeepGEMM
CLI->>MQAData: prepare quantized inputs
MQAData->>TileLang: provide staged tensors
MQAData->>DeepGEMM: provide packed tensors
CLI->>TileLang: run logits kernel
CLI->>DeepGEMM: run comparison kernel
CLI->>CLI: compare results and report timings
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
🧹 Nitpick comments (11)
examples/deepseek_deepgemm/mqa_logits_sm100.py (2)
107-107: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the unused
compressed_logitsparameter.It is never referenced in either kernel body, and all four call sites (
run_fp8,run_fp4, and both benchmark closures) passFalse. It only adds surface to the JIT signature. Same for Line 598.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/deepseek_deepgemm/mqa_logits_sm100.py` at line 107, Remove the unused compressed_logits parameter from the relevant kernel definitions and update all call sites, including run_fp8, run_fp4, and both benchmark closures, to stop passing False. Apply the same change to the duplicate definition near the referenced later location while preserving all other kernel arguments and behavior.
142-146: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDerive the FP4 TMEM width instead of hardcoding 512.
The accumulator region ends at
block_q * heads * num_tmem_stages(384) and the SF columns atsf_kv_col_1 + 4(400);512silently decouples from both. Consider allocating after computing the column offsets, e.g.T.alloc_tmem((half_kv, T.next_power_of_2(sf_kv_col_1 + 4)), ...)or an explicit named constant with a comment on the TMEM column-granularity requirement.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/deepseek_deepgemm/mqa_logits_sm100.py` around lines 142 - 146, Update the TMEM allocation in the surrounding MQA logits setup to derive its width from the computed SF column offsets rather than hardcoding 512. Move c_tmem allocation after sf_kv_col_1 is calculated and use the required TMEM column-granularity rounding of sf_kv_col_1 + 4, preserving accum_dtype and the existing row dimension.examples/deepseek_deepgemm/benchmark_mqa_logits_sm100.py (1)
5-6: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExplain the
sys.meta_pathfiltering.Stripping the
_tilelang_editablefinder is non-obvious and silently becomes a no-op if that internal module is renamed, leaving a confusing import failure. A one-line comment stating why the editable-install finder must be bypassed (and what breaks otherwise) would save the next reader a lot of time.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/deepseek_deepgemm/benchmark_mqa_logits_sm100.py` around lines 5 - 6, Add a concise comment immediately above the sys.meta_path filtering that explains the _tilelang_editable finder must be bypassed to avoid the import failure caused by the editable installation path, and notes the consequence if it remains active.tilelang/language/gemm_op.py (1)
247-247: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument
disable_wsin the docstring.The new keyword-only flag is undocumented; the docstring already explains
use_2cta.Also applies to: 249-261
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/language/gemm_op.py` at line 247, Update the docstring for the function containing the disable_ws parameter to document the new keyword-only flag, alongside the existing use_2cta explanation. Describe what disabling warp specialization changes, without altering the parameter’s behavior or surrounding documentation.tilelang/cuda/op/gemm/gemm_tcgen05.py (1)
53-72: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRewriter applies one offset to every matching call in the function.
_offset_tcgen05_tmem_operandsunconditionally rewrites alltl.ptx_tcgen05_mma_*calls in the lowered body with the samec_col_offset. That holds today because eachlower()produces exactly one MMA site, but it is fragile if a future variant emits several MMA calls for different C regions. A short comment (or an assertion that only one call was rewritten) would make the invariant explicit.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/cuda/op/gemm/gemm_tcgen05.py` around lines 53 - 72, Document the single-MMA-site invariant in `_offset_tcgen05_tmem_operands` by adding a concise comment or an assertion that verifies only one matching `tl.ptx_tcgen05_mma_*` call is rewritten per lowered function. Keep the existing offset application unchanged.testing/python/transform/test_tilelang_transform_lower_shared_tmem.py (1)
252-259: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
test_tmem_fragment_weighted_reduce_sum_correctnessis missing from the__main__list.Pytest still collects it, but direct execution of the file skips the only end-to-end check. Consider adding it (or switching to
tilelang.testing.main()).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/transform/test_tilelang_transform_lower_shared_tmem.py` around lines 252 - 259, The direct-execution test list under __main__ omits test_tmem_fragment_weighted_reduce_sum_correctness. Add this test function to the invocation list so running the file directly executes the end-to-end correctness check, while preserving the existing tests.testing/python/language/test_tilelang_language_tcgen05_gemm.py (1)
31-32: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNormalization weakens the sync-vs-async source equality check.
Dropping brace-only lines and all indentation means block-structure differences between the two codegen paths no longer fail this test. If the intent is only to tolerate the new lexical alloc scope, consider normalizing indentation only and keeping braces.
Also applies to: 102-104, 120-122
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/language/test_tilelang_language_tcgen05_gemm.py` around lines 31 - 32, Update _normalize_cuda_scope_formatting to preserve brace-only lines and block structure while tolerating only the intended indentation differences. Remove the filtering of "{"/"}" lines and normalize indentation consistently for the sync and async source comparisons, including the other uses of this helper.src/cuda/transform/lower_shared_tmem.cc (1)
352-358: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCollapse the duplicated branches.
Both arms produce the identical expression; the condition can be a single
if (buffer_remap_.count(buffer) || var_remap_.count(buffer->data)).♻️ Proposed simplification
- if (buffer_remap_.count(buffer)) { - return BufferLoad(tmem_base_buffer_remap_.at(buffer->data), {0}) + - GetTmemOffset(buffer, indices); - } else if (var_remap_.count(buffer->data)) { + if (buffer_remap_.count(buffer) || var_remap_.count(buffer->data)) { return BufferLoad(tmem_base_buffer_remap_.at(buffer->data), {0}) + GetTmemOffset(buffer, indices); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cuda/transform/lower_shared_tmem.cc` around lines 352 - 358, Collapse the duplicated conditional branches in the surrounding buffer remapping logic into one condition that checks buffer_remap_.count(buffer) or var_remap_.count(buffer->data), while preserving the shared BufferLoad and GetTmemOffset expression.testing/python/cuda/test_cuda_f32x2_intrinsics.py (2)
174-174: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefix the unused block-index unpacks to satisfy Ruff RUF059.
bxandbyare unused in all three new kernel builders.♻️ Proposed fix (apply at each of the three sites)
- with T.Kernel(1, 1, threads=M) as (bx, by): + with T.Kernel(1, 1, threads=M) as (_bx, _by):As per static analysis: "Unpacked variable
bxis never used — Prefix it with an underscore or any other dummy variable pattern (RUF059)".Also applies to: 195-195, 235-235
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/cuda/test_cuda_f32x2_intrinsics.py` at line 174, Update the three kernel builder unpackings in the CUDA intrinsic tests to use underscore-prefixed names for the unused block indices instead of bx and by. Apply the change at each with T.Kernel(...) as (bx, by) site while preserving the existing kernel behavior and thread configuration.Source: Linters/SAST tools
432-432: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNegative assertion silently loses coverage if the loop var is renamed.
"for (int h = 0; h < 8; ++h)"depends on the TIR loop variable keeping the namehthroughAllocVarIDand on the exact spacing emitted byCodeGenTileLangCUDA::VisitStmt_(ForNode). Because it is a negative assertion, any drift makes it pass vacuously rather than fail. Consider matching on a regex for any 8-iteration scalar loop instead.♻️ Proposed tightening
- assert "for (int h = 0; h < 8; ++h)" not in src, "Scalar reduction loop should not be emitted as an 8-iteration loop" + assert not re.search(r"for \(int \w+ = 0; \w+ < 8; \+\+\w+\)", src), "Scalar reduction loop should not be emitted as an 8-iteration loop"Requires
import reat the top of the file.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/cuda/test_cuda_f32x2_intrinsics.py` at line 432, Replace the exact-string negative assertion in the relevant CUDA intrinsic test with a regular-expression check that rejects any emitted scalar loop iterating eight times, regardless of the loop variable name or formatting. Add the required re import and preserve the assertion’s intent of preventing an 8-iteration reduction loop.src/cuda/codegen/codegen_cuda.cc (1)
5325-5382: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider extracting the shared concat-and-select helper.
Lines 5329-5353 duplicate the generic concat-then-shuffle logic already present in the int32
lanes > 4fallback at lines 5405-5427 (element expansion viaPrintVecElemLoad, constant-index validation, bounds check, scalar selection). A small helper returningstd::vector<std::string>of selected scalars would let both paths share it and keep the two index-validation checks in sync.♻️ Sketch of the shared helper
// Collect the shuffle's selected lanes as CUDA-aware scalar expressions. std::vector<std::string> CodeGenTileLangCUDA::CollectShuffleScalars(const ShuffleNode *op) { std::vector<std::string> concat_vec; for (const PrimExpr &vec : op->vectors) { std::string vec_value = PrintExpr(vec); if (vec.dtype().is_scalar()) { concat_vec.push_back(vec_value); } else { for (int i = 0; i < vec.dtype().lanes(); ++i) { std::ostringstream elem; PrintVecElemLoad(vec_value, vec.dtype(), i, elem); concat_vec.push_back(elem.str()); } } } std::vector<std::string> scalars; scalars.reserve(op->indices.size()); for (const PrimExpr &index : op->indices) { const auto *lane = index.as<IntImmNode>(); ICHECK(lane) << "ShuffleNode indices must be constants at codegen time, " "got " << index; ICHECK_GE(lane->value, 0); ICHECK_LT(lane->value, static_cast<int64_t>(concat_vec.size())); scalars.push_back(concat_vec[lane->value]); } return scalars; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cuda/codegen/codegen_cuda.cc` around lines 5325 - 5382, Extract the duplicated concat-and-select logic from the float32 shuffle branch and the int32 lanes-greater-than-four fallback into a shared CodeGenTileLangCUDA helper returning selected scalar expressions, such as CollectShuffleScalars. Have the helper perform vector element expansion via PrintVecElemLoad, constant-index validation, bounds checks, and selection, then replace both existing implementations with calls to it while preserving their type-specific output construction.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/deepseek_deepgemm/mqa_logits_sm100.py`:
- Around line 483-485: Update the logits epilogue stores at the visible
assignment and the corresponding stores around the other reported locations so
each written element is predicated by its query row’s own [KS, KE) range,
storing -inf outside that range. Apply the same per-row masking behavior
consistently across all affected paths, including outputs expected from
clean_logits=True, rather than relying on callers or validation code to mask
afterward.
- Around line 47-51: Update the seq_len_kv validation near the existing
assertions to require divisibility by 256, matching the kernels’ block_kv tile
size. Apply this requirement consistently to the FP8 path and any shared
validation used by run_fp4 or prepare_mqa_data, while preserving the existing
minimum-difference and seq_len checks.
In `@src/cuda/op/copy.cc`:
- Around line 578-593: In the physical extent calculation around
relative_phy_indices, validate that both phy_row_bounds and phy_col_bounds are
finite and have min_value equal to zero before computing extents or narrowing to
int. Add ICHECK guards for these invariants, then preserve the existing extent
and Range construction only after validation so symbolic or nonzero-offset
bounds cannot reach ExpandTcgen05Layout.
- Around line 566-569: In InferTMemLayout, replace the locally constructed
analyzer used for the TMem warpgroup alignment proof with layout_args.analyzer.
Apply the same analyzer to the associated Simplify calls, while preserving the
existing logical-coordinate bindings and proof behavior.
In `@src/transform/layout_inference.cc`:
- Around line 365-370: In the strict-layout seeding loop, narrow the
floating-buffer exclusion so explicitly annotated buffers from
annotated_layout_map_ remain in strict_layout_map. Skip only floating buffers
originating from the step-0 replicated seeding, while preserving the override
behavior for unannotated floating fragment buffers.
- Around line 213-219: Update the floating fragment buffer branch in layout
inference to reject layouts whose destination ReplicateExtent() is one while the
buffer still requires full replication. Preserve the FullyReplicated ownership
invariant through floating accesses, unless this path explicitly proves that a
later pass restores ownership before use; do not call layout_map.Set or
propagate_alias for the invalid transition.
In `@src/transform/loop_vectorize.cc`:
- Around line 452-471: Restrict the self-add-store exemption in
IsLocalScalarSelfAddStore to stores whose buffer indices are
loop-invariant/independent, matching the condition used by
ComputeBufferVectorSize. Keep indexed per-element stores such as acc[i] = acc[i]
+ x[i] on the UpdateVectorSize path so their vectorization constraints and
local_fragment_buffers re-validation remain intact.
- Around line 966-971: Update VisitStmt_ in VectorizeRewriter so unit-extent
loops still pass through the existing kParallel-to-kSerial rewrite in the
non-innermost path before any early return. Use
analyzer_->Simplify(node->extent), matching IsTrivialUnitLoop, when determining
whether the extent is one, while preserving the innermost-loop handling.
In `@src/transform/vectorize_loop.cc`:
- Around line 1652-1687: Before dropping the defining store for vec_buffer in
the stmts loop, add a guard that rejects the transformation when any other
statement accesses vec_buffer at vec_indices, using the existing
ContainsBufferElementAccess helper. Keep the existing store-skipping behavior,
but ensure prefix readers are detected before the defining store is removed;
leave unrelated AllocBufferNode handling unchanged.
In `@testing/python/cuda/test_cuda_f32x2_intrinsics.py`:
- Around line 466-474: Relax the tolerances in
test_correctness_auto_vec_scalar_reduction_chunk_accumulator_f32_sm100 to
accommodate expected float32 reassociation differences between the kernel
reduction and the torch reference, while keeping the existing correctness
assertion and inputs unchanged.
In `@tilelang/cuda/op/gemm/gemm_tcgen05.py`:
- Line 51: Update the TMEM C-offset calculation at c_tmem_offset so it does not
silently floor unaligned values: assert that c_col_offset is a multiple of 32 //
accum_dtype_bits (or otherwise require word-aligned accumulation dimensions),
then retain the existing conversion only for validated inputs.
---
Nitpick comments:
In `@examples/deepseek_deepgemm/benchmark_mqa_logits_sm100.py`:
- Around line 5-6: Add a concise comment immediately above the sys.meta_path
filtering that explains the _tilelang_editable finder must be bypassed to avoid
the import failure caused by the editable installation path, and notes the
consequence if it remains active.
In `@examples/deepseek_deepgemm/mqa_logits_sm100.py`:
- Line 107: Remove the unused compressed_logits parameter from the relevant
kernel definitions and update all call sites, including run_fp8, run_fp4, and
both benchmark closures, to stop passing False. Apply the same change to the
duplicate definition near the referenced later location while preserving all
other kernel arguments and behavior.
- Around line 142-146: Update the TMEM allocation in the surrounding MQA logits
setup to derive its width from the computed SF column offsets rather than
hardcoding 512. Move c_tmem allocation after sf_kv_col_1 is calculated and use
the required TMEM column-granularity rounding of sf_kv_col_1 + 4, preserving
accum_dtype and the existing row dimension.
In `@src/cuda/codegen/codegen_cuda.cc`:
- Around line 5325-5382: Extract the duplicated concat-and-select logic from the
float32 shuffle branch and the int32 lanes-greater-than-four fallback into a
shared CodeGenTileLangCUDA helper returning selected scalar expressions, such as
CollectShuffleScalars. Have the helper perform vector element expansion via
PrintVecElemLoad, constant-index validation, bounds checks, and selection, then
replace both existing implementations with calls to it while preserving their
type-specific output construction.
In `@src/cuda/transform/lower_shared_tmem.cc`:
- Around line 352-358: Collapse the duplicated conditional branches in the
surrounding buffer remapping logic into one condition that checks
buffer_remap_.count(buffer) or var_remap_.count(buffer->data), while preserving
the shared BufferLoad and GetTmemOffset expression.
In `@testing/python/cuda/test_cuda_f32x2_intrinsics.py`:
- Line 174: Update the three kernel builder unpackings in the CUDA intrinsic
tests to use underscore-prefixed names for the unused block indices instead of
bx and by. Apply the change at each with T.Kernel(...) as (bx, by) site while
preserving the existing kernel behavior and thread configuration.
- Line 432: Replace the exact-string negative assertion in the relevant CUDA
intrinsic test with a regular-expression check that rejects any emitted scalar
loop iterating eight times, regardless of the loop variable name or formatting.
Add the required re import and preserve the assertion’s intent of preventing an
8-iteration reduction loop.
In `@testing/python/language/test_tilelang_language_tcgen05_gemm.py`:
- Around line 31-32: Update _normalize_cuda_scope_formatting to preserve
brace-only lines and block structure while tolerating only the intended
indentation differences. Remove the filtering of "{"/"}" lines and normalize
indentation consistently for the sync and async source comparisons, including
the other uses of this helper.
In `@testing/python/transform/test_tilelang_transform_lower_shared_tmem.py`:
- Around line 252-259: The direct-execution test list under __main__ omits
test_tmem_fragment_weighted_reduce_sum_correctness. Add this test function to
the invocation list so running the file directly executes the end-to-end
correctness check, while preserving the existing tests.
In `@tilelang/cuda/op/gemm/gemm_tcgen05.py`:
- Around line 53-72: Document the single-MMA-site invariant in
`_offset_tcgen05_tmem_operands` by adding a concise comment or an assertion that
verifies only one matching `tl.ptx_tcgen05_mma_*` call is rewritten per lowered
function. Keep the existing offset application unchanged.
In `@tilelang/language/gemm_op.py`:
- Line 247: Update the docstring for the function containing the disable_ws
parameter to document the new keyword-only flag, alongside the existing use_2cta
explanation. Describe what disabling warp specialization changes, without
altering the parameter’s behavior or surrounding documentation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 98f5e9b5-1092-4fd1-a720-542235703e40
📒 Files selected for processing (14)
examples/deepseek_deepgemm/benchmark_mqa_logits_sm100.pyexamples/deepseek_deepgemm/mqa_logits_sm100.pysrc/cuda/codegen/codegen_cuda.ccsrc/cuda/op/copy.ccsrc/cuda/transform/lower_shared_tmem.ccsrc/transform/layout_inference.ccsrc/transform/loop_vectorize.ccsrc/transform/vectorize_loop.cctesting/python/cuda/test_cuda_f32x2_intrinsics.pytesting/python/language/test_tilelang_language_tcgen05_gemm.pytesting/python/transform/test_tilelang_transform_inject_tcgen05_fence.pytesting/python/transform/test_tilelang_transform_lower_shared_tmem.pytilelang/cuda/op/gemm/gemm_tcgen05.pytilelang/language/gemm_op.py
| Logits[q_row + qi_epi0, kv_row + bn_epi0] = T.cast( | ||
| logits_epi0[bn_epi0], logits_dtype | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Epilogue writes unmasked values outside each row's [ks, ke) range.
The KV tile range is derived from min(KS)/max(KE) across the whole block_q tile, but the store applies no per-row predicate, so columns inside the tile union yet outside a given row's own range get real logits instead of -inf. The validation path hides this: run_example_case (Lines 1224-1226) and the benchmark both mask by ref == -inf before calc_diff, whereas DeepGEMM's clean_logits=True output is genuinely masked. Either apply the per-row mask when storing, or document loudly that the caller must mask. Same for Lines 561-563, 883-885, 952-954.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/deepseek_deepgemm/mqa_logits_sm100.py` around lines 483 - 485,
Update the logits epilogue stores at the visible assignment and the
corresponding stores around the other reported locations so each written element
is predicated by its query row’s own [KS, KE) range, storing -inf outside that
range. Apply the same per-row masking behavior consistently across all affected
paths, including outputs expected from clean_logits=True, rather than relying on
callers or validation code to mask afterward.
| def test_correctness_auto_vec_scalar_reduction_chunk_accumulator_f32_sm100(): | ||
| func = _make_auto_vec_scalar_reduction_chunks_kernel() | ||
| kernel = tilelang.compile(func, out_idx=[3], target=SM100_TARGET) | ||
| scores = torch.randn(M, 4, 8, device="cuda", dtype=torch.float32) | ||
| weights = torch.randn(M, 4, 8, device="cuda", dtype=torch.float32) | ||
| scale = torch.randn(M, device="cuda", dtype=torch.float32) | ||
| out = kernel(scores, weights, scale) | ||
| ref = (torch.maximum(scores * scale[:, None, None], torch.zeros_like(scores)) * weights).sum(dim=(1, 2)) | ||
| torch.testing.assert_close(out, ref, rtol=1e-5, atol=1e-5) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Tolerance is likely too tight for the reassociated f32 reduction.
The kernel deliberately changes the summation order (packed FMA chains + horizontal fold) versus torch's .sum(dim=(1,2)). Summing 32 products of randn values in f32 under a different association can drift past atol=1e-5 for results of magnitude ~5, making this test intermittently flaky.
💚 Proposed fix
out = kernel(scores, weights, scale)
- ref = (torch.maximum(scores * scale[:, None, None], torch.zeros_like(scores)) * weights).sum(dim=(1, 2))
- torch.testing.assert_close(out, ref, rtol=1e-5, atol=1e-5)
+ # Compute the reference in fp64: the kernel reassociates the reduction
+ # (packed FMA chains + horizontal fold), so an f32 reference of a
+ # different association is not a stable baseline at tight tolerances.
+ ref = (
+ torch.maximum(scores.double() * scale.double()[:, None, None], torch.zeros_like(scores, dtype=torch.float64))
+ * weights.double()
+ ).sum(dim=(1, 2))
+ torch.testing.assert_close(out.double(), ref, rtol=1e-5, atol=1e-4)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def test_correctness_auto_vec_scalar_reduction_chunk_accumulator_f32_sm100(): | |
| func = _make_auto_vec_scalar_reduction_chunks_kernel() | |
| kernel = tilelang.compile(func, out_idx=[3], target=SM100_TARGET) | |
| scores = torch.randn(M, 4, 8, device="cuda", dtype=torch.float32) | |
| weights = torch.randn(M, 4, 8, device="cuda", dtype=torch.float32) | |
| scale = torch.randn(M, device="cuda", dtype=torch.float32) | |
| out = kernel(scores, weights, scale) | |
| ref = (torch.maximum(scores * scale[:, None, None], torch.zeros_like(scores)) * weights).sum(dim=(1, 2)) | |
| torch.testing.assert_close(out, ref, rtol=1e-5, atol=1e-5) | |
| def test_correctness_auto_vec_scalar_reduction_chunk_accumulator_f32_sm100(): | |
| func = _make_auto_vec_scalar_reduction_chunks_kernel() | |
| kernel = tilelang.compile(func, out_idx=[3], target=SM100_TARGET) | |
| scores = torch.randn(M, 4, 8, device="cuda", dtype=torch.float32) | |
| weights = torch.randn(M, 4, 8, device="cuda", dtype=torch.float32) | |
| scale = torch.randn(M, device="cuda", dtype=torch.float32) | |
| out = kernel(scores, weights, scale) | |
| # Compute the reference in fp64: the kernel reassociates the reduction | |
| # (packed FMA chains + horizontal fold), so an f32 reference of a | |
| # different association is not a stable baseline at tight tolerances. | |
| ref = ( | |
| torch.maximum(scores.double() * scale.double()[:, None, None], torch.zeros_like(scores, dtype=torch.float64)) | |
| * weights.double() | |
| ).sum(dim=(1, 2)) | |
| torch.testing.assert_close(out.double(), ref, rtol=1e-5, atol=1e-4) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@testing/python/cuda/test_cuda_f32x2_intrinsics.py` around lines 466 - 474,
Relax the tolerances in
test_correctness_auto_vec_scalar_reduction_chunk_accumulator_f32_sm100 to
accommodate expected float32 reassociation differences between the kernel
reduction and the torch reference, while keeping the existing correctness
assertion and inputs unchanged.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py (1)
448-518: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winPreserve A’s TMEM column offset in the TS variant.
compute_tcgen05_a_desc_params(A_buf)uses theBufferRegionslice, buttcgen05mma_tsthen passes onlyA_buf.buffer.datatotcgen05_ts_atom. Since_warp_mma_tspasses that raw pointer directly andtcgen05_ts_atomcomputesA_tmem_offsetfrom onlyinst_m_idx/ki, a column-sliced A operand such asT.tcgen05_gemm(A_tmem[:, off:off+chunk], B, C, ...)loses the slice base and reads from the start of the TMEM allocation. Thread_resolve_tmem_operand()’s A base offset throughtcgen05_ts_atomso TS slices are generated consistently with the sliced C/SFA/SFB cases.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py` around lines 448 - 518, The TS path drops A’s BufferRegion base offset when passing the raw TMEM pointer, causing column-sliced operands to read from the allocation start. Resolve A with _resolve_tmem_operand() in tcgen05mma_ts, thread its base offset through _warp_mma_ts into tcgen05_ts_atom, and use it when computing A_tmem_offset while preserving existing C offset handling.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py`:
- Around line 448-518: The TS path drops A’s BufferRegion base offset when
passing the raw TMEM pointer, causing column-sliced operands to read from the
allocation start. Resolve A with _resolve_tmem_operand() in tcgen05mma_ts,
thread its base offset through _warp_mma_ts into tcgen05_ts_atom, and use it
when computing A_tmem_offset while preserving existing C offset handling.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 3ecd6e7f-fefe-4e57-959d-eed325e7fa4d
📒 Files selected for processing (5)
src/cuda/op/copy.cctesting/python/language/test_tilelang_language_tcgen05_gemm.pytesting/python/transform/test_tilelang_transform_inject_tcgen05_fence.pytilelang/cuda/intrinsics/macro/tcgen05_macro_generator.pytilelang/cuda/op/gemm/gemm_tcgen05.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/cuda/op/copy.cc
da55531 to
4c5aaff
Compare
4c5aaff to
a83d606
Compare
5a94f6d to
14ffa8a
Compare
dcafe3f to
d5763c4
Compare
b5cbf90 to
c48459f
Compare
1 participant
Summary
C_tmem[3, 128, 128]SFQ_tmem[1, 128, 4]SFKV_tmem[2, 128, 4]tcgen05_cp_warpx4(..., sf_tmem[stage, :, :])without hand-written physical column offsets.Scope
tcgen05.allocarenas #2831 lowers the logical C/SFA/SFB allocations to one 512-columntcgen05.alloc.Performance
B200,
s=2048,skv=4096,h=64,d=128, DeepGEMM 2.6.1 (559d79f),TILELANG_DISABLE_CACHE=1, raw output (clean_logits=Falsefor both implementations), 25 ms warmup / 100 ms measurement. Values are medians of three rounds on the same GPU.Generated SASS has no spill:
Validation